From 8570461969c2d374b8c463150f893bbf99d04384 Mon Sep 17 00:00:00 2001 From: Lorenze Jay Date: Mon, 27 Jan 2025 15:39:58 -0800 Subject: [PATCH] checker --- src/crewai/agents/agent_builder/base_agent.py | 8 ++- tests/agent_test.py | 68 +++++++++---------- 2 files changed, 37 insertions(+), 39 deletions(-) diff --git a/src/crewai/agents/agent_builder/base_agent.py b/src/crewai/agents/agent_builder/base_agent.py index 13bbbb6f5..b01d40a19 100644 --- a/src/crewai/agents/agent_builder/base_agent.py +++ b/src/crewai/agents/agent_builder/base_agent.py @@ -270,15 +270,18 @@ class BaseAgent(ABC, BaseModel): "cache_handler", "llm", "knowledge_sources", + "knowledge_storage", "knowledge", - "formatting_errors", } # Copy llm existing_llm = shallow_copy(self.llm) + print("self.knowledge", self.knowledge) copied_knowledge = shallow_copy(self.knowledge) - + copied_knowledge_storage = shallow_copy(self.knowledge_storage) print("existing_llm", existing_llm) + print("copied_knowledge_storage", copied_knowledge_storage) + print("copied_knowledge", copied_knowledge) # Properly copy knowledge sources if they exist existing_knowledge_sources = None if self.knowledge_sources: @@ -308,6 +311,7 @@ class BaseAgent(ABC, BaseModel): tools=self.tools, knowledge_sources=existing_knowledge_sources, knowledge=copied_knowledge, + knowledge_storage=copied_knowledge_storage, ) return copied_agent diff --git a/tests/agent_test.py b/tests/agent_test.py index 2d9a23238..4ccbb9d03 100644 --- a/tests/agent_test.py +++ b/tests/agent_test.py @@ -12,6 +12,8 @@ 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.source.base_knowledge_source import BaseKnowledgeSource +from crewai.knowledge.storage.knowledge_storage import KnowledgeStorage from crewai.llm import LLM from crewai.tools import tool from crewai.tools.tool_calling import InstructorToolCalling @@ -1609,53 +1611,45 @@ def test_agent_with_knowledge_sources_works_with_copy(): content = "Brandon's favorite color is red and he likes Mexican food." string_source = StringKnowledgeSource(content=content) + # Create a real instance to use as spec + with patch( - "crewai.knowledge.storage.knowledge_storage.KnowledgeStorage" - ) as MockKnowledge: - # Set up the mock instance - mock_knowledge_instance = MockKnowledge.return_value - mock_knowledge_instance.sources = [string_source] - mock_knowledge_instance.query.return_value = [{"content": content}] - mock_knowledge_instance.save.return_value = None - mock_knowledge_instance.initialize_knowledge_storage.return_value = None - mock_knowledge_instance.collection_name = "test_collection" - # mock_knowledge_instance.embedder = { - # "provider": "openai", - # "config": {"model_name": "text-embedding-3-small", "api_key": "123"}, - # } + "crewai.knowledge.source.base_knowledge_source.BaseKnowledgeSource", + autospec=True, + ) as MockKnowledgeSource: + # Set up the mock instance using the real class as spec + mock_knowledge_source_instance = MockKnowledgeSource.return_value + mock_knowledge_source_instance.__class__ = BaseKnowledgeSource + mock_knowledge_source_instance.sources = [string_source] + # mock_knowledge_source_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], - # embedder={ - # "provider": "openai", - # "config": {"model_name": "text-embedding-3-small", "api_key": "123"}, - # }, + knowledge_sources=[ + string_source + ], # Use the real string source instead of mock ) - agent.knowledge_sources = [string_source] - # Actually call copy instead of mocking it - agent_copy = agent.copy() - print("agent_copy", agent_copy) - assert agent_copy.role == agent.role + # Mock the knowledge attribute to avoid validation issues during copy + with patch( + "crewai.knowledge.storage.knowledge_storage.KnowledgeStorage" + ) as MockKnowledgeStorage: + mock_knowledge_storage = MockKnowledgeStorage.return_value + agent.knowledge_storage = mock_knowledge_storage - # if agent.knowledge_sources and agent_copy.knowledge_sources: - # assert len(agent_copy.knowledge_sources) == len(agent.knowledge_sources) - # assert ( - # agent_copy.knowledge_sources[0].content # type: ignore - # == agent.knowledge_sources[0].content # type: ignore - # ) - # assert ( - # agent_copy.knowledge_sources[0].chunk_size - # == agent.knowledge_sources[0].chunk_size - # ) - # assert ( - # agent_copy.knowledge_sources[0].chunk_overlap - # == agent.knowledge_sources[0].chunk_overlap - # ) + agent_copy = agent.copy() + + # Basic assertions + assert agent_copy.role == agent.role + assert agent_copy.goal == agent.goal + assert agent_copy.backstory == agent.backstory + assert len(agent_copy.knowledge_sources) == 1 + assert isinstance(agent_copy.knowledge_sources[0], StringKnowledgeSource) + assert agent_copy.knowledge_sources[0].content == content + assert isinstance(agent_copy.llm, LLM) @pytest.mark.vcr(filter_headers=["authorization"])