Compare commits

...

2 Commits

Author SHA1 Message Date
Devin AI
ac7390b287 Fix linting issues in context-aware knowledge search tests
- Remove unused imports: MagicMock, Knowledge, KnowledgeConfig
- Remove unused result variables in test functions
- Addresses lint failures in CI

Co-Authored-By: João <joao@crewai.com>
2025-08-17 17:31:16 +00:00
Devin AI
6f139dff06 Implement context-aware knowledge search feature
- Add optional context parameter to _get_knowledge_search_query method
- Update knowledge search query generation to include context from previous tasks
- Add new prompt template for context-aware knowledge queries
- Maintain backward compatibility with existing knowledge search
- Add comprehensive tests for context-aware functionality
- Addresses issue #3332

Co-Authored-By: João <joao@crewai.com>
2025-08-17 17:27:49 +00:00
4 changed files with 266 additions and 6 deletions

View File

@@ -346,7 +346,7 @@ class Agent(BaseAgent):
)
try:
self.knowledge_search_query = self._get_knowledge_search_query(
task_prompt
task_prompt, context
)
if self.knowledge_search_query:
# Quering agent specific knowledge
@@ -722,8 +722,8 @@ class Agent(BaseAgent):
def set_fingerprint(self, fingerprint: Fingerprint):
self.security_config.fingerprint = fingerprint
def _get_knowledge_search_query(self, task_prompt: str) -> str | None:
"""Generate a search query for the knowledge base based on the task description."""
def _get_knowledge_search_query(self, task_prompt: str, context: Optional[str] = None) -> str | None:
"""Generate a search query for the knowledge base based on the task description and context."""
crewai_event_bus.emit(
self,
event=KnowledgeQueryStartedEvent(
@@ -731,9 +731,16 @@ class Agent(BaseAgent):
agent=self,
),
)
query = self.i18n.slice("knowledge_search_query").format(
task_prompt=task_prompt
)
if context:
query = self.i18n.slice("knowledge_search_query_with_context").format(
task_prompt=task_prompt, context=context
)
else:
query = self.i18n.slice("knowledge_search_query").format(
task_prompt=task_prompt
)
rewriter_prompt = self.i18n.slice("knowledge_search_query_system_prompt")
if not isinstance(self.llm, BaseLLM):
self._logger.log(

View File

@@ -29,6 +29,7 @@
"lite_agent_system_prompt_without_tools": "You are {role}. {backstory}\nYour personal goal is: {goal}\n\nTo give my best complete final answer to the task respond using the exact following format:\n\nThought: I now can give a great answer\nFinal Answer: Your final answer must be the great and the most complete as possible, it must be outcome described.\n\nI MUST use these formats, my job depends on it!",
"lite_agent_response_format": "\nIMPORTANT: Your final answer MUST contain all the information requested in the following format: {response_format}\n\nIMPORTANT: Ensure the final output does not include any code block markers like ```json or ```python.",
"knowledge_search_query": "The original query is: {task_prompt}.",
"knowledge_search_query_with_context": "The original query is: {task_prompt}.\n\nContext from previous tasks:\n{context}",
"knowledge_search_query_system_prompt": "Your goal is to rewrite the user query so that it is optimized for retrieval from a vector database. Consider how the query will be used to find relevant documents, and aim to make it more specific and context-aware. \n\n Do not include any other text than the rewritten query, especially any preamble or postamble and only add expected output format if its relevant to the rewritten query. \n\n Focus on the key words of the intended task and to retrieve the most relevant information. \n\n There will be some extra context provided that might need to be removed such as expected_output formats structured_outputs and other instructions."
},
"errors": {

View File

@@ -1756,6 +1756,35 @@ def test_agent_with_knowledge_sources_with_query_limit_and_score_threshold():
)
@pytest.mark.vcr(filter_headers=["authorization"])
def test_agent_knowledge_search_with_context_parameter():
"""Test that agent knowledge search accepts context parameter."""
content = "Brandon's favorite color is red and he likes Mexican food."
string_source = StringKnowledgeSource(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],
)
task_prompt = "What is Brandon's favorite color?"
context = "Previous conversation mentioned Brandon's preferences."
with patch.object(agent.llm, 'call') as mock_call:
mock_call.return_value = "Brandon likes red"
result = agent._get_knowledge_search_query(task_prompt, context)
assert result == "Brandon likes red"
mock_call.assert_called_once()
call_args = mock_call.call_args[0][0]
user_message = call_args[1]['content']
assert context in user_message
@pytest.mark.vcr(filter_headers=["authorization"])
def test_agent_with_knowledge_sources_with_query_limit_and_score_threshold_default():
content = "Brandon's favorite color is red and he likes Mexican food."

View File

@@ -0,0 +1,223 @@
"""Test context-aware knowledge search functionality."""
import pytest
from unittest.mock import patch
from crewai import Agent, Task, Crew, LLM
from crewai.knowledge.source.string_knowledge_source import StringKnowledgeSource
@pytest.mark.vcr(filter_headers=["authorization"])
def test_knowledge_search_with_context():
"""Test that knowledge search includes context from previous tasks."""
content = "The company's main product is a CRM system. The system has three modules: Sales, Marketing, and Support."
string_source = StringKnowledgeSource(content=content)
researcher = Agent(
role="Research Analyst",
goal="Research company information",
backstory="You are a research analyst.",
llm=LLM(model="gpt-4o-mini"),
knowledge_sources=[string_source],
)
writer = Agent(
role="Content Writer",
goal="Write content based on research",
backstory="You are a content writer.",
llm=LLM(model="gpt-4o-mini"),
knowledge_sources=[string_source],
)
research_task = Task(
description="Research the company's main product",
expected_output="A summary of the company's main product",
agent=researcher,
)
writing_task = Task(
description="Write a detailed description of the CRM modules",
expected_output="A detailed description of each CRM module",
agent=writer,
context=[research_task],
)
crew = Crew(agents=[researcher, writer], tasks=[research_task, writing_task])
with patch.object(writer, '_get_knowledge_search_query') as mock_search:
mock_search.return_value = "mocked query"
crew.kickoff()
mock_search.assert_called_once()
call_args = mock_search.call_args
assert len(call_args[0]) == 2
assert call_args[0][1] is not None
assert "CRM system" in call_args[0][1] or "product" in call_args[0][1]
@pytest.mark.vcr(filter_headers=["authorization"])
def test_knowledge_search_without_context():
"""Test that knowledge search works without context (backward compatibility)."""
content = "The company's main product is a CRM system."
string_source = StringKnowledgeSource(content=content)
agent = Agent(
role="Research Analyst",
goal="Research company information",
backstory="You are a research analyst.",
llm=LLM(model="gpt-4o-mini"),
knowledge_sources=[string_source],
)
task = Task(
description="Research the company's main product",
expected_output="A summary of the company's main product",
agent=agent,
)
crew = Crew(agents=[agent], tasks=[task])
with patch.object(agent, '_get_knowledge_search_query') as mock_search:
mock_search.return_value = "mocked query"
crew.kickoff()
mock_search.assert_called_once()
call_args = mock_search.call_args
assert len(call_args[0]) == 2
assert call_args[0][1] == ""
@pytest.mark.vcr(filter_headers=["authorization"])
def test_context_aware_knowledge_search_integration():
"""Integration test for context-aware knowledge search."""
knowledge_content = """
Project Alpha is a web application built with React and Node.js.
Project Beta is a mobile application built with React Native.
The team uses Agile methodology with 2-week sprints.
The database is PostgreSQL with Redis for caching.
"""
string_source = StringKnowledgeSource(content=knowledge_content)
project_manager = Agent(
role="Project Manager",
goal="Gather project information",
backstory="You manage software projects.",
llm=LLM(model="gpt-4o-mini"),
knowledge_sources=[string_source],
)
tech_lead = Agent(
role="Technical Lead",
goal="Provide technical details",
backstory="You are a technical expert.",
llm=LLM(model="gpt-4o-mini"),
knowledge_sources=[string_source],
)
project_overview_task = Task(
description="Provide an overview of Project Alpha",
expected_output="Overview of Project Alpha including its technology stack",
agent=project_manager,
)
technical_details_task = Task(
description="Provide technical implementation details for the project",
expected_output="Technical implementation details including database and caching",
agent=tech_lead,
context=[project_overview_task],
)
crew = Crew(agents=[project_manager, tech_lead], tasks=[project_overview_task, technical_details_task])
result = crew.kickoff()
assert result.raw is not None
assert any(keyword in result.raw.lower() for keyword in ["react", "node", "postgresql", "redis"])
def test_knowledge_search_query_template_with_context():
"""Test that the knowledge search query template includes context properly."""
agent = Agent(
role="Test Agent",
goal="Test knowledge search",
backstory="Test agent",
llm=LLM(model="gpt-4o-mini"),
)
task_prompt = "What is the main product?"
context = "Previous research shows the company focuses on CRM solutions."
with patch.object(agent.llm, 'call') as mock_call:
mock_call.return_value = "mocked response"
agent._get_knowledge_search_query(task_prompt, context)
mock_call.assert_called_once()
call_args = mock_call.call_args[0][0]
user_message = call_args[1]['content']
assert task_prompt in user_message
assert context in user_message
assert "Context from previous tasks:" in user_message
def test_knowledge_search_query_template_without_context():
"""Test that the knowledge search query template works without context."""
agent = Agent(
role="Test Agent",
goal="Test knowledge search",
backstory="Test agent",
llm=LLM(model="gpt-4o-mini"),
)
task_prompt = "What is the main product?"
with patch.object(agent.llm, 'call') as mock_call:
mock_call.return_value = "mocked response"
agent._get_knowledge_search_query(task_prompt)
mock_call.assert_called_once()
call_args = mock_call.call_args[0][0]
user_message = call_args[1]['content']
assert task_prompt in user_message
assert "Context from previous tasks:" not in user_message
def test_structured_context_integration():
"""Test context-aware knowledge search with structured context data."""
knowledge_content = """
Error URS-01: User registration service unavailable.
Method getUserStatus returns user account status.
API endpoint /api/users/{id}/status for user status queries.
Database table user_accounts stores user information.
"""
string_source = StringKnowledgeSource(content=knowledge_content)
agent = Agent(
role="Technical Support",
goal="Resolve technical issues",
backstory="You help resolve technical problems.",
llm=LLM(model="gpt-4o-mini"),
knowledge_sources=[string_source],
)
task_prompt = "How to resolve the user status error?"
structured_context = '{"method": "getUserStatus", "error_code": "URS-01", "endpoint": "/api/users/{id}/status"}'
with patch.object(agent.llm, 'call') as mock_call:
mock_call.return_value = "Check getUserStatus method and URS-01 error"
agent._get_knowledge_search_query(task_prompt, structured_context)
mock_call.assert_called_once()
call_args = mock_call.call_args[0][0]
user_message = call_args[1]['content']
assert task_prompt in user_message
assert "getUserStatus" in user_message
assert "URS-01" in user_message
assert "Context from previous tasks:" in user_message