From 1c2238066f78652226eb850d8ee2ec6b1f55603f Mon Sep 17 00:00:00 2001 From: lorenzejay Date: Tue, 6 May 2025 20:26:19 -0700 Subject: [PATCH] refactor: rename knowledge query events and enhance retrieval handling This commit renames the KnowledgeQueryGeneratedEvent to KnowledgeQueryStartedEvent to better reflect its purpose. It also updates the event handling in the EventListener and ConsoleFormatter classes to accommodate the new event structure. Additionally, the retrieval knowledge is now included in the KnowledgeRetrievalCompletedEvent, improving the overall knowledge retrieval process. --- src/crewai/agent.py | 16 ++++- src/crewai/utilities/events/event_listener.py | 29 ++++---- .../utilities/events/knowledge_events.py | 11 +-- .../events/utils/console_formatter.py | 71 ++++++++++++++++--- 4 files changed, 99 insertions(+), 28 deletions(-) diff --git a/src/crewai/agent.py b/src/crewai/agent.py index 016ca17d7..0fb3994f3 100644 --- a/src/crewai/agent.py +++ b/src/crewai/agent.py @@ -34,7 +34,7 @@ from crewai.utilities.events.crewai_event_bus import crewai_event_bus from crewai.utilities.events.knowledge_events import ( KnowledgeQueryCompletedEvent, KnowledgeQueryFailedEvent, - KnowledgeQueryGeneratedEvent, + KnowledgeQueryStartedEvent, KnowledgeRetrievalCompletedEvent, KnowledgeRetrievalStartedEvent, KnowledgeSearchQueryFailedEvent, @@ -295,6 +295,16 @@ class Agent(BaseAgent): event=KnowledgeRetrievalCompletedEvent( query=self.knowledge_search_query, agent=self, + retrieved_knowledge=( + (self.agent_knowledge_context or "") + + ( + "\n" + if self.agent_knowledge_context + and self.crew_knowledge_context + else "" + ) + + (self.crew_knowledge_context or "") + ), ), ) except Exception as e: @@ -603,7 +613,7 @@ class Agent(BaseAgent): """Generate a search query for the knowledge base based on the task description.""" crewai_event_bus.emit( self, - event=KnowledgeQueryGeneratedEvent( + event=KnowledgeQueryStartedEvent( task_prompt=task_prompt, agent=self, ), @@ -635,7 +645,7 @@ class Agent(BaseAgent): crewai_event_bus.emit( self, event=KnowledgeQueryFailedEvent( - query=query, + task_prompt=task_prompt, agent=self, error=str(e), ), diff --git a/src/crewai/utilities/events/event_listener.py b/src/crewai/utilities/events/event_listener.py index bc16f9eec..a76b87964 100644 --- a/src/crewai/utilities/events/event_listener.py +++ b/src/crewai/utilities/events/event_listener.py @@ -11,7 +11,7 @@ from crewai.utilities.events.base_event_listener import BaseEventListener from crewai.utilities.events.knowledge_events import ( KnowledgeQueryCompletedEvent, KnowledgeQueryFailedEvent, - KnowledgeQueryGeneratedEvent, + KnowledgeQueryStartedEvent, KnowledgeRetrievalCompletedEvent, KnowledgeRetrievalStartedEvent, KnowledgeSearchQueryFailedEvent, @@ -65,6 +65,8 @@ class EventListener(BaseEventListener): execution_spans: Dict[Task, Any] = Field(default_factory=dict) next_chunk = 0 text_stream = StringIO() + knowledge_retrieval_in_progress = False + knowledge_query_in_progress = False def __new__(cls): if cls._instance is None: @@ -354,6 +356,11 @@ class EventListener(BaseEventListener): def on_knowledge_retrieval_started( source, event: KnowledgeRetrievalStartedEvent ): + if self.knowledge_retrieval_in_progress: + return + + self.knowledge_retrieval_in_progress = True + self.formatter.handle_knowledge_retrieval_started( self.formatter.current_agent_branch, self.formatter.current_crew_tree, @@ -363,18 +370,19 @@ class EventListener(BaseEventListener): def on_knowledge_retrieval_completed( source, event: KnowledgeRetrievalCompletedEvent ): + if not self.knowledge_retrieval_in_progress: + return + + self.knowledge_retrieval_in_progress = False self.formatter.handle_knowledge_retrieval_completed( self.formatter.current_agent_branch, self.formatter.current_crew_tree, + event.retrieved_knowledge, ) - @crewai_event_bus.on(KnowledgeQueryGeneratedEvent) - def on_knowledge_query_generated(source, event: KnowledgeQueryGeneratedEvent): - self.formatter.handle_knowledge_query_generated( - self.formatter.current_agent_branch, - event.task_prompt, - self.formatter.current_crew_tree, - ) + @crewai_event_bus.on(KnowledgeQueryStartedEvent) + def on_knowledge_query_started(source, event: KnowledgeQueryStartedEvent): + pass @crewai_event_bus.on(KnowledgeQueryFailedEvent) def on_knowledge_query_failed(source, event: KnowledgeQueryFailedEvent): @@ -386,10 +394,7 @@ class EventListener(BaseEventListener): @crewai_event_bus.on(KnowledgeQueryCompletedEvent) def on_knowledge_query_completed(source, event: KnowledgeQueryCompletedEvent): - self.formatter.handle_knowledge_query_completed( - self.formatter.current_agent_branch, - self.formatter.current_crew_tree, - ) + pass @crewai_event_bus.on(KnowledgeSearchQueryFailedEvent) def on_knowledge_search_query_failed( diff --git a/src/crewai/utilities/events/knowledge_events.py b/src/crewai/utilities/events/knowledge_events.py index 16a932dce..3913e9e64 100644 --- a/src/crewai/utilities/events/knowledge_events.py +++ b/src/crewai/utilities/events/knowledge_events.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from crewai.agents.agent_builder.base_agent import BaseAgent from crewai.utilities.events.base_events import BaseEvent @@ -20,20 +20,21 @@ class KnowledgeRetrievalCompletedEvent(BaseEvent): query: str type: str = "knowledge_search_query_completed" agent: BaseAgent + retrieved_knowledge: Any -class KnowledgeQueryGeneratedEvent(BaseEvent): - """Event emitted when a knowledge query is generated.""" +class KnowledgeQueryStartedEvent(BaseEvent): + """Event emitted when a knowledge query is started.""" task_prompt: str - type: str = "knowledge_query_generated" + type: str = "knowledge_query_started" agent: BaseAgent class KnowledgeQueryFailedEvent(BaseEvent): """Event emitted when a knowledge query fails.""" - query: str + task_prompt: str type: str = "knowledge_query_failed" agent: BaseAgent error: str diff --git a/src/crewai/utilities/events/utils/console_formatter.py b/src/crewai/utilities/events/utils/console_formatter.py index f2b15d9ea..b9adc9fda 100644 --- a/src/crewai/utilities/events/utils/console_formatter.py +++ b/src/crewai/utilities/events/utils/console_formatter.py @@ -816,6 +816,7 @@ class ConsoleFormatter: self, agent_branch: Optional[Tree], crew_tree: Optional[Tree], + retrieved_knowledge: Any, ) -> None: """Handle knowledge retrieval completed event.""" if not self.verbose: @@ -824,21 +825,74 @@ class ConsoleFormatter: branch_to_use = self.current_lite_agent_branch or agent_branch tree_to_use = branch_to_use or crew_tree + if branch_to_use is None and tree_to_use is not None: + branch_to_use = tree_to_use + if branch_to_use is None or tree_to_use is None: + if retrieved_knowledge: + knowledge_text = str(retrieved_knowledge) + if len(knowledge_text) > 500: + knowledge_text = knowledge_text[:497] + "..." + + knowledge_panel = Panel( + Text(knowledge_text, style="white"), + title="📚 Retrieved Knowledge", + border_style="green", + padding=(1, 2), + ) + self.print(knowledge_panel) + self.print() return None - knowledge_branch = branch_to_use.add("") - self.update_tree_label( - knowledge_branch, "✅", "Knowledge Retrieval Completed", "green" - ) + knowledge_branch_found = False + for child in branch_to_use.children: + if "Knowledge Retrieval Started" in str(child.label): + self.update_tree_label( + child, "✅", "Knowledge Retrieval Completed", "green" + ) + knowledge_branch_found = True + break + + if not knowledge_branch_found: + for child in branch_to_use.children: + if ( + "Knowledge Retrieval" in str(child.label) + and "Started" not in str(child.label) + and "Completed" not in str(child.label) + ): + self.update_tree_label( + child, "✅", "Knowledge Retrieval Completed", "green" + ) + knowledge_branch_found = True + break + + if not knowledge_branch_found: + knowledge_branch = branch_to_use.add("") + self.update_tree_label( + knowledge_branch, "✅", "Knowledge Retrieval Completed", "green" + ) self.print(tree_to_use) + + if retrieved_knowledge: + knowledge_text = str(retrieved_knowledge) + if len(knowledge_text) > 500: + knowledge_text = knowledge_text[:497] + "..." + + knowledge_panel = Panel( + Text(knowledge_text, style="white"), + title="📚 Retrieved Knowledge", + border_style="green", + padding=(1, 2), + ) + self.print(knowledge_panel) + self.print() - def handle_knowledge_query_generated( + def handle_knowledge_query_started( self, agent_branch: Optional[Tree], - query: str, + task_prompt: str, crew_tree: Optional[Tree], ) -> None: """Handle knowledge query generated event.""" @@ -847,12 +901,13 @@ class ConsoleFormatter: branch_to_use = self.current_lite_agent_branch or agent_branch tree_to_use = branch_to_use or crew_tree - if branch_to_use is None or tree_to_use is None: return None query_branch = branch_to_use.add("") - self.update_tree_label(query_branch, "🔎", f"Query: {query[:50]}...", "yellow") + self.update_tree_label( + query_branch, "🔎", f"Query: {task_prompt[:50]}...", "yellow" + ) self.print(tree_to_use) self.print()