From 509fb375cabb6dc2677d1729f1c18f71762cb1d2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 5 Feb 2025 10:43:58 +0000 Subject: [PATCH] feat: Update embedding configurator to support configurable default providers Co-Authored-By: Joe Moura --- .../utilities/embedding_configurator.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/src/crewai/utilities/embedding_configurator.py b/src/crewai/utilities/embedding_configurator.py index 44e832ec2..b86541a38 100644 --- a/src/crewai/utilities/embedding_configurator.py +++ b/src/crewai/utilities/embedding_configurator.py @@ -47,13 +47,19 @@ class EmbeddingConfigurator: @staticmethod def _create_default_embedding_function(): - from chromadb.utils.embedding_functions.openai_embedding_function import ( - OpenAIEmbeddingFunction, - ) - - return OpenAIEmbeddingFunction( - api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small" - ) + from crewai.utilities.constants import DEFAULT_EMBEDDING_MODEL, DEFAULT_EMBEDDING_PROVIDER + + provider = os.getenv("CREWAI_EMBEDDING_PROVIDER", DEFAULT_EMBEDDING_PROVIDER) + model = os.getenv("CREWAI_EMBEDDING_MODEL", DEFAULT_EMBEDDING_MODEL) + + if provider == "openai": + from chromadb.utils.embedding_functions.openai_embedding_function import OpenAIEmbeddingFunction + return OpenAIEmbeddingFunction(api_key=os.getenv("OPENAI_API_KEY"), model_name=model) + elif provider == "ollama": + from chromadb.utils.embedding_functions.ollama_embedding_function import OllamaEmbeddingFunction + return OllamaEmbeddingFunction(url=os.getenv("CREWAI_OLLAMA_URL", "http://localhost:11434/api/embeddings"), model_name=model) + else: + raise ValueError(f"Unsupported default embedding provider: {provider}. Set CREWAI_EMBEDDING_PROVIDER to 'openai' or 'ollama'") @staticmethod def _configure_openai(config, model_name):