mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 00:28:31 +00:00
fix: support multiple URL config keys for Ollama embedder (#1994)
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
@@ -82,12 +82,33 @@ class EmbeddingConfigurator:
|
|||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _configure_ollama(config, model_name):
|
def _configure_ollama(config, model_name):
|
||||||
|
"""Configure Ollama embedder with flexible URL configuration.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: Configuration dictionary that supports multiple URL keys:
|
||||||
|
- url: Legacy key (default: http://localhost:11434/api/embeddings)
|
||||||
|
- api_url: Alternative key following HuggingFace pattern
|
||||||
|
- base_url: Alternative key
|
||||||
|
- api_base: Alternative key following Azure pattern
|
||||||
|
model_name: Name of the Ollama model to use
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
OllamaEmbeddingFunction: Configured embedder instance
|
||||||
|
"""
|
||||||
from chromadb.utils.embedding_functions.ollama_embedding_function import (
|
from chromadb.utils.embedding_functions.ollama_embedding_function import (
|
||||||
OllamaEmbeddingFunction,
|
OllamaEmbeddingFunction,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
url = (
|
||||||
|
config.get("url")
|
||||||
|
or config.get("api_url")
|
||||||
|
or config.get("base_url")
|
||||||
|
or config.get("api_base")
|
||||||
|
or "http://localhost:11434/api/embeddings"
|
||||||
|
)
|
||||||
|
|
||||||
return OllamaEmbeddingFunction(
|
return OllamaEmbeddingFunction(
|
||||||
url=config.get("url", "http://localhost:11434/api/embeddings"),
|
url=url,
|
||||||
model_name=model_name,
|
model_name=model_name,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
56
tests/embedder_test.py
Normal file
56
tests/embedder_test.py
Normal file
@@ -0,0 +1,56 @@
|
|||||||
|
import pytest
|
||||||
|
from unittest.mock import patch
|
||||||
|
from crewai.utilities.embedding_configurator import EmbeddingConfigurator
|
||||||
|
|
||||||
|
def test_ollama_embedder_url_config():
|
||||||
|
configurator = EmbeddingConfigurator()
|
||||||
|
|
||||||
|
test_cases = [
|
||||||
|
# Test default URL
|
||||||
|
{
|
||||||
|
"config": {"provider": "ollama", "config": {"model": "test-model"}},
|
||||||
|
"expected_url": "http://localhost:11434/api/embeddings"
|
||||||
|
},
|
||||||
|
# Test legacy url key
|
||||||
|
{
|
||||||
|
"config": {"provider": "ollama", "config": {"model": "test-model", "url": "http://custom:11434"}},
|
||||||
|
"expected_url": "http://custom:11434"
|
||||||
|
},
|
||||||
|
# Test api_url key
|
||||||
|
{
|
||||||
|
"config": {"provider": "ollama", "config": {"model": "test-model", "api_url": "http://api:11434"}},
|
||||||
|
"expected_url": "http://api:11434"
|
||||||
|
},
|
||||||
|
# Test base_url key
|
||||||
|
{
|
||||||
|
"config": {"provider": "ollama", "config": {"model": "test-model", "base_url": "http://base:11434"}},
|
||||||
|
"expected_url": "http://base:11434"
|
||||||
|
},
|
||||||
|
# Test api_base key
|
||||||
|
{
|
||||||
|
"config": {"provider": "ollama", "config": {"model": "test-model", "api_base": "http://base-api:11434"}},
|
||||||
|
"expected_url": "http://base-api:11434"
|
||||||
|
},
|
||||||
|
# Test URL precedence order
|
||||||
|
{
|
||||||
|
"config": {
|
||||||
|
"provider": "ollama",
|
||||||
|
"config": {
|
||||||
|
"model": "test-model",
|
||||||
|
"url": "http://url:11434",
|
||||||
|
"api_url": "http://api:11434",
|
||||||
|
"base_url": "http://base:11434",
|
||||||
|
"api_base": "http://base-api:11434"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"expected_url": "http://url:11434" # url key should have highest precedence
|
||||||
|
}
|
||||||
|
]
|
||||||
|
|
||||||
|
for test_case in test_cases:
|
||||||
|
with patch("chromadb.utils.embedding_functions.ollama_embedding_function.OllamaEmbeddingFunction") as mock_ollama:
|
||||||
|
configurator.configure_embedder(test_case["config"])
|
||||||
|
mock_ollama.assert_called_once()
|
||||||
|
_, kwargs = mock_ollama.call_args
|
||||||
|
assert kwargs["url"] == test_case["expected_url"]
|
||||||
|
mock_ollama.reset_mock()
|
||||||
Reference in New Issue
Block a user