added test

This commit is contained in:
Lorenze Jay
2024-11-26 12:06:48 -08:00
parent b3deac2a2b
commit 3f87bf3ada

View File

@@ -7,16 +7,18 @@ from unittest.mock import patch
import pytest
from crewai import Agent, Crew, Task
from crewai.knowledge.source.base_knowledge_source import BaseKnowledgeSource
from crewai.agents.cache import CacheHandler
from crewai.agents.crew_agent_executor import CrewAgentExecutor
from crewai.agents.parser import AgentAction, CrewAgentParser, OutputParserException
from crewai.knowledge.source.string_knowledge_source import StringKnowledgeSource
from crewai.knowledge.knowledge import Knowledge
from crewai.llm import LLM
from crewai.tools import tool
from crewai.tools.tool_calling import InstructorToolCalling
from crewai.tools.tool_usage import ToolUsage
from crewai.tools.tool_usage_events import ToolUsageFinished
from crewai.utilities import RPMController
from crewai.knowledge.source.string_knowledge_source import StringKnowledgeSource
from crewai.utilities.events import Emitter
@@ -1584,16 +1586,16 @@ def test_agent_with_knowledge_sources():
string_source = StringKnowledgeSource(
content=content, metadata={"preference": "personal"}
)
with patch('crewai.knowledge.storage.knowledge_storage.KnowledgeStorage') as MockKnowledge:
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,
"metadata": {"preference": "personal"}
}]
mock_knowledge_instance.query.return_value = [
{"content": content, "metadata": {"preference": "personal"}}
]
agent = Agent(
role="Information Agent",
goal="Provide information based on knowledge sources",
@@ -1614,3 +1616,21 @@ def test_agent_with_knowledge_sources():
# Assert that the agent provides the correct information
assert "blue" in result.raw.lower()
def test_agent_with_knowledge_sources_context():
content = "Brandon's favorite color is blue and he likes Mexican food."
string_source = StringKnowledgeSource(
content=content, metadata={"preference": "personal"}
)
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=[string_source],
)
# Test that agent is properly initialized with knowledge sources
assert isinstance(agent.knowledge, Knowledge)
assert len(agent.knowledge.sources) == 1
assert isinstance(agent.knowledge.sources[0], BaseKnowledgeSource)
assert agent.knowledge.sources[0].metadata == {"preference": "personal"}