mirror of
https://github.com/crewAIInc/crewAI.git
synced 2025-12-16 12:28:30 +00:00
Compare commits
4 Commits
devin/1765
...
devin/1757
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
094b3fe660 | ||
|
|
7ff4062d4a | ||
|
|
e2f4ea3956 | ||
|
|
6fda55518d |
@@ -1,25 +1,25 @@
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any, ClassVar
|
||||
|
||||
from rich.console import Console
|
||||
from rich.live import Live
|
||||
from rich.panel import Panel
|
||||
from rich.syntax import Syntax
|
||||
from rich.text import Text
|
||||
from rich.tree import Tree
|
||||
from rich.live import Live
|
||||
from rich.syntax import Syntax
|
||||
|
||||
|
||||
class ConsoleFormatter:
|
||||
current_crew_tree: Optional[Tree] = None
|
||||
current_task_branch: Optional[Tree] = None
|
||||
current_agent_branch: Optional[Tree] = None
|
||||
current_tool_branch: Optional[Tree] = None
|
||||
current_flow_tree: Optional[Tree] = None
|
||||
current_method_branch: Optional[Tree] = None
|
||||
current_lite_agent_branch: Optional[Tree] = None
|
||||
tool_usage_counts: Dict[str, int] = {}
|
||||
current_reasoning_branch: Optional[Tree] = None # Track reasoning status
|
||||
current_crew_tree: Tree | None = None
|
||||
current_task_branch: Tree | None = None
|
||||
current_agent_branch: Tree | None = None
|
||||
current_tool_branch: Tree | None = None
|
||||
current_flow_tree: Tree | None = None
|
||||
current_method_branch: Tree | None = None
|
||||
current_lite_agent_branch: Tree | None = None
|
||||
tool_usage_counts: ClassVar[dict[str, int]] = {}
|
||||
current_reasoning_branch: Tree | None = None # Track reasoning status
|
||||
_live_paused: bool = False
|
||||
current_llm_tool_tree: Optional[Tree] = None
|
||||
current_llm_tool_tree: Tree | None = None
|
||||
|
||||
def __init__(self, verbose: bool = False):
|
||||
self.console = Console(width=None)
|
||||
@@ -29,7 +29,7 @@ class ConsoleFormatter:
|
||||
# instance so the previous render is replaced instead of writing a new one.
|
||||
# Once any non-Tree renderable is printed we stop the Live session so the
|
||||
# final Tree persists on the terminal.
|
||||
self._live: Optional[Live] = None
|
||||
self._live: Live | None = None
|
||||
|
||||
def create_panel(self, content: Text, title: str, style: str = "blue") -> Panel:
|
||||
"""Create a standardized panel with consistent styling."""
|
||||
@@ -45,7 +45,7 @@ class ConsoleFormatter:
|
||||
title: str,
|
||||
name: str,
|
||||
status_style: str = "blue",
|
||||
tool_args: Dict[str, Any] | str = "",
|
||||
tool_args: dict[str, Any] | str = "",
|
||||
**fields,
|
||||
) -> Text:
|
||||
"""Create standardized status content with consistent formatting."""
|
||||
@@ -70,7 +70,7 @@ class ConsoleFormatter:
|
||||
prefix: str,
|
||||
name: str,
|
||||
style: str = "blue",
|
||||
status: Optional[str] = None,
|
||||
status: str | None = None,
|
||||
) -> None:
|
||||
"""Update tree label with consistent formatting."""
|
||||
label = Text()
|
||||
@@ -115,7 +115,7 @@ class ConsoleFormatter:
|
||||
self._live.update(tree, refresh=True)
|
||||
return # Nothing else to do
|
||||
|
||||
# Case 2: blank line while a live session is running – ignore so we
|
||||
# Case 2: blank line while a live session is running - ignore so we
|
||||
# don't break the in-place rendering behaviour
|
||||
if len(args) == 0 and self._live:
|
||||
return
|
||||
@@ -156,7 +156,7 @@ class ConsoleFormatter:
|
||||
|
||||
def update_crew_tree(
|
||||
self,
|
||||
tree: Optional[Tree],
|
||||
tree: Tree | None,
|
||||
crew_name: str,
|
||||
source_id: str,
|
||||
status: str = "completed",
|
||||
@@ -196,7 +196,7 @@ class ConsoleFormatter:
|
||||
|
||||
self.print_panel(content, title, style)
|
||||
|
||||
def create_crew_tree(self, crew_name: str, source_id: str) -> Optional[Tree]:
|
||||
def create_crew_tree(self, crew_name: str, source_id: str) -> Tree | None:
|
||||
"""Create and initialize a new crew tree with initial status."""
|
||||
if not self.verbose:
|
||||
return None
|
||||
@@ -220,8 +220,8 @@ class ConsoleFormatter:
|
||||
return tree
|
||||
|
||||
def create_task_branch(
|
||||
self, crew_tree: Optional[Tree], task_id: str, task_name: Optional[str] = None
|
||||
) -> Optional[Tree]:
|
||||
self, crew_tree: Tree | None, task_id: str, task_name: str | None = None
|
||||
) -> Tree | None:
|
||||
"""Create and initialize a task branch."""
|
||||
if not self.verbose:
|
||||
return None
|
||||
@@ -255,11 +255,11 @@ class ConsoleFormatter:
|
||||
|
||||
def update_task_status(
|
||||
self,
|
||||
crew_tree: Optional[Tree],
|
||||
crew_tree: Tree | None,
|
||||
task_id: str,
|
||||
agent_role: str,
|
||||
status: str = "completed",
|
||||
task_name: Optional[str] = None,
|
||||
task_name: str | None = None,
|
||||
) -> None:
|
||||
"""Update task status in the tree."""
|
||||
if not self.verbose or crew_tree is None:
|
||||
@@ -306,8 +306,8 @@ class ConsoleFormatter:
|
||||
self.print_panel(content, panel_title, style)
|
||||
|
||||
def create_agent_branch(
|
||||
self, task_branch: Optional[Tree], agent_role: str, crew_tree: Optional[Tree]
|
||||
) -> Optional[Tree]:
|
||||
self, task_branch: Tree | None, agent_role: str, crew_tree: Tree | None
|
||||
) -> Tree | None:
|
||||
"""Create and initialize an agent branch."""
|
||||
if not self.verbose or not task_branch or not crew_tree:
|
||||
return None
|
||||
@@ -325,9 +325,9 @@ class ConsoleFormatter:
|
||||
|
||||
def update_agent_status(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
agent_branch: Tree | None,
|
||||
agent_role: str,
|
||||
crew_tree: Optional[Tree],
|
||||
crew_tree: Tree | None,
|
||||
status: str = "completed",
|
||||
) -> None:
|
||||
"""Update agent status in the tree."""
|
||||
@@ -336,7 +336,7 @@ class ConsoleFormatter:
|
||||
# altering the tree. Keeping it a no-op avoids duplicate status lines.
|
||||
return
|
||||
|
||||
def create_flow_tree(self, flow_name: str, flow_id: str) -> Optional[Tree]:
|
||||
def create_flow_tree(self, flow_name: str, flow_id: str) -> Tree | None:
|
||||
"""Create and initialize a flow tree."""
|
||||
content = self.create_status_content(
|
||||
"Starting Flow Execution", flow_name, "blue", ID=flow_id
|
||||
@@ -356,7 +356,7 @@ class ConsoleFormatter:
|
||||
|
||||
return flow_tree
|
||||
|
||||
def start_flow(self, flow_name: str, flow_id: str) -> Optional[Tree]:
|
||||
def start_flow(self, flow_name: str, flow_id: str) -> Tree | None:
|
||||
"""Initialize a flow execution tree."""
|
||||
flow_tree = Tree("")
|
||||
flow_label = Text()
|
||||
@@ -376,7 +376,7 @@ class ConsoleFormatter:
|
||||
|
||||
def update_flow_status(
|
||||
self,
|
||||
flow_tree: Optional[Tree],
|
||||
flow_tree: Tree | None,
|
||||
flow_name: str,
|
||||
flow_id: str,
|
||||
status: str = "completed",
|
||||
@@ -423,11 +423,11 @@ class ConsoleFormatter:
|
||||
|
||||
def update_method_status(
|
||||
self,
|
||||
method_branch: Optional[Tree],
|
||||
flow_tree: Optional[Tree],
|
||||
method_branch: Tree | None,
|
||||
flow_tree: Tree | None,
|
||||
method_name: str,
|
||||
status: str = "running",
|
||||
) -> Optional[Tree]:
|
||||
) -> Tree | None:
|
||||
"""Update method status in the flow tree."""
|
||||
if not flow_tree:
|
||||
return None
|
||||
@@ -480,7 +480,7 @@ class ConsoleFormatter:
|
||||
def handle_llm_tool_usage_started(
|
||||
self,
|
||||
tool_name: str,
|
||||
tool_args: Dict[str, Any] | str,
|
||||
tool_args: dict[str, Any] | str,
|
||||
):
|
||||
# Create status content for the tool usage
|
||||
content = self.create_status_content(
|
||||
@@ -520,11 +520,11 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_tool_usage_started(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
agent_branch: Tree | None,
|
||||
tool_name: str,
|
||||
crew_tree: Optional[Tree],
|
||||
tool_args: Dict[str, Any] | str = "",
|
||||
) -> Optional[Tree]:
|
||||
crew_tree: Tree | None,
|
||||
tool_args: dict[str, Any] | str = "",
|
||||
) -> Tree | None:
|
||||
"""Handle tool usage started event."""
|
||||
if not self.verbose:
|
||||
return None
|
||||
@@ -569,9 +569,9 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_tool_usage_finished(
|
||||
self,
|
||||
tool_branch: Optional[Tree],
|
||||
tool_branch: Tree | None,
|
||||
tool_name: str,
|
||||
crew_tree: Optional[Tree],
|
||||
crew_tree: Tree | None,
|
||||
) -> None:
|
||||
"""Handle tool usage finished event."""
|
||||
if not self.verbose or tool_branch is None:
|
||||
@@ -600,10 +600,10 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_tool_usage_error(
|
||||
self,
|
||||
tool_branch: Optional[Tree],
|
||||
tool_branch: Tree | None,
|
||||
tool_name: str,
|
||||
error: str,
|
||||
crew_tree: Optional[Tree],
|
||||
crew_tree: Tree | None,
|
||||
) -> None:
|
||||
"""Handle tool usage error event."""
|
||||
if not self.verbose:
|
||||
@@ -631,9 +631,9 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_llm_call_started(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
crew_tree: Optional[Tree],
|
||||
) -> Optional[Tree]:
|
||||
agent_branch: Tree | None,
|
||||
crew_tree: Tree | None,
|
||||
) -> Tree | None:
|
||||
"""Handle LLM call started event."""
|
||||
if not self.verbose:
|
||||
return None
|
||||
@@ -672,9 +672,9 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_llm_call_completed(
|
||||
self,
|
||||
tool_branch: Optional[Tree],
|
||||
agent_branch: Optional[Tree],
|
||||
crew_tree: Optional[Tree],
|
||||
tool_branch: Tree | None,
|
||||
agent_branch: Tree | None,
|
||||
crew_tree: Tree | None,
|
||||
) -> None:
|
||||
"""Handle LLM call completed event."""
|
||||
if not self.verbose:
|
||||
@@ -736,7 +736,7 @@ class ConsoleFormatter:
|
||||
self.print()
|
||||
|
||||
def handle_llm_call_failed(
|
||||
self, tool_branch: Optional[Tree], error: str, crew_tree: Optional[Tree]
|
||||
self, tool_branch: Tree | None, error: str, crew_tree: Tree | None
|
||||
) -> None:
|
||||
"""Handle LLM call failed event."""
|
||||
if not self.verbose:
|
||||
@@ -789,7 +789,7 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_crew_test_started(
|
||||
self, crew_name: str, source_id: str, n_iterations: int
|
||||
) -> Optional[Tree]:
|
||||
) -> Tree | None:
|
||||
"""Handle crew test started event."""
|
||||
if not self.verbose:
|
||||
return None
|
||||
@@ -823,7 +823,7 @@ class ConsoleFormatter:
|
||||
return test_tree
|
||||
|
||||
def handle_crew_test_completed(
|
||||
self, flow_tree: Optional[Tree], crew_name: str
|
||||
self, flow_tree: Tree | None, crew_name: str
|
||||
) -> None:
|
||||
"""Handle crew test completed event."""
|
||||
if not self.verbose:
|
||||
@@ -913,7 +913,7 @@ class ConsoleFormatter:
|
||||
self.print_panel(failure_content, "Test Failure", "red")
|
||||
self.print()
|
||||
|
||||
def create_lite_agent_branch(self, lite_agent_role: str) -> Optional[Tree]:
|
||||
def create_lite_agent_branch(self, lite_agent_role: str) -> Tree | None:
|
||||
"""Create and initialize a lite agent branch."""
|
||||
if not self.verbose:
|
||||
return None
|
||||
@@ -935,10 +935,10 @@ class ConsoleFormatter:
|
||||
|
||||
def update_lite_agent_status(
|
||||
self,
|
||||
lite_agent_branch: Optional[Tree],
|
||||
lite_agent_branch: Tree | None,
|
||||
lite_agent_role: str,
|
||||
status: str = "completed",
|
||||
**fields: Dict[str, Any],
|
||||
**fields: dict[str, Any],
|
||||
) -> None:
|
||||
"""Update lite agent status in the tree."""
|
||||
if not self.verbose or lite_agent_branch is None:
|
||||
@@ -981,7 +981,7 @@ class ConsoleFormatter:
|
||||
lite_agent_role: str,
|
||||
status: str = "started",
|
||||
error: Any = None,
|
||||
**fields: Dict[str, Any],
|
||||
**fields: dict[str, Any],
|
||||
) -> None:
|
||||
"""Handle lite agent execution events with consistent formatting."""
|
||||
if not self.verbose:
|
||||
@@ -1006,9 +1006,9 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_knowledge_retrieval_started(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
crew_tree: Optional[Tree],
|
||||
) -> Optional[Tree]:
|
||||
agent_branch: Tree | None,
|
||||
crew_tree: Tree | None,
|
||||
) -> Tree | None:
|
||||
"""Handle knowledge retrieval started event."""
|
||||
if not self.verbose:
|
||||
return None
|
||||
@@ -1034,13 +1034,13 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_knowledge_retrieval_completed(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
crew_tree: Optional[Tree],
|
||||
agent_branch: Tree | None,
|
||||
crew_tree: Tree | None,
|
||||
retrieved_knowledge: Any,
|
||||
) -> None:
|
||||
"""Handle knowledge retrieval completed event."""
|
||||
if not self.verbose:
|
||||
return None
|
||||
return
|
||||
|
||||
branch_to_use = self.current_lite_agent_branch or agent_branch
|
||||
tree_to_use = branch_to_use or crew_tree
|
||||
@@ -1062,7 +1062,7 @@ class ConsoleFormatter:
|
||||
)
|
||||
self.print(knowledge_panel)
|
||||
self.print()
|
||||
return None
|
||||
return
|
||||
|
||||
knowledge_branch_found = False
|
||||
for child in branch_to_use.children:
|
||||
@@ -1111,18 +1111,18 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_knowledge_query_started(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
agent_branch: Tree | None,
|
||||
task_prompt: str,
|
||||
crew_tree: Optional[Tree],
|
||||
crew_tree: Tree | None,
|
||||
) -> None:
|
||||
"""Handle knowledge query generated event."""
|
||||
if not self.verbose:
|
||||
return None
|
||||
return
|
||||
|
||||
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
|
||||
return
|
||||
|
||||
query_branch = branch_to_use.add("")
|
||||
self.update_tree_label(
|
||||
@@ -1134,9 +1134,9 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_knowledge_query_failed(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
agent_branch: Tree | None,
|
||||
error: str,
|
||||
crew_tree: Optional[Tree],
|
||||
crew_tree: Tree | None,
|
||||
) -> None:
|
||||
"""Handle knowledge query failed event."""
|
||||
if not self.verbose:
|
||||
@@ -1159,18 +1159,18 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_knowledge_query_completed(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
crew_tree: Optional[Tree],
|
||||
agent_branch: Tree | None,
|
||||
crew_tree: Tree | None,
|
||||
) -> None:
|
||||
"""Handle knowledge query completed event."""
|
||||
if not self.verbose:
|
||||
return None
|
||||
return
|
||||
|
||||
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
|
||||
return
|
||||
|
||||
query_branch = branch_to_use.add("")
|
||||
self.update_tree_label(query_branch, "✅", "Knowledge Query Completed", "green")
|
||||
@@ -1180,9 +1180,9 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_knowledge_search_query_failed(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
agent_branch: Tree | None,
|
||||
error: str,
|
||||
crew_tree: Optional[Tree],
|
||||
crew_tree: Tree | None,
|
||||
) -> None:
|
||||
"""Handle knowledge search query failed event."""
|
||||
if not self.verbose:
|
||||
@@ -1207,10 +1207,10 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_reasoning_started(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
agent_branch: Tree | None,
|
||||
attempt: int,
|
||||
crew_tree: Optional[Tree],
|
||||
) -> Optional[Tree]:
|
||||
crew_tree: Tree | None,
|
||||
) -> Tree | None:
|
||||
"""Handle agent reasoning started (or refinement) event."""
|
||||
if not self.verbose:
|
||||
return None
|
||||
@@ -1249,7 +1249,7 @@ class ConsoleFormatter:
|
||||
self,
|
||||
plan: str,
|
||||
ready: bool,
|
||||
crew_tree: Optional[Tree],
|
||||
crew_tree: Tree | None,
|
||||
) -> None:
|
||||
"""Handle agent reasoning completed event."""
|
||||
if not self.verbose:
|
||||
@@ -1292,7 +1292,7 @@ class ConsoleFormatter:
|
||||
def handle_reasoning_failed(
|
||||
self,
|
||||
error: str,
|
||||
crew_tree: Optional[Tree],
|
||||
crew_tree: Tree | None,
|
||||
) -> None:
|
||||
"""Handle agent reasoning failure event."""
|
||||
if not self.verbose:
|
||||
@@ -1329,7 +1329,7 @@ class ConsoleFormatter:
|
||||
def handle_agent_logs_started(
|
||||
self,
|
||||
agent_role: str,
|
||||
task_description: Optional[str] = None,
|
||||
task_description: str | None = None,
|
||||
verbose: bool = False,
|
||||
) -> None:
|
||||
"""Handle agent logs started event."""
|
||||
@@ -1367,10 +1367,11 @@ class ConsoleFormatter:
|
||||
if not verbose:
|
||||
return
|
||||
|
||||
from crewai.agents.parser import AgentAction, AgentFinish
|
||||
import json
|
||||
import re
|
||||
|
||||
from crewai.agents.parser import AgentAction, AgentFinish
|
||||
|
||||
agent_role = agent_role.partition("\n")[0]
|
||||
|
||||
if isinstance(formatted_answer, AgentAction):
|
||||
@@ -1437,8 +1438,17 @@ class ConsoleFormatter:
|
||||
|
||||
# Create tool output content with better formatting
|
||||
output_text = str(formatted_answer.result)
|
||||
if len(output_text) > 2000:
|
||||
output_text = output_text[:1997] + "..."
|
||||
if len(output_text) > 5000:
|
||||
if output_text.count("\n") > 10: # Multi-line structured data
|
||||
lines = output_text.split("\n")
|
||||
truncated_lines = lines[:10]
|
||||
remaining_lines = len(lines) - 10
|
||||
output_text = (
|
||||
"\n".join(truncated_lines)
|
||||
+ f"\n... and {remaining_lines} more rows"
|
||||
)
|
||||
else:
|
||||
output_text = output_text[:4997] + "..."
|
||||
|
||||
output_panel = Panel(
|
||||
Text(output_text, style="bright_green"),
|
||||
@@ -1473,9 +1483,9 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_memory_retrieval_started(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
crew_tree: Optional[Tree],
|
||||
) -> Optional[Tree]:
|
||||
agent_branch: Tree | None,
|
||||
crew_tree: Tree | None,
|
||||
) -> Tree | None:
|
||||
if not self.verbose:
|
||||
return None
|
||||
|
||||
@@ -1497,13 +1507,13 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_memory_retrieval_completed(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
crew_tree: Optional[Tree],
|
||||
agent_branch: Tree | None,
|
||||
crew_tree: Tree | None,
|
||||
memory_content: str,
|
||||
retrieval_time_ms: float,
|
||||
) -> None:
|
||||
if not self.verbose:
|
||||
return None
|
||||
return
|
||||
|
||||
branch_to_use = self.current_lite_agent_branch or agent_branch
|
||||
tree_to_use = branch_to_use or crew_tree
|
||||
@@ -1528,7 +1538,7 @@ class ConsoleFormatter:
|
||||
|
||||
if branch_to_use is None or tree_to_use is None:
|
||||
add_panel()
|
||||
return None
|
||||
return
|
||||
|
||||
memory_branch_found = False
|
||||
for child in branch_to_use.children:
|
||||
@@ -1565,13 +1575,13 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_memory_query_completed(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
agent_branch: Tree | None,
|
||||
source_type: str,
|
||||
query_time_ms: float,
|
||||
crew_tree: Optional[Tree],
|
||||
crew_tree: Tree | None,
|
||||
) -> None:
|
||||
if not self.verbose:
|
||||
return None
|
||||
return
|
||||
|
||||
branch_to_use = self.current_lite_agent_branch or agent_branch
|
||||
tree_to_use = branch_to_use or crew_tree
|
||||
@@ -1580,15 +1590,15 @@ class ConsoleFormatter:
|
||||
branch_to_use = tree_to_use
|
||||
|
||||
if branch_to_use is None:
|
||||
return None
|
||||
return
|
||||
|
||||
memory_type = source_type.replace("_", " ").title()
|
||||
|
||||
for child in branch_to_use.children:
|
||||
if "Memory Retrieval" in str(child.label):
|
||||
for child in child.children:
|
||||
sources_branch = child
|
||||
if "Sources Used" in str(child.label):
|
||||
for inner_child in child.children:
|
||||
sources_branch = inner_child
|
||||
if "Sources Used" in str(inner_child.label):
|
||||
sources_branch.add(f"✅ {memory_type} ({query_time_ms:.2f}ms)")
|
||||
break
|
||||
else:
|
||||
@@ -1598,13 +1608,13 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_memory_query_failed(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
crew_tree: Optional[Tree],
|
||||
agent_branch: Tree | None,
|
||||
crew_tree: Tree | None,
|
||||
error: str,
|
||||
source_type: str,
|
||||
) -> None:
|
||||
if not self.verbose:
|
||||
return None
|
||||
return
|
||||
|
||||
branch_to_use = self.current_lite_agent_branch or agent_branch
|
||||
tree_to_use = branch_to_use or crew_tree
|
||||
@@ -1613,15 +1623,15 @@ class ConsoleFormatter:
|
||||
branch_to_use = tree_to_use
|
||||
|
||||
if branch_to_use is None:
|
||||
return None
|
||||
return
|
||||
|
||||
memory_type = source_type.replace("_", " ").title()
|
||||
|
||||
for child in branch_to_use.children:
|
||||
if "Memory Retrieval" in str(child.label):
|
||||
for child in child.children:
|
||||
sources_branch = child
|
||||
if "Sources Used" in str(child.label):
|
||||
for inner_child in child.children:
|
||||
sources_branch = inner_child
|
||||
if "Sources Used" in str(inner_child.label):
|
||||
sources_branch.add(f"❌ {memory_type} - Error: {error}")
|
||||
break
|
||||
else:
|
||||
@@ -1630,16 +1640,16 @@ class ConsoleFormatter:
|
||||
break
|
||||
|
||||
def handle_memory_save_started(
|
||||
self, agent_branch: Optional[Tree], crew_tree: Optional[Tree]
|
||||
self, agent_branch: Tree | None, crew_tree: Tree | None
|
||||
) -> None:
|
||||
if not self.verbose:
|
||||
return None
|
||||
return
|
||||
|
||||
branch_to_use = agent_branch or self.current_lite_agent_branch
|
||||
tree_to_use = branch_to_use or crew_tree
|
||||
|
||||
if tree_to_use is None:
|
||||
return None
|
||||
return
|
||||
|
||||
for child in tree_to_use.children:
|
||||
if "Memory Update" in str(child.label):
|
||||
@@ -1655,19 +1665,19 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_memory_save_completed(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
crew_tree: Optional[Tree],
|
||||
agent_branch: Tree | None,
|
||||
crew_tree: Tree | None,
|
||||
save_time_ms: float,
|
||||
source_type: str,
|
||||
) -> None:
|
||||
if not self.verbose:
|
||||
return None
|
||||
return
|
||||
|
||||
branch_to_use = agent_branch or self.current_lite_agent_branch
|
||||
tree_to_use = branch_to_use or crew_tree
|
||||
|
||||
if tree_to_use is None:
|
||||
return None
|
||||
return
|
||||
|
||||
memory_type = source_type.replace("_", " ").title()
|
||||
content = f"✅ {memory_type} Memory Saved ({save_time_ms:.2f}ms)"
|
||||
@@ -1685,19 +1695,19 @@ class ConsoleFormatter:
|
||||
|
||||
def handle_memory_save_failed(
|
||||
self,
|
||||
agent_branch: Optional[Tree],
|
||||
agent_branch: Tree | None,
|
||||
error: str,
|
||||
source_type: str,
|
||||
crew_tree: Optional[Tree],
|
||||
crew_tree: Tree | None,
|
||||
) -> None:
|
||||
if not self.verbose:
|
||||
return None
|
||||
return
|
||||
|
||||
branch_to_use = agent_branch or self.current_lite_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
|
||||
return
|
||||
|
||||
memory_type = source_type.replace("_", " ").title()
|
||||
content = f"❌ {memory_type} Memory Save Failed"
|
||||
@@ -1738,7 +1748,7 @@ class ConsoleFormatter:
|
||||
def handle_guardrail_completed(
|
||||
self,
|
||||
success: bool,
|
||||
error: Optional[str],
|
||||
error: str | None,
|
||||
retry_count: int,
|
||||
) -> None:
|
||||
"""Display guardrail evaluation result.
|
||||
|
||||
@@ -5,12 +5,20 @@ import time
|
||||
from difflib import SequenceMatcher
|
||||
from json import JSONDecodeError
|
||||
from textwrap import dedent
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
|
||||
from typing import TYPE_CHECKING, Any, Union
|
||||
|
||||
import json5
|
||||
from json_repair import repair_json
|
||||
from json_repair import repair_json # type: ignore
|
||||
|
||||
from crewai.agents.tools_handler import ToolsHandler
|
||||
from crewai.events.event_bus import crewai_event_bus
|
||||
from crewai.events.types.tool_usage_events import (
|
||||
ToolSelectionErrorEvent,
|
||||
ToolUsageErrorEvent,
|
||||
ToolUsageFinishedEvent,
|
||||
ToolUsageStartedEvent,
|
||||
ToolValidateInputErrorEvent,
|
||||
)
|
||||
from crewai.task import Task
|
||||
from crewai.telemetry import Telemetry
|
||||
from crewai.tools.structured_tool import CrewStructuredTool
|
||||
@@ -20,14 +28,6 @@ from crewai.utilities.agent_utils import (
|
||||
get_tool_names,
|
||||
render_text_description_and_args,
|
||||
)
|
||||
from crewai.events.event_bus import crewai_event_bus
|
||||
from crewai.events.types.tool_usage_events import (
|
||||
ToolSelectionErrorEvent,
|
||||
ToolUsageErrorEvent,
|
||||
ToolUsageFinishedEvent,
|
||||
ToolUsageStartedEvent,
|
||||
ToolValidateInputErrorEvent,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from crewai.agents.agent_builder.base_agent import BaseAgent
|
||||
@@ -44,7 +44,7 @@ OPENAI_BIGGER_MODELS = [
|
||||
]
|
||||
|
||||
|
||||
class ToolUsageErrorException(Exception):
|
||||
class ToolUsageError(Exception):
|
||||
"""Exception raised for errors in the tool usage."""
|
||||
|
||||
def __init__(self, message: str) -> None:
|
||||
@@ -68,13 +68,13 @@ class ToolUsage:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tools_handler: Optional[ToolsHandler],
|
||||
tools: List[CrewStructuredTool],
|
||||
task: Optional[Task],
|
||||
tools_handler: ToolsHandler | None,
|
||||
tools: list[CrewStructuredTool],
|
||||
task: Task | None,
|
||||
function_calling_llm: Any,
|
||||
agent: Optional[Union["BaseAgent", "LiteAgent"]] = None,
|
||||
agent: Union["BaseAgent", "LiteAgent"] | None = None,
|
||||
action: Any = None,
|
||||
fingerprint_context: Optional[Dict[str, str]] = None,
|
||||
fingerprint_context: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
self._i18n: I18N = agent.i18n if agent else I18N()
|
||||
self._printer: Printer = Printer()
|
||||
@@ -105,9 +105,9 @@ class ToolUsage:
|
||||
return self._tool_calling(tool_string)
|
||||
|
||||
def use(
|
||||
self, calling: Union[ToolCalling, InstructorToolCalling], tool_string: str
|
||||
self, calling: ToolCalling | InstructorToolCalling, tool_string: str
|
||||
) -> str:
|
||||
if isinstance(calling, ToolUsageErrorException):
|
||||
if isinstance(calling, ToolUsageError):
|
||||
error = calling.message
|
||||
if self.agent and self.agent.verbose:
|
||||
self._printer.print(content=f"\n\n{error}\n", color="red")
|
||||
@@ -130,8 +130,7 @@ class ToolUsage:
|
||||
and tool.name == self._i18n.tools("add_image")["name"] # type: ignore
|
||||
):
|
||||
try:
|
||||
result = self._use(tool_string=tool_string, tool=tool, calling=calling)
|
||||
return result
|
||||
return self._use(tool_string=tool_string, tool=tool, calling=calling)
|
||||
|
||||
except Exception as e:
|
||||
error = getattr(e, "message", str(e))
|
||||
@@ -147,7 +146,7 @@ class ToolUsage:
|
||||
self,
|
||||
tool_string: str,
|
||||
tool: CrewStructuredTool,
|
||||
calling: Union[ToolCalling, InstructorToolCalling],
|
||||
calling: ToolCalling | InstructorToolCalling,
|
||||
) -> str:
|
||||
if self._check_tool_repeated_usage(calling=calling): # type: ignore # _check_tool_repeated_usage of "ToolUsage" does not return a value (it only ever returns None)
|
||||
try:
|
||||
@@ -159,8 +158,7 @@ class ToolUsage:
|
||||
tool_name=tool.name,
|
||||
attempts=self._run_attempts,
|
||||
)
|
||||
result = self._format_result(result=result) # type: ignore # "_format_result" of "ToolUsage" does not return a value (it only ever returns None)
|
||||
return result # type: ignore # Fix the return type of this function
|
||||
return self._format_result(result=result) # type: ignore # "_format_result" of "ToolUsage" does not return a value (it only ever returns None)
|
||||
|
||||
except Exception:
|
||||
if self.task:
|
||||
@@ -176,7 +174,7 @@ class ToolUsage:
|
||||
"agent": self.agent,
|
||||
}
|
||||
|
||||
if self.agent.fingerprint:
|
||||
if hasattr(self.agent, 'fingerprint') and self.agent.fingerprint:
|
||||
event_data.update(self.agent.fingerprint)
|
||||
if self.task:
|
||||
event_data["task_name"] = self.task.name or self.task.description
|
||||
@@ -185,13 +183,14 @@ class ToolUsage:
|
||||
|
||||
started_at = time.time()
|
||||
from_cache = False
|
||||
result = None # type: ignore
|
||||
|
||||
if self.tools_handler and self.tools_handler.cache:
|
||||
result = self.tools_handler.cache.read(
|
||||
tool=calling.tool_name, input=calling.arguments
|
||||
) # type: ignore
|
||||
from_cache = result is not None
|
||||
cache_result = self.tools_handler.cache.read(
|
||||
tool=calling.tool_name, input=str(calling.arguments)
|
||||
)
|
||||
from_cache = cache_result is not None
|
||||
if cache_result is not None:
|
||||
result = cache_result
|
||||
|
||||
available_tool = next(
|
||||
(
|
||||
@@ -207,8 +206,7 @@ class ToolUsage:
|
||||
try:
|
||||
result = usage_limit_error
|
||||
self._telemetry.tool_usage_error(llm=self.function_calling_llm)
|
||||
result = self._format_result(result=result)
|
||||
return result
|
||||
return self._format_result(result=result)
|
||||
except Exception:
|
||||
if self.task:
|
||||
self.task.increment_tools_errors()
|
||||
@@ -255,7 +253,7 @@ class ToolUsage:
|
||||
error_message = self._i18n.errors("tool_usage_exception").format(
|
||||
error=e, tool=tool.name, tool_inputs=tool.description
|
||||
)
|
||||
error = ToolUsageErrorException(
|
||||
error = ToolUsageError(
|
||||
f"\n{error_message}.\nMoving on then. {self._i18n.slice('format').format(tool_names=self.tools_names)}"
|
||||
).message
|
||||
if self.task:
|
||||
@@ -331,6 +329,15 @@ class ToolUsage:
|
||||
self.task.used_tools += 1
|
||||
if self._should_remember_format():
|
||||
result = self._remember_format(result=result)
|
||||
|
||||
if isinstance(result, (list, dict)):
|
||||
import json
|
||||
|
||||
try:
|
||||
return json.dumps(result, indent=2, ensure_ascii=False)
|
||||
except (TypeError, ValueError):
|
||||
return str(result)
|
||||
|
||||
return str(result)
|
||||
|
||||
def _should_remember_format(self) -> bool:
|
||||
@@ -346,7 +353,7 @@ class ToolUsage:
|
||||
return result
|
||||
|
||||
def _check_tool_repeated_usage(
|
||||
self, calling: Union[ToolCalling, InstructorToolCalling]
|
||||
self, calling: ToolCalling | InstructorToolCalling
|
||||
) -> bool:
|
||||
if not self.tools_handler:
|
||||
return False
|
||||
@@ -393,7 +400,7 @@ class ToolUsage:
|
||||
return tool
|
||||
if self.task:
|
||||
self.task.increment_tools_errors()
|
||||
tool_selection_data: Dict[str, Any] = {
|
||||
tool_selection_data: dict[str, Any] = {
|
||||
"agent_key": getattr(self.agent, "key", None) if self.agent else None,
|
||||
"agent_role": getattr(self.agent, "role", None) if self.agent else None,
|
||||
"tool_name": tool_name,
|
||||
@@ -410,27 +417,24 @@ class ToolUsage:
|
||||
),
|
||||
)
|
||||
raise Exception(error)
|
||||
else:
|
||||
error = f"I forgot the Action name, these are the only available Actions: {self.tools_description}"
|
||||
crewai_event_bus.emit(
|
||||
self,
|
||||
ToolSelectionErrorEvent(
|
||||
**tool_selection_data,
|
||||
error=error,
|
||||
),
|
||||
)
|
||||
raise Exception(error)
|
||||
error = f"I forgot the Action name, these are the only available Actions: {self.tools_description}"
|
||||
crewai_event_bus.emit(
|
||||
self,
|
||||
ToolSelectionErrorEvent(
|
||||
**tool_selection_data,
|
||||
error=error,
|
||||
),
|
||||
)
|
||||
raise Exception(error)
|
||||
|
||||
def _render(self) -> str:
|
||||
"""Render the tool name and description in plain text."""
|
||||
descriptions = []
|
||||
for tool in self.tools:
|
||||
descriptions.append(tool.description)
|
||||
descriptions = [tool.description for tool in self.tools]
|
||||
return "\n--\n".join(descriptions)
|
||||
|
||||
def _function_calling(
|
||||
self, tool_string: str
|
||||
) -> Union[ToolCalling, InstructorToolCalling]:
|
||||
) -> ToolCalling | InstructorToolCalling:
|
||||
model = (
|
||||
InstructorToolCalling
|
||||
if self.function_calling_llm.supports_function_calling()
|
||||
@@ -453,13 +457,13 @@ class ToolUsage:
|
||||
)
|
||||
tool_object = converter.to_pydantic()
|
||||
if not isinstance(tool_object, (ToolCalling, InstructorToolCalling)):
|
||||
raise ToolUsageErrorException("Failed to parse tool calling")
|
||||
raise ToolUsageError("Failed to parse tool calling")
|
||||
|
||||
return tool_object
|
||||
|
||||
def _original_tool_calling(
|
||||
self, tool_string: str, raise_error: bool = False
|
||||
) -> Union[ToolCalling, InstructorToolCalling, ToolUsageErrorException]:
|
||||
) -> ToolCalling | InstructorToolCalling | ToolUsageError:
|
||||
tool_name = self.action.tool
|
||||
tool = self._select_tool(tool_name)
|
||||
try:
|
||||
@@ -468,18 +472,16 @@ class ToolUsage:
|
||||
except Exception:
|
||||
if raise_error:
|
||||
raise
|
||||
else:
|
||||
return ToolUsageErrorException(
|
||||
f"{self._i18n.errors('tool_arguments_error')}"
|
||||
)
|
||||
return ToolUsageError(
|
||||
f"{self._i18n.errors('tool_arguments_error')}"
|
||||
)
|
||||
|
||||
if not isinstance(arguments, dict):
|
||||
if raise_error:
|
||||
raise
|
||||
else:
|
||||
return ToolUsageErrorException(
|
||||
f"{self._i18n.errors('tool_arguments_error')}"
|
||||
)
|
||||
return ToolUsageError(
|
||||
f"{self._i18n.errors('tool_arguments_error')}"
|
||||
)
|
||||
|
||||
return ToolCalling(
|
||||
tool_name=tool.name,
|
||||
@@ -488,15 +490,14 @@ class ToolUsage:
|
||||
|
||||
def _tool_calling(
|
||||
self, tool_string: str
|
||||
) -> Union[ToolCalling, InstructorToolCalling, ToolUsageErrorException]:
|
||||
) -> ToolCalling | InstructorToolCalling | ToolUsageError:
|
||||
try:
|
||||
try:
|
||||
return self._original_tool_calling(tool_string, raise_error=True)
|
||||
except Exception:
|
||||
if self.function_calling_llm:
|
||||
return self._function_calling(tool_string)
|
||||
else:
|
||||
return self._original_tool_calling(tool_string)
|
||||
return self._original_tool_calling(tool_string)
|
||||
except Exception as e:
|
||||
self._run_attempts += 1
|
||||
if self._run_attempts > self._max_parsing_attempts:
|
||||
@@ -505,12 +506,12 @@ class ToolUsage:
|
||||
self.task.increment_tools_errors()
|
||||
if self.agent and self.agent.verbose:
|
||||
self._printer.print(content=f"\n\n{e}\n", color="red")
|
||||
return ToolUsageErrorException( # type: ignore # Incompatible return value type (got "ToolUsageErrorException", expected "ToolCalling | InstructorToolCalling")
|
||||
return ToolUsageError( # type: ignore # Incompatible return value type (got "ToolUsageError", expected "ToolCalling | InstructorToolCalling")
|
||||
f"{self._i18n.errors('tool_usage_error').format(error=e)}\nMoving on then. {self._i18n.slice('format').format(tool_names=self.tools_names)}"
|
||||
)
|
||||
return self._tool_calling(tool_string)
|
||||
|
||||
def _validate_tool_input(self, tool_input: Optional[str]) -> Dict[str, Any]:
|
||||
def _validate_tool_input(self, tool_input: str | None) -> dict[str, Any]:
|
||||
if tool_input is None:
|
||||
return {}
|
||||
|
||||
@@ -534,7 +535,7 @@ class ToolUsage:
|
||||
return arguments
|
||||
except (ValueError, SyntaxError):
|
||||
repaired_input = repair_json(tool_input)
|
||||
pass # Continue to the next parsing attempt
|
||||
# Continue to the next parsing attempt
|
||||
|
||||
# Attempt 3: Parse as JSON5
|
||||
try:
|
||||
@@ -586,7 +587,7 @@ class ToolUsage:
|
||||
def on_tool_error(
|
||||
self,
|
||||
tool: Any,
|
||||
tool_calling: Union[ToolCalling, InstructorToolCalling],
|
||||
tool_calling: ToolCalling | InstructorToolCalling,
|
||||
e: Exception,
|
||||
) -> None:
|
||||
event_data = self._prepare_event_data(tool, tool_calling)
|
||||
@@ -595,7 +596,7 @@ class ToolUsage:
|
||||
def on_tool_use_finished(
|
||||
self,
|
||||
tool: Any,
|
||||
tool_calling: Union[ToolCalling, InstructorToolCalling],
|
||||
tool_calling: ToolCalling | InstructorToolCalling,
|
||||
from_cache: bool,
|
||||
started_at: float,
|
||||
result: Any,
|
||||
@@ -616,7 +617,7 @@ class ToolUsage:
|
||||
crewai_event_bus.emit(self, ToolUsageFinishedEvent(**event_data))
|
||||
|
||||
def _prepare_event_data(
|
||||
self, tool: Any, tool_calling: Union[ToolCalling, InstructorToolCalling]
|
||||
self, tool: Any, tool_calling: ToolCalling | InstructorToolCalling
|
||||
) -> dict:
|
||||
event_data = {
|
||||
"run_attempts": self._run_attempts,
|
||||
|
||||
@@ -1,24 +1,24 @@
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from crewai.agents.parser import AgentAction
|
||||
from crewai.security import Fingerprint
|
||||
from crewai.tools.structured_tool import CrewStructuredTool
|
||||
from crewai.tools.tool_types import ToolResult
|
||||
from crewai.tools.tool_usage import ToolUsage, ToolUsageErrorException
|
||||
from crewai.tools.tool_usage import ToolUsage, ToolUsageError
|
||||
from crewai.utilities.i18n import I18N
|
||||
|
||||
|
||||
def execute_tool_and_check_finality(
|
||||
agent_action: AgentAction,
|
||||
tools: List[CrewStructuredTool],
|
||||
tools: list[CrewStructuredTool],
|
||||
i18n: I18N,
|
||||
agent_key: Optional[str] = None,
|
||||
agent_role: Optional[str] = None,
|
||||
tools_handler: Optional[Any] = None,
|
||||
task: Optional[Any] = None,
|
||||
agent: Optional[Any] = None,
|
||||
function_calling_llm: Optional[Any] = None,
|
||||
fingerprint_context: Optional[Dict[str, str]] = None,
|
||||
agent_key: str | None = None,
|
||||
agent_role: str | None = None,
|
||||
tools_handler: Any | None = None,
|
||||
task: Any | None = None,
|
||||
agent: Any | None = None,
|
||||
function_calling_llm: Any | None = None,
|
||||
fingerprint_context: dict[str, str] | None = None,
|
||||
) -> ToolResult:
|
||||
"""Execute a tool and check if the result should be treated as a final answer.
|
||||
|
||||
@@ -50,7 +50,7 @@ def execute_tool_and_check_finality(
|
||||
fingerprint_obj = Fingerprint.from_dict(fingerprint_context)
|
||||
agent.set_fingerprint(fingerprint_obj)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to set fingerprint: {e}")
|
||||
raise ValueError(f"Failed to set fingerprint: {e}") from e
|
||||
|
||||
# Create tool usage instance
|
||||
tool_usage = ToolUsage(
|
||||
@@ -65,7 +65,7 @@ def execute_tool_and_check_finality(
|
||||
# Parse tool calling
|
||||
tool_calling = tool_usage.parse_tool_calling(agent_action.text)
|
||||
|
||||
if isinstance(tool_calling, ToolUsageErrorException):
|
||||
if isinstance(tool_calling, ToolUsageError):
|
||||
return ToolResult(tool_calling.message, False)
|
||||
|
||||
# Check if tool name matches
|
||||
|
||||
235
tests/test_mcp_tool_output.py
Normal file
235
tests/test_mcp_tool_output.py
Normal file
@@ -0,0 +1,235 @@
|
||||
import json
|
||||
from typing import Any, ClassVar
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
from crewai.agent import Agent
|
||||
from crewai.agents.agent_builder.base_agent import BaseAgent
|
||||
from crewai.crew import Crew
|
||||
from crewai.project import CrewBase, agent, crew, task
|
||||
from crewai.task import Task
|
||||
from crewai.tools import tool
|
||||
|
||||
|
||||
@tool
|
||||
def mock_bigquery_single_row():
|
||||
"""Mock BigQuery tool that returns a single row"""
|
||||
return {"id": 1, "name": "John", "age": 30}
|
||||
|
||||
|
||||
@tool
|
||||
def mock_bigquery_multiple_rows():
|
||||
"""Mock BigQuery tool that returns multiple rows"""
|
||||
return [
|
||||
{"id": 1, "name": "John", "age": 30},
|
||||
{"id": 2, "name": "Jane", "age": 25},
|
||||
{"id": 3, "name": "Bob", "age": 35},
|
||||
{"id": 4, "name": "Alice", "age": 28},
|
||||
]
|
||||
|
||||
|
||||
@tool
|
||||
def mock_bigquery_large_dataset():
|
||||
"""Mock BigQuery tool that returns a large dataset"""
|
||||
return [{"id": i, "name": f"User{i}", "value": f"data_{i}"} for i in range(100)]
|
||||
|
||||
|
||||
@tool
|
||||
def mock_bigquery_nested_data():
|
||||
"""Mock BigQuery tool that returns nested data structures"""
|
||||
return [
|
||||
{
|
||||
"id": 1,
|
||||
"user": {"name": "John", "email": "john@example.com"},
|
||||
"orders": [
|
||||
{"order_id": 101, "amount": 50.0},
|
||||
{"order_id": 102, "amount": 75.0},
|
||||
],
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"user": {"name": "Jane", "email": "jane@example.com"},
|
||||
"orders": [{"order_id": 103, "amount": 100.0}],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@CrewBase
|
||||
class MCPTestCrew:
|
||||
agents_config = "config/agents.yaml"
|
||||
tasks_config = "config/tasks.yaml"
|
||||
mcp_server_params: ClassVar[dict[str, Any]] = {"host": "localhost", "port": 8000}
|
||||
mcp_connect_timeout = 120
|
||||
|
||||
agents: list[BaseAgent]
|
||||
tasks: list[Task]
|
||||
|
||||
@agent
|
||||
def data_analyst(self):
|
||||
return Agent(
|
||||
role="Data Analyst",
|
||||
goal="Analyze data from various sources",
|
||||
backstory="Expert in data analysis and BigQuery",
|
||||
tools=[mock_bigquery_single_row, mock_bigquery_multiple_rows],
|
||||
)
|
||||
|
||||
@agent
|
||||
def mcp_agent(self):
|
||||
return Agent(
|
||||
role="MCP Agent",
|
||||
goal="Use MCP tools to fetch data",
|
||||
backstory="Agent that uses MCP tools",
|
||||
tools=self.get_mcp_tools(),
|
||||
)
|
||||
|
||||
@task
|
||||
def analyze_single_row(self):
|
||||
return Task(
|
||||
description="Use mock_bigquery_single_row tool to get data",
|
||||
expected_output="Single row of data",
|
||||
agent=self.data_analyst(),
|
||||
)
|
||||
|
||||
@task
|
||||
def analyze_multiple_rows(self):
|
||||
return Task(
|
||||
description="Use mock_bigquery_multiple_rows tool to get data",
|
||||
expected_output="Multiple rows of data",
|
||||
agent=self.data_analyst(),
|
||||
)
|
||||
|
||||
@crew
|
||||
def crew(self):
|
||||
return Crew(agents=self.agents, tasks=self.tasks, verbose=True)
|
||||
|
||||
|
||||
def test_single_row_tool_output():
|
||||
"""Test that single row tool output works correctly"""
|
||||
result = mock_bigquery_single_row.invoke({})
|
||||
assert isinstance(result, dict)
|
||||
assert result["id"] == 1
|
||||
assert result["name"] == "John"
|
||||
assert result["age"] == 30
|
||||
|
||||
|
||||
def test_multiple_rows_tool_output():
|
||||
"""Test that multiple rows tool output is preserved"""
|
||||
result = mock_bigquery_multiple_rows.invoke({})
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 4
|
||||
assert result[0]["id"] == 1
|
||||
assert result[1]["id"] == 2
|
||||
assert result[2]["id"] == 3
|
||||
assert result[3]["id"] == 4
|
||||
|
||||
|
||||
def test_large_dataset_tool_output():
|
||||
"""Test that large datasets are handled correctly"""
|
||||
result = mock_bigquery_large_dataset.invoke({})
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 100
|
||||
assert result[0]["id"] == 0
|
||||
assert result[99]["id"] == 99
|
||||
|
||||
|
||||
def test_nested_data_tool_output():
|
||||
"""Test that nested data structures are preserved"""
|
||||
result = mock_bigquery_nested_data.invoke({})
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 2
|
||||
assert result[0]["user"]["name"] == "John"
|
||||
assert len(result[0]["orders"]) == 2
|
||||
assert result[1]["user"]["name"] == "Jane"
|
||||
assert len(result[1]["orders"]) == 1
|
||||
|
||||
|
||||
def test_tool_result_formatting():
|
||||
"""Test that tool results are properly formatted as strings"""
|
||||
from crewai.tools.tool_usage import ToolUsage
|
||||
|
||||
tool_usage = ToolUsage()
|
||||
|
||||
single_result = mock_bigquery_single_row.invoke({})
|
||||
formatted_single = tool_usage._format_result(single_result)
|
||||
assert isinstance(formatted_single, str)
|
||||
parsed_single = json.loads(formatted_single)
|
||||
assert parsed_single["id"] == 1
|
||||
|
||||
multi_result = mock_bigquery_multiple_rows.invoke({})
|
||||
formatted_multi = tool_usage._format_result(multi_result)
|
||||
assert isinstance(formatted_multi, str)
|
||||
parsed_multi = json.loads(formatted_multi)
|
||||
assert len(parsed_multi) == 4
|
||||
assert parsed_multi[0]["id"] == 1
|
||||
assert parsed_multi[3]["id"] == 4
|
||||
|
||||
|
||||
def test_mcp_crew_with_mock_tools():
|
||||
"""Test MCP crew integration with mock tools"""
|
||||
with patch("embedchain.client.Client.setup"):
|
||||
from crewai_tools import MCPServerAdapter
|
||||
from crewai_tools.adapters.mcp_adapter import ToolCollection
|
||||
|
||||
mock_adapter = Mock(spec=MCPServerAdapter)
|
||||
mock_adapter.tools = ToolCollection([mock_bigquery_multiple_rows])
|
||||
|
||||
with patch("crewai_tools.MCPServerAdapter", return_value=mock_adapter):
|
||||
crew = MCPTestCrew()
|
||||
mcp_agent = crew.mcp_agent()
|
||||
assert mock_bigquery_multiple_rows in mcp_agent.tools
|
||||
|
||||
|
||||
def test_tool_output_preserves_structure():
|
||||
"""Test that tool output preserves data structure through the processing pipeline"""
|
||||
from crewai.tools.tool_usage import ToolUsage
|
||||
|
||||
tool_usage = ToolUsage()
|
||||
|
||||
bigquery_result = [
|
||||
{"id": 1, "name": "John", "revenue": 1000.50},
|
||||
{"id": 2, "name": "Jane", "revenue": 2500.75},
|
||||
{"id": 3, "name": "Bob", "revenue": 1750.25},
|
||||
]
|
||||
|
||||
formatted_result = tool_usage._format_result(bigquery_result)
|
||||
|
||||
assert isinstance(formatted_result, str)
|
||||
|
||||
parsed_result = json.loads(formatted_result)
|
||||
assert len(parsed_result) == 3
|
||||
assert parsed_result[0]["id"] == 1
|
||||
assert parsed_result[1]["name"] == "Jane"
|
||||
assert parsed_result[2]["revenue"] == 1750.25
|
||||
|
||||
|
||||
def test_tool_output_backward_compatibility():
|
||||
"""Test that simple string/number outputs still work"""
|
||||
from crewai.tools.tool_usage import ToolUsage
|
||||
|
||||
tool_usage = ToolUsage()
|
||||
|
||||
string_result = "Simple string result"
|
||||
formatted_string = tool_usage._format_result(string_result)
|
||||
assert formatted_string == "Simple string result"
|
||||
|
||||
number_result = 42
|
||||
formatted_number = tool_usage._format_result(number_result)
|
||||
assert formatted_number == "42"
|
||||
|
||||
bool_result = True
|
||||
formatted_bool = tool_usage._format_result(bool_result)
|
||||
assert formatted_bool == "True"
|
||||
|
||||
|
||||
def test_malformed_data_handling():
|
||||
"""Test that malformed data is handled gracefully"""
|
||||
from crewai.tools.tool_usage import ToolUsage
|
||||
|
||||
tool_usage = ToolUsage()
|
||||
|
||||
class NonSerializable:
|
||||
def __str__(self):
|
||||
return "NonSerializable object"
|
||||
|
||||
non_serializable = NonSerializable()
|
||||
formatted_result = tool_usage._format_result(non_serializable)
|
||||
assert formatted_result == "NonSerializable object"
|
||||
Reference in New Issue
Block a user