From 476d7433f9991938bfc76c298767136e96ed6bfa Mon Sep 17 00:00:00 2001
From: Brandon Hancock <brandon@brandonhancock.io>
Date: Thu, 12 Dec 2024 14:12:35 -0500
Subject: [PATCH] more fixes

---
 .../knowledge/storage/knowledge_storage.py    | 13 +++++
 tests/agent_test.py                           | 48 +++++++++----------
 2 files changed, 37 insertions(+), 24 deletions(-)

diff --git a/src/crewai/knowledge/storage/knowledge_storage.py b/src/crewai/knowledge/storage/knowledge_storage.py
index 13e466faf..4a70c5997 100644
--- a/src/crewai/knowledge/storage/knowledge_storage.py
+++ b/src/crewai/knowledge/storage/knowledge_storage.py
@@ -177,6 +177,19 @@ class KnowledgeStorage(BaseKnowledgeStorage):
         except Exception as e:
             Logger(verbose=True).log("error", f"Failed to upsert documents: {e}", "red")
             raise
+
+    def _create_default_embedding_function(self):
+        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"
+        )
+
+    def _set_embedder_config(
+        self, embedder_config: Optional[Dict[str, Any]] = None
+    ) -> None:
         """Set the embedding configuration for the knowledge storage.
 
         Args:
diff --git a/tests/agent_test.py b/tests/agent_test.py
index 164cfc096..048cc4807 100644
--- a/tests/agent_test.py
+++ b/tests/agent_test.py
@@ -1592,36 +1592,36 @@ def test_agent_execute_task_with_ollama():
     assert "AI" in result or "artificial intelligence" in result.lower()
 
 
-@pytest.mark.vcr(filter_headers=["authorization"])
+# @pytest.mark.vcr(filter_headers=["authorization"])
 def test_agent_with_knowledge_sources():
     # Create a knowledge source with some content
     content = "Brandon's favorite color is red and he likes Mexican food."
     string_source = StringKnowledgeSource(content=content)
 
-    # with patch(
-    #     "crewai.knowledge.storage.knowledge_storage.KnowledgeStorage"
-    # ) as MockKnowledge:
-    #     mock_knowledge_instance = MockKnowledge.return_value
-    #     mock_knowledge_instance.sources = [string_source]
-    #     mock_knowledge_instance.query.return_value = [{"content": content}]
+    with patch(
+        "crewai.knowledge.storage.knowledge_storage.KnowledgeStorage"
+    ) as MockKnowledge:
+        mock_knowledge_instance = MockKnowledge.return_value
+        mock_knowledge_instance.sources = [string_source]
+        mock_knowledge_instance.query.return_value = [{"content": content}]
 
-    agent = Agent(
-        role="Information Agent",
-        goal="Provide information based on knowledge sources",
-        backstory="You have access to specific knowledge sources.",
-        llm=LLM(model="gpt-4o-mini"),
-        knowledge_sources=[string_source],
-    )
+        agent = Agent(
+            role="Information Agent",
+            goal="Provide information based on knowledge sources",
+            backstory="You have access to specific knowledge sources.",
+            llm=LLM(model="gpt-4o-mini"),
+            knowledge_sources=[string_source],
+        )
 
-    # Create a task that requires the agent to use the knowledge
-    task = Task(
-        description="What is Brandon's favorite color?",
-        expected_output="Brandon's favorite color.",
-        agent=agent,
-    )
+        # Create a task that requires the agent to use the knowledge
+        task = Task(
+            description="What is Brandon's favorite color?",
+            expected_output="Brandon's favorite color.",
+            agent=agent,
+        )
 
-    crew = Crew(agents=[agent], tasks=[task])
-    result = crew.kickoff()
+        crew = Crew(agents=[agent], tasks=[task])
+        result = crew.kickoff()
 
-    # Assert that the agent provides the correct information
-    assert "red" in result.raw.lower()
+        # Assert that the agent provides the correct information
+        assert "red" in result.raw.lower()