diff --git a/tests/agent_test.py b/tests/agent_test.py index fb6ab22b8..5163a8e55 100644 --- a/tests/agent_test.py +++ b/tests/agent_test.py @@ -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"}