more fixes

This commit is contained in:
Brandon Hancock
2024-12-12 14:12:35 -05:00
parent 8ad6bf5cd1
commit 476d7433f9
2 changed files with 37 additions and 24 deletions

View File

@@ -177,6 +177,19 @@ class KnowledgeStorage(BaseKnowledgeStorage):
except Exception as e: except Exception as e:
Logger(verbose=True).log("error", f"Failed to upsert documents: {e}", "red") Logger(verbose=True).log("error", f"Failed to upsert documents: {e}", "red")
raise 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. """Set the embedding configuration for the knowledge storage.
Args: Args:

View File

@@ -1592,36 +1592,36 @@ def test_agent_execute_task_with_ollama():
assert "AI" in result or "artificial intelligence" in result.lower() 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(): def test_agent_with_knowledge_sources():
# Create a knowledge source with some content # Create a knowledge source with some content
content = "Brandon's favorite color is red and he likes Mexican food." content = "Brandon's favorite color is red and he likes Mexican food."
string_source = StringKnowledgeSource(content=content) string_source = StringKnowledgeSource(content=content)
# with patch( with patch(
# "crewai.knowledge.storage.knowledge_storage.KnowledgeStorage" "crewai.knowledge.storage.knowledge_storage.KnowledgeStorage"
# ) as MockKnowledge: ) as MockKnowledge:
# mock_knowledge_instance = MockKnowledge.return_value mock_knowledge_instance = MockKnowledge.return_value
# mock_knowledge_instance.sources = [string_source] mock_knowledge_instance.sources = [string_source]
# mock_knowledge_instance.query.return_value = [{"content": content}] mock_knowledge_instance.query.return_value = [{"content": content}]
agent = Agent( agent = Agent(
role="Information Agent", role="Information Agent",
goal="Provide information based on knowledge sources", goal="Provide information based on knowledge sources",
backstory="You have access to specific knowledge sources.", backstory="You have access to specific knowledge sources.",
llm=LLM(model="gpt-4o-mini"), llm=LLM(model="gpt-4o-mini"),
knowledge_sources=[string_source], knowledge_sources=[string_source],
) )
# Create a task that requires the agent to use the knowledge # Create a task that requires the agent to use the knowledge
task = Task( task = Task(
description="What is Brandon's favorite color?", description="What is Brandon's favorite color?",
expected_output="Brandon's favorite color.", expected_output="Brandon's favorite color.",
agent=agent, agent=agent,
) )
crew = Crew(agents=[agent], tasks=[task]) crew = Crew(agents=[agent], tasks=[task])
result = crew.kickoff() result = crew.kickoff()
# Assert that the agent provides the correct information # Assert that the agent provides the correct information
assert "red" in result.raw.lower() assert "red" in result.raw.lower()