mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-04-28 05:42:46 +00:00
Compare commits
8 Commits
devin/1740
...
better-tel
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
17a19dee0c | ||
|
|
75a84a55c2 | ||
|
|
7460906712 | ||
|
|
84c809eee2 | ||
|
|
9bb46f158c | ||
|
|
2d4a7701e6 | ||
|
|
9c040c9e97 | ||
|
|
2d07c8d2e4 |
@@ -136,21 +136,17 @@ crewai test -n 5 -m gpt-3.5-turbo
|
||||
|
||||
### 8. Run
|
||||
|
||||
Run the crew or flow.
|
||||
Run the crew.
|
||||
|
||||
```shell Terminal
|
||||
crewai run
|
||||
```
|
||||
|
||||
<Note>
|
||||
Starting from version 0.103.0, the `crewai run` command can be used to run both standard crews and flows. For flows, it automatically detects the type from pyproject.toml and runs the appropriate command. This is now the recommended way to run both crews and flows.
|
||||
</Note>
|
||||
|
||||
<Note>
|
||||
Make sure to run these commands from the directory where your CrewAI project is set up.
|
||||
Some commands may require additional configuration or setup within your project structure.
|
||||
</Note>
|
||||
|
||||
|
||||
### 9. Chat
|
||||
|
||||
Starting in version `0.98.0`, when you run the `crewai chat` command, you start an interactive session with your crew. The AI assistant will guide you by asking for necessary inputs to execute the crew. Once all inputs are provided, the crew will execute its tasks.
|
||||
@@ -179,6 +175,7 @@ def crew(self) -> Crew:
|
||||
```
|
||||
</Note>
|
||||
|
||||
|
||||
### 10. API Keys
|
||||
|
||||
When running ```crewai create crew``` command, the CLI will first show you the top 5 most common LLM providers and ask you to select one.
|
||||
|
||||
@@ -150,12 +150,12 @@ final_output = flow.kickoff()
|
||||
|
||||
print("---- Final Output ----")
|
||||
print(final_output)
|
||||
```
|
||||
````
|
||||
|
||||
```text Output
|
||||
---- Final Output ----
|
||||
Second method received: Output from first_method
|
||||
```
|
||||
````
|
||||
|
||||
</CodeGroup>
|
||||
|
||||
@@ -738,34 +738,3 @@ Also, check out our YouTube video on how to use flows in CrewAI below!
|
||||
referrerpolicy="strict-origin-when-cross-origin"
|
||||
allowfullscreen
|
||||
></iframe>
|
||||
|
||||
## Running Flows
|
||||
|
||||
There are two ways to run a flow:
|
||||
|
||||
### Using the Flow API
|
||||
|
||||
You can run a flow programmatically by creating an instance of your flow class and calling the `kickoff()` method:
|
||||
|
||||
```python
|
||||
flow = ExampleFlow()
|
||||
result = flow.kickoff()
|
||||
```
|
||||
|
||||
### Using the CLI
|
||||
|
||||
Starting from version 0.103.0, you can run flows using the `crewai run` command:
|
||||
|
||||
```shell
|
||||
crewai run
|
||||
```
|
||||
|
||||
This command automatically detects if your project is a flow (based on the `type = "flow"` setting in your pyproject.toml) and runs it accordingly. This is the recommended way to run flows from the command line.
|
||||
|
||||
For backward compatibility, you can also use:
|
||||
|
||||
```shell
|
||||
crewai flow kickoff
|
||||
```
|
||||
|
||||
However, the `crewai run` command is now the preferred method as it works for both crews and flows.
|
||||
|
||||
@@ -506,7 +506,7 @@ my_crew = Crew(
|
||||
)
|
||||
```
|
||||
|
||||
### Resetting Memory via cli
|
||||
### Resetting Memory
|
||||
|
||||
```shell
|
||||
crewai reset-memories [OPTIONS]
|
||||
@@ -520,46 +520,8 @@ crewai reset-memories [OPTIONS]
|
||||
| `-s`, `--short` | Reset SHORT TERM memory. | Flag (boolean) | False |
|
||||
| `-e`, `--entities` | Reset ENTITIES memory. | Flag (boolean) | False |
|
||||
| `-k`, `--kickoff-outputs` | Reset LATEST KICKOFF TASK OUTPUTS. | Flag (boolean) | False |
|
||||
| `-kn`, `--knowledge` | Reset KNOWLEDEGE storage | Flag (boolean) | False |
|
||||
| `-a`, `--all` | Reset ALL memories. | Flag (boolean) | False |
|
||||
|
||||
Note: To use the cli command you need to have your crew in a file called crew.py in the same directory.
|
||||
|
||||
|
||||
|
||||
|
||||
### Resetting Memory via crew object
|
||||
|
||||
```python
|
||||
|
||||
my_crew = Crew(
|
||||
agents=[...],
|
||||
tasks=[...],
|
||||
process=Process.sequential,
|
||||
memory=True,
|
||||
verbose=True,
|
||||
embedder={
|
||||
"provider": "custom",
|
||||
"config": {
|
||||
"embedder": CustomEmbedder()
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
my_crew.reset_memories(command_type = 'all') # Resets all the memory
|
||||
```
|
||||
|
||||
#### Resetting Memory Options
|
||||
|
||||
| Command Type | Description |
|
||||
| :----------------- | :------------------------------- |
|
||||
| `long` | Reset LONG TERM memory. |
|
||||
| `short` | Reset SHORT TERM memory. |
|
||||
| `entities` | Reset ENTITIES memory. |
|
||||
| `kickoff_outputs` | Reset LATEST KICKOFF TASK OUTPUTS. |
|
||||
| `knowledge` | Reset KNOWLEDGE memory. |
|
||||
| `all` | Reset ALL memories. |
|
||||
|
||||
|
||||
## Benefits of Using CrewAI's Memory System
|
||||
|
||||
|
||||
@@ -54,8 +54,7 @@ coding_agent = Agent(
|
||||
# Create a task that requires code execution
|
||||
data_analysis_task = Task(
|
||||
description="Analyze the given dataset and calculate the average age of participants. Ages: {ages}",
|
||||
agent=coding_agent,
|
||||
expected_output="The average age of the participants."
|
||||
agent=coding_agent
|
||||
)
|
||||
|
||||
# Create a crew and add the task
|
||||
@@ -117,4 +116,4 @@ async def async_multiple_crews():
|
||||
|
||||
# Run the async function
|
||||
asyncio.run(async_multiple_crews())
|
||||
```
|
||||
```
|
||||
@@ -114,6 +114,7 @@ class Agent(BaseAgent):
|
||||
|
||||
@model_validator(mode="after")
|
||||
def post_init_setup(self):
|
||||
self._set_knowledge()
|
||||
self.agent_ops_agent_name = self.role
|
||||
|
||||
self.llm = create_llm(self.llm)
|
||||
@@ -133,11 +134,8 @@ class Agent(BaseAgent):
|
||||
self.cache_handler = CacheHandler()
|
||||
self.set_cache_handler(self.cache_handler)
|
||||
|
||||
def set_knowledge(self, crew_embedder: Optional[Dict[str, Any]] = None):
|
||||
def _set_knowledge(self):
|
||||
try:
|
||||
if self.embedder is None and crew_embedder:
|
||||
self.embedder = crew_embedder
|
||||
|
||||
if self.knowledge_sources:
|
||||
full_pattern = re.compile(r"[^a-zA-Z0-9\-_\r\n]|(\.\.)")
|
||||
knowledge_agent_name = f"{re.sub(full_pattern, '_', self.role)}"
|
||||
|
||||
@@ -351,6 +351,3 @@ class BaseAgent(ABC, BaseModel):
|
||||
if not self._rpm_controller:
|
||||
self._rpm_controller = rpm_controller
|
||||
self.create_agent_executor()
|
||||
|
||||
def set_knowledge(self, crew_embedder: Optional[Dict[str, Any]] = None):
|
||||
pass
|
||||
|
||||
@@ -124,15 +124,14 @@ class CrewAgentParser:
|
||||
)
|
||||
|
||||
def _extract_thought(self, text: str) -> str:
|
||||
thought_index = text.find("\n\nAction")
|
||||
if thought_index == -1:
|
||||
thought_index = text.find("\n\nFinal Answer")
|
||||
if thought_index == -1:
|
||||
return ""
|
||||
thought = text[:thought_index].strip()
|
||||
# Remove any triple backticks from the thought string
|
||||
thought = thought.replace("```", "").strip()
|
||||
return thought
|
||||
regex = r"(.*?)(?:\n\nAction|\n\nFinal Answer)"
|
||||
thought_match = re.search(regex, text, re.DOTALL)
|
||||
if thought_match:
|
||||
thought = thought_match.group(1).strip()
|
||||
# Remove any triple backticks from the thought string
|
||||
thought = thought.replace("```", "").strip()
|
||||
return thought
|
||||
return ""
|
||||
|
||||
def _clean_action(self, text: str) -> str:
|
||||
"""Clean action string by removing non-essential formatting characters."""
|
||||
|
||||
@@ -203,6 +203,7 @@ def install(context):
|
||||
@crewai.command()
|
||||
def run():
|
||||
"""Run the Crew."""
|
||||
click.echo("Running the Crew")
|
||||
run_crew()
|
||||
|
||||
|
||||
|
||||
@@ -216,43 +216,10 @@ MODELS = {
|
||||
"watsonx/ibm/granite-3-8b-instruct",
|
||||
],
|
||||
"bedrock": [
|
||||
"bedrock/us.amazon.nova-pro-v1:0",
|
||||
"bedrock/us.amazon.nova-micro-v1:0",
|
||||
"bedrock/us.amazon.nova-lite-v1:0",
|
||||
"bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0",
|
||||
"bedrock/us.anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||
"bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/us.anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
"bedrock/us.anthropic.claude-3-opus-20240229-v1:0",
|
||||
"bedrock/us.anthropic.claude-3-haiku-20240307-v1:0",
|
||||
"bedrock/us.meta.llama3-2-11b-instruct-v1:0",
|
||||
"bedrock/us.meta.llama3-2-3b-instruct-v1:0",
|
||||
"bedrock/us.meta.llama3-2-90b-instruct-v1:0",
|
||||
"bedrock/us.meta.llama3-2-1b-instruct-v1:0",
|
||||
"bedrock/us.meta.llama3-1-8b-instruct-v1:0",
|
||||
"bedrock/us.meta.llama3-1-70b-instruct-v1:0",
|
||||
"bedrock/us.meta.llama3-3-70b-instruct-v1:0",
|
||||
"bedrock/us.meta.llama3-1-405b-instruct-v1:0",
|
||||
"bedrock/eu.anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"bedrock/eu.anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
"bedrock/eu.anthropic.claude-3-haiku-20240307-v1:0",
|
||||
"bedrock/eu.meta.llama3-2-3b-instruct-v1:0",
|
||||
"bedrock/eu.meta.llama3-2-1b-instruct-v1:0",
|
||||
"bedrock/apac.anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"bedrock/apac.anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||
"bedrock/apac.anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
"bedrock/apac.anthropic.claude-3-haiku-20240307-v1:0",
|
||||
"bedrock/amazon.nova-pro-v1:0",
|
||||
"bedrock/amazon.nova-micro-v1:0",
|
||||
"bedrock/amazon.nova-lite-v1:0",
|
||||
"bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"bedrock/anthropic.claude-3-5-haiku-20241022-v1:0",
|
||||
"bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||
"bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
"bedrock/anthropic.claude-3-opus-20240229-v1:0",
|
||||
"bedrock/anthropic.claude-3-haiku-20240307-v1:0",
|
||||
"bedrock/anthropic.claude-3-opus-20240229-v1:0",
|
||||
"bedrock/anthropic.claude-v2:1",
|
||||
"bedrock/anthropic.claude-v2",
|
||||
"bedrock/anthropic.claude-instant-v1",
|
||||
@@ -267,6 +234,8 @@ MODELS = {
|
||||
"bedrock/ai21.j2-mid-v1",
|
||||
"bedrock/ai21.j2-ultra-v1",
|
||||
"bedrock/ai21.jamba-instruct-v1:0",
|
||||
"bedrock/meta.llama2-13b-chat-v1",
|
||||
"bedrock/meta.llama2-70b-chat-v1",
|
||||
"bedrock/mistral.mistral-7b-instruct-v0:2",
|
||||
"bedrock/mistral.mixtral-8x7b-instruct-v0:1",
|
||||
],
|
||||
|
||||
@@ -1,6 +1,4 @@
|
||||
import subprocess
|
||||
from enum import Enum
|
||||
from typing import List, Optional
|
||||
|
||||
import click
|
||||
from packaging import version
|
||||
@@ -9,24 +7,16 @@ from crewai.cli.utils import read_toml
|
||||
from crewai.cli.version import get_crewai_version
|
||||
|
||||
|
||||
class CrewType(Enum):
|
||||
STANDARD = "standard"
|
||||
FLOW = "flow"
|
||||
|
||||
|
||||
def run_crew() -> None:
|
||||
"""
|
||||
Run the crew or flow by running a command in the UV environment.
|
||||
|
||||
Starting from version 0.103.0, this command can be used to run both
|
||||
standard crews and flows. For flows, it detects the type from pyproject.toml
|
||||
and automatically runs the appropriate command.
|
||||
Run the crew by running a command in the UV environment.
|
||||
"""
|
||||
command = ["uv", "run", "run_crew"]
|
||||
crewai_version = get_crewai_version()
|
||||
min_required_version = "0.71.0"
|
||||
|
||||
pyproject_data = read_toml()
|
||||
|
||||
# Check for legacy poetry configuration
|
||||
if pyproject_data.get("tool", {}).get("poetry") and (
|
||||
version.parse(crewai_version) < version.parse(min_required_version)
|
||||
):
|
||||
@@ -36,54 +26,18 @@ def run_crew() -> None:
|
||||
fg="red",
|
||||
)
|
||||
|
||||
# Determine crew type
|
||||
is_flow = pyproject_data.get("tool", {}).get("crewai", {}).get("type") == "flow"
|
||||
crew_type = CrewType.FLOW if is_flow else CrewType.STANDARD
|
||||
|
||||
# Display appropriate message
|
||||
click.echo(f"Running the {'Flow' if is_flow else 'Crew'}")
|
||||
|
||||
# Execute the appropriate command
|
||||
execute_command(crew_type)
|
||||
|
||||
|
||||
def execute_command(crew_type: CrewType) -> None:
|
||||
"""
|
||||
Execute the appropriate command based on crew type.
|
||||
|
||||
Args:
|
||||
crew_type: The type of crew to run
|
||||
"""
|
||||
command = ["uv", "run", "kickoff" if crew_type == CrewType.FLOW else "run_crew"]
|
||||
|
||||
try:
|
||||
subprocess.run(command, capture_output=False, text=True, check=True)
|
||||
|
||||
except subprocess.CalledProcessError as e:
|
||||
handle_error(e, crew_type)
|
||||
click.echo(f"An error occurred while running the crew: {e}", err=True)
|
||||
click.echo(e.output, err=True, nl=True)
|
||||
|
||||
if pyproject_data.get("tool", {}).get("poetry"):
|
||||
click.secho(
|
||||
"It's possible that you are using an old version of crewAI that uses poetry, please run `crewai update` to update your pyproject.toml to use uv.",
|
||||
fg="yellow",
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
click.echo(f"An unexpected error occurred: {e}", err=True)
|
||||
|
||||
|
||||
def handle_error(error: subprocess.CalledProcessError, crew_type: CrewType) -> None:
|
||||
"""
|
||||
Handle subprocess errors with appropriate messaging.
|
||||
|
||||
Args:
|
||||
error: The subprocess error that occurred
|
||||
crew_type: The type of crew that was being run
|
||||
"""
|
||||
entity_type = "flow" if crew_type == CrewType.FLOW else "crew"
|
||||
click.echo(f"An error occurred while running the {entity_type}: {error}", err=True)
|
||||
|
||||
if error.output:
|
||||
click.echo(error.output, err=True, nl=True)
|
||||
|
||||
pyproject_data = read_toml()
|
||||
if pyproject_data.get("tool", {}).get("poetry"):
|
||||
click.secho(
|
||||
"It's possible that you are using an old version of crewAI that uses poetry, "
|
||||
"please run `crewai update` to update your pyproject.toml to use uv.",
|
||||
fg="yellow",
|
||||
)
|
||||
|
||||
@@ -30,13 +30,13 @@ crewai install
|
||||
|
||||
## Running the Project
|
||||
|
||||
To kickstart your flow and begin execution, run this from the root folder of your project:
|
||||
To kickstart your crew of AI agents and begin task execution, run this from the root folder of your project:
|
||||
|
||||
```bash
|
||||
crewai run
|
||||
```
|
||||
|
||||
This command initializes the {{name}} Flow as defined in your configuration.
|
||||
This command initializes the {{name}} Crew, assembling the agents and assigning them tasks as defined in your configuration.
|
||||
|
||||
This example, unmodified, will run the create a `report.md` file with the output of a research on LLMs in the root folder.
|
||||
|
||||
|
||||
@@ -257,11 +257,11 @@ def get_crew(crew_path: str = "crew.py", require: bool = False) -> Crew | None:
|
||||
import os
|
||||
|
||||
for root, _, files in os.walk("."):
|
||||
if crew_path in files:
|
||||
crew_os_path = os.path.join(root, crew_path)
|
||||
if "crew.py" in files:
|
||||
crew_path = os.path.join(root, "crew.py")
|
||||
try:
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"crew_module", crew_os_path
|
||||
"crew_module", crew_path
|
||||
)
|
||||
if not spec or not spec.loader:
|
||||
continue
|
||||
@@ -273,11 +273,9 @@ def get_crew(crew_path: str = "crew.py", require: bool = False) -> Crew | None:
|
||||
for attr_name in dir(module):
|
||||
attr = getattr(module, attr_name)
|
||||
try:
|
||||
if isinstance(attr, Crew) and hasattr(attr, "kickoff"):
|
||||
print(
|
||||
f"Found valid crew object in attribute '{attr_name}' at {crew_os_path}."
|
||||
)
|
||||
return attr
|
||||
if callable(attr) and hasattr(attr, "crew"):
|
||||
crew_instance = attr().crew()
|
||||
return crew_instance
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing attribute {attr_name}: {e}")
|
||||
|
||||
@@ -37,6 +37,7 @@ from crewai.tasks.conditional_task import ConditionalTask
|
||||
from crewai.tasks.task_output import TaskOutput
|
||||
from crewai.tools.agent_tools.agent_tools import AgentTools
|
||||
from crewai.tools.base_tool import Tool
|
||||
from crewai.traces.unified_trace_controller import init_crew_main_trace
|
||||
from crewai.types.usage_metrics import UsageMetrics
|
||||
from crewai.utilities import I18N, FileHandler, Logger, RPMController
|
||||
from crewai.utilities.constants import TRAINING_DATA_FILE
|
||||
@@ -570,6 +571,7 @@ class Crew(BaseModel):
|
||||
CrewTrainingHandler(filename).clear()
|
||||
raise
|
||||
|
||||
@init_crew_main_trace
|
||||
def kickoff(
|
||||
self,
|
||||
inputs: Optional[Dict[str, Any]] = None,
|
||||
@@ -600,7 +602,6 @@ class Crew(BaseModel):
|
||||
agent.i18n = i18n
|
||||
# type: ignore[attr-defined] # Argument 1 to "_interpolate_inputs" of "Crew" has incompatible type "dict[str, Any] | None"; expected "dict[str, Any]"
|
||||
agent.crew = self # type: ignore[attr-defined]
|
||||
agent.set_knowledge(crew_embedder=self.embedder)
|
||||
# TODO: Create an AgentFunctionCalling protocol for future refactoring
|
||||
if not agent.function_calling_llm: # type: ignore # "BaseAgent" has no attribute "function_calling_llm"
|
||||
agent.function_calling_llm = self.function_calling_llm # type: ignore # "BaseAgent" has no attribute "function_calling_llm"
|
||||
|
||||
@@ -22,6 +22,10 @@ from pydantic import BaseModel, Field, ValidationError
|
||||
from crewai.flow.flow_visualizer import plot_flow
|
||||
from crewai.flow.persistence.base import FlowPersistence
|
||||
from crewai.flow.utils import get_possible_return_constants
|
||||
from crewai.traces.unified_trace_controller import (
|
||||
init_flow_main_trace,
|
||||
trace_flow_step,
|
||||
)
|
||||
from crewai.utilities.events.crewai_event_bus import crewai_event_bus
|
||||
from crewai.utilities.events.flow_events import (
|
||||
FlowCreatedEvent,
|
||||
@@ -721,6 +725,7 @@ class Flow(Generic[T], metaclass=FlowMeta):
|
||||
|
||||
return asyncio.run(run_flow())
|
||||
|
||||
@init_flow_main_trace
|
||||
async def kickoff_async(self, inputs: Optional[Dict[str, Any]] = None) -> Any:
|
||||
"""
|
||||
Start the flow execution asynchronously.
|
||||
@@ -777,17 +782,18 @@ class Flow(Generic[T], metaclass=FlowMeta):
|
||||
f"Flow started with ID: {self.flow_id}", color="bold_magenta"
|
||||
)
|
||||
|
||||
if inputs is not None and "id" not in inputs:
|
||||
self._initialize_state(inputs)
|
||||
if not self._start_methods:
|
||||
raise ValueError("No start method defined")
|
||||
|
||||
# Execute all start methods concurrently.
|
||||
tasks = [
|
||||
self._execute_start_method(start_method)
|
||||
for start_method in self._start_methods
|
||||
]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
final_output = self._method_outputs[-1] if self._method_outputs else None
|
||||
|
||||
# Emit FlowFinishedEvent after all processing is complete.
|
||||
crewai_event_bus.emit(
|
||||
self,
|
||||
FlowFinishedEvent(
|
||||
@@ -796,7 +802,6 @@ class Flow(Generic[T], metaclass=FlowMeta):
|
||||
result=final_output,
|
||||
),
|
||||
)
|
||||
|
||||
return final_output
|
||||
|
||||
async def _execute_start_method(self, start_method_name: str) -> None:
|
||||
@@ -822,6 +827,7 @@ class Flow(Generic[T], metaclass=FlowMeta):
|
||||
)
|
||||
await self._execute_listeners(start_method_name, result)
|
||||
|
||||
@trace_flow_step
|
||||
async def _execute_method(
|
||||
self, method_name: str, method: Callable, *args: Any, **kwargs: Any
|
||||
) -> Any:
|
||||
@@ -894,45 +900,35 @@ class Flow(Generic[T], metaclass=FlowMeta):
|
||||
Notes
|
||||
-----
|
||||
- Routers are executed sequentially to maintain flow control
|
||||
- Each router's result becomes a new trigger_method
|
||||
- Each router's result becomes the new trigger_method
|
||||
- Normal listeners are executed in parallel for efficiency
|
||||
- Listeners can receive the trigger method's result as a parameter
|
||||
"""
|
||||
# First, handle routers repeatedly until no router triggers anymore
|
||||
router_results = []
|
||||
current_trigger = trigger_method
|
||||
|
||||
while True:
|
||||
routers_triggered = self._find_triggered_methods(
|
||||
current_trigger, router_only=True
|
||||
trigger_method, router_only=True
|
||||
)
|
||||
if not routers_triggered:
|
||||
break
|
||||
|
||||
for router_name in routers_triggered:
|
||||
await self._execute_single_listener(router_name, result)
|
||||
# After executing router, the router's result is the path
|
||||
router_result = self._method_outputs[-1]
|
||||
if router_result: # Only add non-None results
|
||||
router_results.append(router_result)
|
||||
current_trigger = (
|
||||
router_result # Update for next iteration of router chain
|
||||
)
|
||||
# The last router executed sets the trigger_method
|
||||
# The router result is the last element in self._method_outputs
|
||||
trigger_method = self._method_outputs[-1]
|
||||
|
||||
# Now execute normal listeners for all router results and the original trigger
|
||||
all_triggers = [trigger_method] + router_results
|
||||
|
||||
for current_trigger in all_triggers:
|
||||
if current_trigger: # Skip None results
|
||||
listeners_triggered = self._find_triggered_methods(
|
||||
current_trigger, router_only=False
|
||||
)
|
||||
if listeners_triggered:
|
||||
tasks = [
|
||||
self._execute_single_listener(listener_name, result)
|
||||
for listener_name in listeners_triggered
|
||||
]
|
||||
await asyncio.gather(*tasks)
|
||||
# Now that no more routers are triggered by current trigger_method,
|
||||
# execute normal listeners
|
||||
listeners_triggered = self._find_triggered_methods(
|
||||
trigger_method, router_only=False
|
||||
)
|
||||
if listeners_triggered:
|
||||
tasks = [
|
||||
self._execute_single_listener(listener_name, result)
|
||||
for listener_name in listeners_triggered
|
||||
]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
def _find_triggered_methods(
|
||||
self, trigger_method: str, router_only: bool
|
||||
|
||||
@@ -4,7 +4,7 @@ SQLite-based implementation of flow state persistence.
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
@@ -34,7 +34,6 @@ class SQLiteFlowPersistence(FlowPersistence):
|
||||
ValueError: If db_path is invalid
|
||||
"""
|
||||
from crewai.utilities.paths import db_storage_path
|
||||
|
||||
# Get path from argument or default location
|
||||
path = db_path or str(Path(db_storage_path()) / "flow_states.db")
|
||||
|
||||
@@ -47,8 +46,7 @@ class SQLiteFlowPersistence(FlowPersistence):
|
||||
def init_db(self) -> None:
|
||||
"""Create the necessary tables if they don't exist."""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS flow_states (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
flow_uuid TEXT NOT NULL,
|
||||
@@ -56,15 +54,12 @@ class SQLiteFlowPersistence(FlowPersistence):
|
||||
timestamp DATETIME NOT NULL,
|
||||
state_json TEXT NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
""")
|
||||
# Add index for faster UUID lookups
|
||||
conn.execute(
|
||||
"""
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_flow_states_uuid
|
||||
ON flow_states(flow_uuid)
|
||||
"""
|
||||
)
|
||||
""")
|
||||
|
||||
def save_state(
|
||||
self,
|
||||
@@ -90,22 +85,19 @@ class SQLiteFlowPersistence(FlowPersistence):
|
||||
)
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
conn.execute("""
|
||||
INSERT INTO flow_states (
|
||||
flow_uuid,
|
||||
method_name,
|
||||
timestamp,
|
||||
state_json
|
||||
) VALUES (?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
flow_uuid,
|
||||
method_name,
|
||||
datetime.now(timezone.utc).isoformat(),
|
||||
json.dumps(state_dict),
|
||||
),
|
||||
)
|
||||
""", (
|
||||
flow_uuid,
|
||||
method_name,
|
||||
datetime.utcnow().isoformat(),
|
||||
json.dumps(state_dict),
|
||||
))
|
||||
|
||||
def load_state(self, flow_uuid: str) -> Optional[Dict[str, Any]]:
|
||||
"""Load the most recent state for a given flow UUID.
|
||||
@@ -117,16 +109,13 @@ class SQLiteFlowPersistence(FlowPersistence):
|
||||
The most recent state as a dictionary, or None if no state exists
|
||||
"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
cursor = conn.execute("""
|
||||
SELECT state_json
|
||||
FROM flow_states
|
||||
WHERE flow_uuid = ?
|
||||
ORDER BY id DESC
|
||||
LIMIT 1
|
||||
""",
|
||||
(flow_uuid,),
|
||||
)
|
||||
""", (flow_uuid,))
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row:
|
||||
|
||||
@@ -16,8 +16,7 @@ Example
|
||||
import ast
|
||||
import inspect
|
||||
import textwrap
|
||||
from collections import defaultdict, deque
|
||||
from typing import Any, Deque, Dict, List, Optional, Set, Union
|
||||
from typing import Any, Dict, List, Optional, Set, Union
|
||||
|
||||
|
||||
def get_possible_return_constants(function: Any) -> Optional[List[str]]:
|
||||
@@ -119,7 +118,7 @@ def calculate_node_levels(flow: Any) -> Dict[str, int]:
|
||||
- Processes router paths separately
|
||||
"""
|
||||
levels: Dict[str, int] = {}
|
||||
queue: Deque[str] = deque()
|
||||
queue: List[str] = []
|
||||
visited: Set[str] = set()
|
||||
pending_and_listeners: Dict[str, Set[str]] = {}
|
||||
|
||||
@@ -129,35 +128,28 @@ def calculate_node_levels(flow: Any) -> Dict[str, int]:
|
||||
levels[method_name] = 0
|
||||
queue.append(method_name)
|
||||
|
||||
# Precompute listener dependencies
|
||||
or_listeners = defaultdict(list)
|
||||
and_listeners = defaultdict(set)
|
||||
for listener_name, (condition_type, trigger_methods) in flow._listeners.items():
|
||||
if condition_type == "OR":
|
||||
for method in trigger_methods:
|
||||
or_listeners[method].append(listener_name)
|
||||
elif condition_type == "AND":
|
||||
and_listeners[listener_name] = set(trigger_methods)
|
||||
|
||||
# Breadth-first traversal to assign levels
|
||||
while queue:
|
||||
current = queue.popleft()
|
||||
current = queue.pop(0)
|
||||
current_level = levels[current]
|
||||
visited.add(current)
|
||||
|
||||
for listener_name in or_listeners[current]:
|
||||
if listener_name not in levels or levels[listener_name] > current_level + 1:
|
||||
levels[listener_name] = current_level + 1
|
||||
if listener_name not in visited:
|
||||
queue.append(listener_name)
|
||||
|
||||
for listener_name, required_methods in and_listeners.items():
|
||||
if current in required_methods:
|
||||
for listener_name, (condition_type, trigger_methods) in flow._listeners.items():
|
||||
if condition_type == "OR":
|
||||
if current in trigger_methods:
|
||||
if (
|
||||
listener_name not in levels
|
||||
or levels[listener_name] > current_level + 1
|
||||
):
|
||||
levels[listener_name] = current_level + 1
|
||||
if listener_name not in visited:
|
||||
queue.append(listener_name)
|
||||
elif condition_type == "AND":
|
||||
if listener_name not in pending_and_listeners:
|
||||
pending_and_listeners[listener_name] = set()
|
||||
pending_and_listeners[listener_name].add(current)
|
||||
|
||||
if required_methods == pending_and_listeners[listener_name]:
|
||||
if current in trigger_methods:
|
||||
pending_and_listeners[listener_name].add(current)
|
||||
if set(trigger_methods) == pending_and_listeners[listener_name]:
|
||||
if (
|
||||
listener_name not in levels
|
||||
or levels[listener_name] > current_level + 1
|
||||
@@ -167,7 +159,22 @@ def calculate_node_levels(flow: Any) -> Dict[str, int]:
|
||||
queue.append(listener_name)
|
||||
|
||||
# Handle router connections
|
||||
process_router_paths(flow, current, current_level, levels, queue)
|
||||
if current in flow._routers:
|
||||
router_method_name = current
|
||||
paths = flow._router_paths.get(router_method_name, [])
|
||||
for path in paths:
|
||||
for listener_name, (
|
||||
condition_type,
|
||||
trigger_methods,
|
||||
) in flow._listeners.items():
|
||||
if path in trigger_methods:
|
||||
if (
|
||||
listener_name not in levels
|
||||
or levels[listener_name] > current_level + 1
|
||||
):
|
||||
levels[listener_name] = current_level + 1
|
||||
if listener_name not in visited:
|
||||
queue.append(listener_name)
|
||||
|
||||
return levels
|
||||
|
||||
@@ -220,7 +227,10 @@ def build_ancestor_dict(flow: Any) -> Dict[str, Set[str]]:
|
||||
|
||||
|
||||
def dfs_ancestors(
|
||||
node: str, ancestors: Dict[str, Set[str]], visited: Set[str], flow: Any
|
||||
node: str,
|
||||
ancestors: Dict[str, Set[str]],
|
||||
visited: Set[str],
|
||||
flow: Any
|
||||
) -> None:
|
||||
"""
|
||||
Perform depth-first search to build ancestor relationships.
|
||||
@@ -264,9 +274,7 @@ def dfs_ancestors(
|
||||
dfs_ancestors(listener_name, ancestors, visited, flow)
|
||||
|
||||
|
||||
def is_ancestor(
|
||||
node: str, ancestor_candidate: str, ancestors: Dict[str, Set[str]]
|
||||
) -> bool:
|
||||
def is_ancestor(node: str, ancestor_candidate: str, ancestors: Dict[str, Set[str]]) -> bool:
|
||||
"""
|
||||
Check if one node is an ancestor of another.
|
||||
|
||||
@@ -331,9 +339,7 @@ def build_parent_children_dict(flow: Any) -> Dict[str, List[str]]:
|
||||
return parent_children
|
||||
|
||||
|
||||
def get_child_index(
|
||||
parent: str, child: str, parent_children: Dict[str, List[str]]
|
||||
) -> int:
|
||||
def get_child_index(parent: str, child: str, parent_children: Dict[str, List[str]]) -> int:
|
||||
"""
|
||||
Get the index of a child node in its parent's sorted children list.
|
||||
|
||||
@@ -354,23 +360,3 @@ def get_child_index(
|
||||
children = parent_children.get(parent, [])
|
||||
children.sort()
|
||||
return children.index(child)
|
||||
|
||||
|
||||
def process_router_paths(flow, current, current_level, levels, queue):
|
||||
"""
|
||||
Handle the router connections for the current node.
|
||||
"""
|
||||
if current in flow._routers:
|
||||
paths = flow._router_paths.get(current, [])
|
||||
for path in paths:
|
||||
for listener_name, (
|
||||
condition_type,
|
||||
trigger_methods,
|
||||
) in flow._listeners.items():
|
||||
if path in trigger_methods:
|
||||
if (
|
||||
listener_name not in levels
|
||||
or levels[listener_name] > current_level + 1
|
||||
):
|
||||
levels[listener_name] = current_level + 1
|
||||
queue.append(listener_name)
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
@@ -5,7 +6,17 @@ import sys
|
||||
import threading
|
||||
import warnings
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Dict, List, Literal, Optional, Type, Union, cast
|
||||
from typing import (
|
||||
Any,
|
||||
Dict,
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Tuple,
|
||||
Type,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel
|
||||
@@ -26,10 +37,12 @@ with warnings.catch_warnings():
|
||||
from litellm.utils import get_supported_openai_params, supports_response_schema
|
||||
|
||||
|
||||
from crewai.traces.unified_trace_controller import trace_llm_call
|
||||
from crewai.utilities.events import crewai_event_bus
|
||||
from crewai.utilities.exceptions.context_window_exceeding_exception import (
|
||||
LLMContextLengthExceededException,
|
||||
)
|
||||
from crewai.utilities.protocols import AgentExecutorProtocol
|
||||
|
||||
load_dotenv()
|
||||
|
||||
@@ -64,7 +77,6 @@ LLM_CONTEXT_WINDOW_SIZES = {
|
||||
"gpt-4-turbo": 128000,
|
||||
"o1-preview": 128000,
|
||||
"o1-mini": 128000,
|
||||
"o3-mini": 200000, # Based on official o3-mini specifications
|
||||
# gemini
|
||||
"gemini-2.0-flash": 1048576,
|
||||
"gemini-1.5-pro": 2097152,
|
||||
@@ -174,6 +186,7 @@ class LLM:
|
||||
self.context_window_size = 0
|
||||
self.reasoning_effort = reasoning_effort
|
||||
self.additional_params = kwargs
|
||||
self._message_history: List[Dict[str, str]] = []
|
||||
self.is_anthropic = self._is_anthropic_model(model)
|
||||
|
||||
litellm.drop_params = True
|
||||
@@ -189,6 +202,12 @@ class LLM:
|
||||
self.set_callbacks(callbacks)
|
||||
self.set_env_callbacks()
|
||||
|
||||
@trace_llm_call
|
||||
def _call_llm(self, params: Dict[str, Any]) -> Any:
|
||||
with suppress_warnings():
|
||||
response = litellm.completion(**params)
|
||||
return response
|
||||
|
||||
def _is_anthropic_model(self, model: str) -> bool:
|
||||
"""Determine if the model is from Anthropic provider.
|
||||
|
||||
@@ -307,7 +326,7 @@ class LLM:
|
||||
params = {k: v for k, v in params.items() if v is not None}
|
||||
|
||||
# --- 2) Make the completion call
|
||||
response = litellm.completion(**params)
|
||||
response = self._call_llm(params)
|
||||
response_message = cast(Choices, cast(ModelResponse, response).choices)[
|
||||
0
|
||||
].message
|
||||
@@ -486,23 +505,10 @@ class LLM:
|
||||
"""
|
||||
Returns the context window size, using 75% of the maximum to avoid
|
||||
cutting off messages mid-thread.
|
||||
|
||||
Raises:
|
||||
ValueError: If a model's context window size is outside valid bounds (1024-2097152)
|
||||
"""
|
||||
if self.context_window_size != 0:
|
||||
return self.context_window_size
|
||||
|
||||
MIN_CONTEXT = 1024
|
||||
MAX_CONTEXT = 2097152 # Current max from gemini-1.5-pro
|
||||
|
||||
# Validate all context window sizes
|
||||
for key, value in LLM_CONTEXT_WINDOW_SIZES.items():
|
||||
if value < MIN_CONTEXT or value > MAX_CONTEXT:
|
||||
raise ValueError(
|
||||
f"Context window for {key} must be between {MIN_CONTEXT} and {MAX_CONTEXT}"
|
||||
)
|
||||
|
||||
self.context_window_size = int(
|
||||
DEFAULT_CONTEXT_WINDOW_SIZE * CONTEXT_WINDOW_USAGE_RATIO
|
||||
)
|
||||
@@ -564,3 +570,95 @@ class LLM:
|
||||
|
||||
litellm.success_callback = success_callbacks
|
||||
litellm.failure_callback = failure_callbacks
|
||||
|
||||
def _get_execution_context(self) -> Tuple[Optional[Any], Optional[Any]]:
|
||||
"""Get the agent and task from the execution context.
|
||||
|
||||
Returns:
|
||||
tuple: (agent, task) from any AgentExecutor context, or (None, None) if not found
|
||||
"""
|
||||
frame = inspect.currentframe()
|
||||
caller_frame = frame.f_back if frame else None
|
||||
agent = None
|
||||
task = None
|
||||
|
||||
# Add a maximum depth to prevent infinite loops
|
||||
max_depth = 100 # Reasonable limit for call stack depth
|
||||
current_depth = 0
|
||||
|
||||
while caller_frame and current_depth < max_depth:
|
||||
if "self" in caller_frame.f_locals:
|
||||
caller_self = caller_frame.f_locals["self"]
|
||||
if isinstance(caller_self, AgentExecutorProtocol):
|
||||
agent = caller_self.agent
|
||||
task = caller_self.task
|
||||
break
|
||||
caller_frame = caller_frame.f_back
|
||||
current_depth += 1
|
||||
|
||||
return agent, task
|
||||
|
||||
def _get_new_messages(self, messages: List[Dict[str, str]]) -> List[Dict[str, str]]:
|
||||
"""Get only the new messages that haven't been processed before."""
|
||||
if not hasattr(self, "_message_history"):
|
||||
self._message_history = []
|
||||
|
||||
new_messages = []
|
||||
for message in messages:
|
||||
message_key = (message["role"], message["content"])
|
||||
if message_key not in [
|
||||
(m["role"], m["content"]) for m in self._message_history
|
||||
]:
|
||||
new_messages.append(message)
|
||||
self._message_history.append(message)
|
||||
return new_messages
|
||||
|
||||
def _get_new_tool_results(self, agent) -> List[Dict]:
|
||||
"""Get only the new tool results that haven't been processed before."""
|
||||
if not agent or not agent.tools_results:
|
||||
return []
|
||||
|
||||
if not hasattr(self, "_tool_results_history"):
|
||||
self._tool_results_history: List[Dict] = []
|
||||
|
||||
new_tool_results = []
|
||||
|
||||
for result in agent.tools_results:
|
||||
# Process tool arguments to extract actual values
|
||||
processed_args = {}
|
||||
if isinstance(result["tool_args"], dict):
|
||||
for key, value in result["tool_args"].items():
|
||||
if isinstance(value, dict) and "type" in value:
|
||||
# Skip metadata and just store the actual value
|
||||
continue
|
||||
processed_args[key] = value
|
||||
|
||||
# Create a clean result with processed arguments
|
||||
clean_result = {
|
||||
"tool_name": result["tool_name"],
|
||||
"tool_args": processed_args,
|
||||
"result": result["result"],
|
||||
"content": result.get("content", ""),
|
||||
"start_time": result.get("start_time", ""),
|
||||
}
|
||||
|
||||
# Check if this exact tool execution exists in history
|
||||
is_duplicate = False
|
||||
for history_result in self._tool_results_history:
|
||||
if (
|
||||
clean_result["tool_name"] == history_result["tool_name"]
|
||||
and str(clean_result["tool_args"])
|
||||
== str(history_result["tool_args"])
|
||||
and str(clean_result["result"]) == str(history_result["result"])
|
||||
and clean_result["content"] == history_result.get("content", "")
|
||||
and clean_result["start_time"]
|
||||
== history_result.get("start_time", "")
|
||||
):
|
||||
is_duplicate = True
|
||||
break
|
||||
|
||||
if not is_duplicate:
|
||||
new_tool_results.append(clean_result)
|
||||
self._tool_results_history.append(clean_result)
|
||||
|
||||
return new_tool_results
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from pydantic import PrivateAttr
|
||||
@@ -53,34 +52,10 @@ class ShortTermMemory(Memory):
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
agent: Optional[str] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Save a memory item to the storage.
|
||||
|
||||
Args:
|
||||
value: The data to save.
|
||||
metadata: Optional metadata to associate with the memory.
|
||||
agent: Optional agent identifier.
|
||||
|
||||
Raises:
|
||||
ValueError: If the item's timestamp is in the future.
|
||||
"""
|
||||
import logging
|
||||
|
||||
item = ShortTermMemoryItem(data=value, metadata=metadata, agent=agent)
|
||||
|
||||
if item.timestamp > datetime.now():
|
||||
raise ValueError("Cannot save memory item with future timestamp")
|
||||
|
||||
logging.debug(f"Saving memory item with timestamp: {item.timestamp}")
|
||||
|
||||
if self._memory_provider == "mem0":
|
||||
item.data = f"Remember the following insights from Agent run: {item.data}"
|
||||
|
||||
# Include timestamp in metadata
|
||||
if item.metadata is None:
|
||||
item.metadata = {}
|
||||
item.metadata["timestamp"] = item.timestamp.isoformat()
|
||||
|
||||
super().save(value=item.data, metadata=item.metadata, agent=item.agent)
|
||||
|
||||
def search(
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
|
||||
@@ -8,11 +7,7 @@ class ShortTermMemoryItem:
|
||||
data: Any,
|
||||
agent: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
timestamp: Optional[datetime] = None,
|
||||
):
|
||||
if timestamp is not None and timestamp > datetime.now():
|
||||
raise ValueError("Timestamp cannot be in the future")
|
||||
self.data = data
|
||||
self.agent = agent
|
||||
self.metadata = metadata if metadata is not None else {}
|
||||
self.timestamp = timestamp if timestamp is not None else datetime.now()
|
||||
|
||||
@@ -114,32 +114,13 @@ class RAGStorage(BaseRAGStorage):
|
||||
limit: int = 3,
|
||||
filter: Optional[dict] = None,
|
||||
score_threshold: float = 0.35,
|
||||
recency_weight: float = 0.3,
|
||||
time_decay_days: float = 1.0,
|
||||
) -> List[Any]:
|
||||
"""
|
||||
Search for entries in the storage based on semantic similarity and recency.
|
||||
|
||||
Args:
|
||||
query: The search query string.
|
||||
limit: Maximum number of results to return.
|
||||
filter: Optional filter to apply to the search.
|
||||
score_threshold: Minimum score threshold for results.
|
||||
recency_weight: Weight given to recency vs. semantic similarity (0.0-1.0).
|
||||
Higher values prioritize recent memories more strongly.
|
||||
time_decay_days: Number of days over which recency factor decays to zero.
|
||||
Smaller values make older memories lose relevance faster.
|
||||
|
||||
Returns:
|
||||
List of search results, each containing id, metadata, context, and score.
|
||||
Results are sorted by combined semantic similarity and recency score.
|
||||
"""
|
||||
if not hasattr(self, "app"):
|
||||
self._initialize_app()
|
||||
|
||||
try:
|
||||
with suppress_logging():
|
||||
response = self.collection.query(query_texts=query, n_results=limit * 2) # Get more results to allow for recency filtering
|
||||
response = self.collection.query(query_texts=query, n_results=limit)
|
||||
|
||||
results = []
|
||||
for i in range(len(response["ids"][0])):
|
||||
@@ -149,27 +130,10 @@ class RAGStorage(BaseRAGStorage):
|
||||
"context": response["documents"][0][i],
|
||||
"score": response["distances"][0][i],
|
||||
}
|
||||
|
||||
# Apply recency boost if timestamp exists in metadata
|
||||
if "timestamp" in result["metadata"]:
|
||||
try:
|
||||
from datetime import datetime
|
||||
timestamp = datetime.fromisoformat(result["metadata"]["timestamp"])
|
||||
now = datetime.now()
|
||||
# Calculate recency factor (newer = higher score)
|
||||
time_diff_seconds = (now - timestamp).total_seconds()
|
||||
recency_factor = max(0, 1 - (time_diff_seconds / (time_decay_days * 24 * 60 * 60)))
|
||||
# Adjust score with recency factor
|
||||
result["score"] = result["score"] * (1 - recency_weight) + recency_factor * recency_weight
|
||||
except (ValueError, TypeError):
|
||||
pass # If timestamp parsing fails, use original score
|
||||
|
||||
if result["score"] >= score_threshold:
|
||||
results.append(result)
|
||||
|
||||
# Sort by adjusted score (higher is better)
|
||||
results.sort(key=lambda x: x["score"], reverse=True)
|
||||
return results[:limit] # Return only the requested number of results
|
||||
return results
|
||||
except Exception as e:
|
||||
logging.error(f"Error during {self.type} search: {str(e)}")
|
||||
return []
|
||||
|
||||
@@ -2,6 +2,7 @@ import ast
|
||||
import datetime
|
||||
import json
|
||||
import time
|
||||
from datetime import UTC
|
||||
from difflib import SequenceMatcher
|
||||
from json import JSONDecodeError
|
||||
from textwrap import dedent
|
||||
@@ -117,7 +118,10 @@ class ToolUsage:
|
||||
self._printer.print(content=f"\n\n{error}\n", color="red")
|
||||
return error
|
||||
|
||||
if isinstance(tool, CrewStructuredTool) and tool.name == self._i18n.tools("add_image")["name"]: # type: ignore
|
||||
if (
|
||||
isinstance(tool, CrewStructuredTool)
|
||||
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
|
||||
@@ -154,6 +158,7 @@ class ToolUsage:
|
||||
self.task.increment_tools_errors()
|
||||
|
||||
started_at = time.time()
|
||||
started_at_trace = datetime.datetime.now(UTC)
|
||||
from_cache = False
|
||||
|
||||
result = None # type: ignore # Incompatible types in assignment (expression has type "None", variable has type "str")
|
||||
@@ -181,7 +186,9 @@ class ToolUsage:
|
||||
|
||||
if calling.arguments:
|
||||
try:
|
||||
acceptable_args = tool.args_schema.model_json_schema()["properties"].keys() # type: ignore
|
||||
acceptable_args = tool.args_schema.model_json_schema()[
|
||||
"properties"
|
||||
].keys() # type: ignore
|
||||
arguments = {
|
||||
k: v
|
||||
for k, v in calling.arguments.items()
|
||||
@@ -202,7 +209,7 @@ class ToolUsage:
|
||||
error=e, tool=tool.name, tool_inputs=tool.description
|
||||
)
|
||||
error = ToolUsageErrorException(
|
||||
f'\n{error_message}.\nMoving on then. {self._i18n.slice("format").format(tool_names=self.tools_names)}'
|
||||
f"\n{error_message}.\nMoving on then. {self._i18n.slice('format').format(tool_names=self.tools_names)}"
|
||||
).message
|
||||
self.task.increment_tools_errors()
|
||||
if self.agent.verbose:
|
||||
@@ -237,6 +244,7 @@ class ToolUsage:
|
||||
"result": result,
|
||||
"tool_name": tool.name,
|
||||
"tool_args": calling.arguments,
|
||||
"start_time": started_at_trace,
|
||||
}
|
||||
|
||||
self.on_tool_use_finished(
|
||||
@@ -380,7 +388,7 @@ class ToolUsage:
|
||||
raise
|
||||
else:
|
||||
return ToolUsageErrorException(
|
||||
f'{self._i18n.errors("tool_arguments_error")}'
|
||||
f"{self._i18n.errors('tool_arguments_error')}"
|
||||
)
|
||||
|
||||
if not isinstance(arguments, dict):
|
||||
@@ -388,7 +396,7 @@ class ToolUsage:
|
||||
raise
|
||||
else:
|
||||
return ToolUsageErrorException(
|
||||
f'{self._i18n.errors("tool_arguments_error")}'
|
||||
f"{self._i18n.errors('tool_arguments_error')}"
|
||||
)
|
||||
|
||||
return ToolCalling(
|
||||
@@ -416,7 +424,7 @@ class ToolUsage:
|
||||
if 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")
|
||||
f'{self._i18n.errors("tool_usage_error").format(error=e)}\nMoving on then. {self._i18n.slice("format").format(tool_names=self.tools_names)}'
|
||||
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)
|
||||
|
||||
|
||||
0
src/crewai/traces/__init__.py
Normal file
0
src/crewai/traces/__init__.py
Normal file
39
src/crewai/traces/context.py
Normal file
39
src/crewai/traces/context.py
Normal file
@@ -0,0 +1,39 @@
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from typing import Generator
|
||||
|
||||
|
||||
class TraceContext:
|
||||
"""Maintains the current trace context throughout the execution stack.
|
||||
|
||||
This class provides a context manager for tracking trace execution across
|
||||
async and sync code paths using ContextVars.
|
||||
"""
|
||||
|
||||
_context: ContextVar = ContextVar("trace_context", default=None)
|
||||
|
||||
@classmethod
|
||||
def get_current(cls):
|
||||
"""Get the current trace context.
|
||||
|
||||
Returns:
|
||||
Optional[UnifiedTraceController]: The current trace controller or None if not set.
|
||||
"""
|
||||
return cls._context.get()
|
||||
|
||||
@classmethod
|
||||
@contextmanager
|
||||
def set_current(cls, trace):
|
||||
"""Set the current trace context within a context manager.
|
||||
|
||||
Args:
|
||||
trace: The trace controller to set as current.
|
||||
|
||||
Yields:
|
||||
UnifiedTraceController: The current trace controller.
|
||||
"""
|
||||
token = cls._context.set(trace)
|
||||
try:
|
||||
yield trace
|
||||
finally:
|
||||
cls._context.reset(token)
|
||||
19
src/crewai/traces/enums.py
Normal file
19
src/crewai/traces/enums.py
Normal file
@@ -0,0 +1,19 @@
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class TraceType(Enum):
|
||||
LLM_CALL = "llm_call"
|
||||
TOOL_CALL = "tool_call"
|
||||
FLOW_STEP = "flow_step"
|
||||
START_CALL = "start_call"
|
||||
|
||||
|
||||
class RunType(Enum):
|
||||
KICKOFF = "kickoff"
|
||||
TRAIN = "train"
|
||||
TEST = "test"
|
||||
|
||||
|
||||
class CrewType(Enum):
|
||||
CREW = "crew"
|
||||
FLOW = "flow"
|
||||
89
src/crewai/traces/models.py
Normal file
89
src/crewai/traces/models.py
Normal file
@@ -0,0 +1,89 @@
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class ToolCall(BaseModel):
|
||||
"""Model representing a tool call during execution"""
|
||||
|
||||
name: str
|
||||
arguments: Dict[str, Any]
|
||||
output: str
|
||||
start_time: datetime
|
||||
end_time: Optional[datetime] = None
|
||||
latency_ms: Optional[int] = None
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
class LLMRequest(BaseModel):
|
||||
"""Model representing the LLM request details"""
|
||||
|
||||
model: str
|
||||
messages: List[Dict[str, str]]
|
||||
temperature: Optional[float] = None
|
||||
max_tokens: Optional[int] = None
|
||||
stop_sequences: Optional[List[str]] = None
|
||||
additional_params: Dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class LLMResponse(BaseModel):
|
||||
"""Model representing the LLM response details"""
|
||||
|
||||
content: str
|
||||
finish_reason: Optional[str] = None
|
||||
|
||||
|
||||
class FlowStepIO(BaseModel):
|
||||
"""Model representing flow step input/output details"""
|
||||
|
||||
function_name: str
|
||||
inputs: Dict[str, Any] = Field(default_factory=dict)
|
||||
outputs: Any
|
||||
metadata: Dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class CrewTrace(BaseModel):
|
||||
"""Model for tracking detailed information about LLM interactions and Flow steps"""
|
||||
|
||||
deployment_instance_id: Optional[str] = Field(
|
||||
description="ID of the deployment instance"
|
||||
)
|
||||
trace_id: str = Field(description="Unique identifier for this trace")
|
||||
run_id: str = Field(description="Identifier for the execution run")
|
||||
agent_role: Optional[str] = Field(description="Role of the agent")
|
||||
task_id: Optional[str] = Field(description="ID of the current task being executed")
|
||||
task_name: Optional[str] = Field(description="Name of the current task")
|
||||
task_description: Optional[str] = Field(
|
||||
description="Description of the current task"
|
||||
)
|
||||
trace_type: str = Field(description="Type of the trace")
|
||||
crew_type: str = Field(description="Type of the crew")
|
||||
run_type: str = Field(description="Type of the run")
|
||||
|
||||
# Timing information
|
||||
start_time: Optional[datetime] = None
|
||||
end_time: Optional[datetime] = None
|
||||
latency_ms: Optional[int] = None
|
||||
|
||||
# Request/Response for LLM calls
|
||||
request: Optional[LLMRequest] = None
|
||||
response: Optional[LLMResponse] = None
|
||||
|
||||
# Input/Output for Flow steps
|
||||
flow_step: Optional[FlowStepIO] = None
|
||||
|
||||
# Tool usage
|
||||
tool_calls: List[ToolCall] = Field(default_factory=list)
|
||||
|
||||
# Metrics
|
||||
tokens_used: Optional[int] = None
|
||||
prompt_tokens: Optional[int] = None
|
||||
completion_tokens: Optional[int] = None
|
||||
cost: Optional[float] = None
|
||||
|
||||
# Additional metadata
|
||||
status: str = "running" # running, completed, error
|
||||
error: Optional[str] = None
|
||||
metadata: Dict[str, Any] = Field(default_factory=dict)
|
||||
tags: List[str] = Field(default_factory=list)
|
||||
543
src/crewai/traces/unified_trace_controller.py
Normal file
543
src/crewai/traces/unified_trace_controller.py
Normal file
@@ -0,0 +1,543 @@
|
||||
import inspect
|
||||
import os
|
||||
from datetime import UTC, datetime
|
||||
from functools import wraps
|
||||
from typing import Any, Awaitable, Callable, Dict, List, Optional
|
||||
from uuid import uuid4
|
||||
|
||||
from crewai.traces.context import TraceContext
|
||||
from crewai.traces.enums import CrewType, RunType, TraceType
|
||||
from crewai.traces.models import (
|
||||
CrewTrace,
|
||||
FlowStepIO,
|
||||
LLMRequest,
|
||||
LLMResponse,
|
||||
ToolCall,
|
||||
)
|
||||
|
||||
|
||||
class UnifiedTraceController:
|
||||
"""Controls and manages trace execution and recording.
|
||||
|
||||
This class handles the lifecycle of traces including creation, execution tracking,
|
||||
and recording of results for various types of operations (LLM calls, tool calls, flow steps).
|
||||
"""
|
||||
|
||||
_task_traces: Dict[str, List["UnifiedTraceController"]] = {}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
trace_type: TraceType,
|
||||
run_type: RunType,
|
||||
crew_type: CrewType,
|
||||
run_id: str,
|
||||
deployment_instance_id: str = os.environ.get(
|
||||
"CREWAI_DEPLOYMENT_INSTANCE_ID", ""
|
||||
),
|
||||
parent_trace_id: Optional[str] = None,
|
||||
agent_role: Optional[str] = "unknown",
|
||||
task_name: Optional[str] = None,
|
||||
task_description: Optional[str] = None,
|
||||
task_id: Optional[str] = None,
|
||||
flow_step: Dict[str, Any] = {},
|
||||
tool_calls: List[ToolCall] = [],
|
||||
**context: Any,
|
||||
) -> None:
|
||||
"""Initialize a new trace controller.
|
||||
|
||||
Args:
|
||||
trace_type: Type of trace being recorded.
|
||||
run_type: Type of run being executed.
|
||||
crew_type: Type of crew executing the trace.
|
||||
run_id: Unique identifier for the run.
|
||||
deployment_instance_id: Optional deployment instance identifier.
|
||||
parent_trace_id: Optional parent trace identifier for nested traces.
|
||||
agent_role: Role of the agent executing the trace.
|
||||
task_name: Optional name of the task being executed.
|
||||
task_description: Optional description of the task.
|
||||
task_id: Optional unique identifier for the task.
|
||||
flow_step: Optional flow step information.
|
||||
tool_calls: Optional list of tool calls made during execution.
|
||||
**context: Additional context parameters.
|
||||
"""
|
||||
self.trace_id = str(uuid4())
|
||||
self.run_id = run_id
|
||||
self.parent_trace_id = parent_trace_id
|
||||
self.trace_type = trace_type
|
||||
self.run_type = run_type
|
||||
self.crew_type = crew_type
|
||||
self.context = context
|
||||
self.agent_role = agent_role
|
||||
self.task_name = task_name
|
||||
self.task_description = task_description
|
||||
self.task_id = task_id
|
||||
self.deployment_instance_id = deployment_instance_id
|
||||
self.children: List[Dict[str, Any]] = []
|
||||
self.start_time: Optional[datetime] = None
|
||||
self.end_time: Optional[datetime] = None
|
||||
self.error: Optional[str] = None
|
||||
self.tool_calls = tool_calls
|
||||
self.flow_step = flow_step
|
||||
self.status: str = "running"
|
||||
|
||||
# Add trace to task's trace collection if task_id is present
|
||||
if task_id:
|
||||
self._add_to_task_traces()
|
||||
|
||||
def _add_to_task_traces(self) -> None:
|
||||
"""Add this trace to the task's trace collection."""
|
||||
if not hasattr(UnifiedTraceController, "_task_traces"):
|
||||
UnifiedTraceController._task_traces = {}
|
||||
|
||||
if self.task_id is None:
|
||||
return
|
||||
|
||||
if self.task_id not in UnifiedTraceController._task_traces:
|
||||
UnifiedTraceController._task_traces[self.task_id] = []
|
||||
|
||||
UnifiedTraceController._task_traces[self.task_id].append(self)
|
||||
|
||||
@classmethod
|
||||
def get_task_traces(cls, task_id: str) -> List["UnifiedTraceController"]:
|
||||
"""Get all traces for a specific task.
|
||||
|
||||
Args:
|
||||
task_id: The ID of the task to get traces for
|
||||
|
||||
Returns:
|
||||
List of traces associated with the task
|
||||
"""
|
||||
return cls._task_traces.get(task_id, [])
|
||||
|
||||
@classmethod
|
||||
def clear_task_traces(cls, task_id: str) -> None:
|
||||
"""Clear traces for a specific task.
|
||||
|
||||
Args:
|
||||
task_id: The ID of the task to clear traces for
|
||||
"""
|
||||
if hasattr(cls, "_task_traces") and task_id in cls._task_traces:
|
||||
del cls._task_traces[task_id]
|
||||
|
||||
def _get_current_trace(self) -> "UnifiedTraceController":
|
||||
return TraceContext.get_current()
|
||||
|
||||
def start_trace(self) -> "UnifiedTraceController":
|
||||
"""Start the trace execution.
|
||||
|
||||
Returns:
|
||||
UnifiedTraceController: Self for method chaining.
|
||||
"""
|
||||
self.start_time = datetime.now(UTC)
|
||||
return self
|
||||
|
||||
def end_trace(self, result: Any = None, error: Optional[str] = None) -> None:
|
||||
"""End the trace execution and record results.
|
||||
|
||||
Args:
|
||||
result: Optional result from the trace execution.
|
||||
error: Optional error message if the trace failed.
|
||||
"""
|
||||
self.end_time = datetime.now(UTC)
|
||||
self.status = "error" if error else "completed"
|
||||
self.error = error
|
||||
self._record_trace(result)
|
||||
|
||||
def add_child_trace(self, child_trace: Dict[str, Any]) -> None:
|
||||
"""Add a child trace to this trace's execution history.
|
||||
|
||||
Args:
|
||||
child_trace: The child trace information to add.
|
||||
"""
|
||||
self.children.append(child_trace)
|
||||
|
||||
def to_crew_trace(self) -> CrewTrace:
|
||||
"""Convert to CrewTrace format for storage.
|
||||
|
||||
Returns:
|
||||
CrewTrace: The trace data in CrewTrace format.
|
||||
"""
|
||||
latency_ms = None
|
||||
|
||||
if self.tool_calls and hasattr(self.tool_calls[0], "start_time"):
|
||||
self.start_time = self.tool_calls[0].start_time
|
||||
|
||||
if self.start_time and self.end_time:
|
||||
latency_ms = int((self.end_time - self.start_time).total_seconds() * 1000)
|
||||
|
||||
request = None
|
||||
response = None
|
||||
flow_step_obj = None
|
||||
|
||||
if self.trace_type in [TraceType.LLM_CALL, TraceType.TOOL_CALL]:
|
||||
request = LLMRequest(
|
||||
model=self.context.get("model", "unknown"),
|
||||
messages=self.context.get("messages", []),
|
||||
temperature=self.context.get("temperature"),
|
||||
max_tokens=self.context.get("max_tokens"),
|
||||
stop_sequences=self.context.get("stop_sequences"),
|
||||
)
|
||||
if "response" in self.context:
|
||||
response = LLMResponse(
|
||||
content=self.context["response"].get("content", ""),
|
||||
finish_reason=self.context["response"].get("finish_reason"),
|
||||
)
|
||||
|
||||
elif self.trace_type == TraceType.FLOW_STEP:
|
||||
flow_step_obj = FlowStepIO(
|
||||
function_name=self.flow_step.get("function_name", "unknown"),
|
||||
inputs=self.flow_step.get("inputs", {}),
|
||||
outputs={"result": self.context.get("response")},
|
||||
metadata=self.flow_step.get("metadata", {}),
|
||||
)
|
||||
|
||||
return CrewTrace(
|
||||
deployment_instance_id=self.deployment_instance_id,
|
||||
trace_id=self.trace_id,
|
||||
task_id=self.task_id,
|
||||
run_id=self.run_id,
|
||||
agent_role=self.agent_role,
|
||||
task_name=self.task_name,
|
||||
task_description=self.task_description,
|
||||
trace_type=self.trace_type.value,
|
||||
crew_type=self.crew_type.value,
|
||||
run_type=self.run_type.value,
|
||||
start_time=self.start_time,
|
||||
end_time=self.end_time,
|
||||
latency_ms=latency_ms,
|
||||
request=request,
|
||||
response=response,
|
||||
flow_step=flow_step_obj,
|
||||
tool_calls=self.tool_calls,
|
||||
tokens_used=self.context.get("tokens_used"),
|
||||
prompt_tokens=self.context.get("prompt_tokens"),
|
||||
completion_tokens=self.context.get("completion_tokens"),
|
||||
status=self.status,
|
||||
error=self.error,
|
||||
)
|
||||
|
||||
def _record_trace(self, result: Any = None) -> None:
|
||||
"""Record the trace.
|
||||
|
||||
This method is called when a trace is completed. It ensures the trace
|
||||
is properly recorded and associated with its task if applicable.
|
||||
|
||||
Args:
|
||||
result: Optional result to include in the trace
|
||||
"""
|
||||
if result:
|
||||
self.context["response"] = result
|
||||
|
||||
# Add to task traces if this trace belongs to a task
|
||||
if self.task_id:
|
||||
self._add_to_task_traces()
|
||||
|
||||
|
||||
def should_trace() -> bool:
|
||||
"""Check if tracing is enabled via environment variable."""
|
||||
return os.getenv("CREWAI_ENABLE_TRACING", "false").lower() == "true"
|
||||
|
||||
|
||||
# Crew main trace
|
||||
def init_crew_main_trace(func: Callable[..., Any]) -> Callable[..., Any]:
|
||||
"""Decorator to initialize and track the main crew execution trace.
|
||||
|
||||
This decorator sets up the trace context for the main crew execution,
|
||||
handling both synchronous and asynchronous crew operations.
|
||||
|
||||
Args:
|
||||
func: The crew function to be traced.
|
||||
|
||||
Returns:
|
||||
Wrapped function that creates and manages the main crew trace context.
|
||||
"""
|
||||
|
||||
@wraps(func)
|
||||
def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
if not should_trace():
|
||||
return func(self, *args, **kwargs)
|
||||
|
||||
trace = build_crew_main_trace(self)
|
||||
with TraceContext.set_current(trace):
|
||||
try:
|
||||
return func(self, *args, **kwargs)
|
||||
except Exception as e:
|
||||
trace.end_trace(error=str(e))
|
||||
raise
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def build_crew_main_trace(self: Any) -> "UnifiedTraceController":
|
||||
"""Build the main trace controller for a crew execution.
|
||||
|
||||
This function creates a trace controller configured for the main crew execution,
|
||||
handling different run types (kickoff, test, train) and maintaining context.
|
||||
|
||||
Args:
|
||||
self: The crew instance.
|
||||
|
||||
Returns:
|
||||
UnifiedTraceController: The configured trace controller for the crew.
|
||||
"""
|
||||
run_type = RunType.KICKOFF
|
||||
if hasattr(self, "_test") and self._test:
|
||||
run_type = RunType.TEST
|
||||
elif hasattr(self, "_train") and self._train:
|
||||
run_type = RunType.TRAIN
|
||||
|
||||
current_trace = TraceContext.get_current()
|
||||
|
||||
trace = UnifiedTraceController(
|
||||
trace_type=TraceType.LLM_CALL,
|
||||
run_type=run_type,
|
||||
crew_type=current_trace.crew_type if current_trace else CrewType.CREW,
|
||||
run_id=current_trace.run_id if current_trace else str(self.id),
|
||||
parent_trace_id=current_trace.trace_id if current_trace else None,
|
||||
)
|
||||
return trace
|
||||
|
||||
|
||||
# Flow main trace
|
||||
def init_flow_main_trace(
|
||||
func: Callable[..., Awaitable[Any]],
|
||||
) -> Callable[..., Awaitable[Any]]:
|
||||
"""Decorator to initialize and track the main flow execution trace.
|
||||
|
||||
Args:
|
||||
func: The async flow function to be traced.
|
||||
|
||||
Returns:
|
||||
Wrapped async function that creates and manages the main flow trace context.
|
||||
"""
|
||||
|
||||
@wraps(func)
|
||||
async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
if not should_trace():
|
||||
return await func(self, *args, **kwargs)
|
||||
|
||||
trace = build_flow_main_trace(self, *args, **kwargs)
|
||||
with TraceContext.set_current(trace):
|
||||
try:
|
||||
return await func(self, *args, **kwargs)
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def build_flow_main_trace(
|
||||
self: Any, *args: Any, **kwargs: Any
|
||||
) -> "UnifiedTraceController":
|
||||
"""Build the main trace controller for a flow execution.
|
||||
|
||||
Args:
|
||||
self: The flow instance.
|
||||
*args: Variable positional arguments.
|
||||
**kwargs: Variable keyword arguments.
|
||||
|
||||
Returns:
|
||||
UnifiedTraceController: The configured trace controller for the flow.
|
||||
"""
|
||||
current_trace = TraceContext.get_current()
|
||||
trace = UnifiedTraceController(
|
||||
trace_type=TraceType.FLOW_STEP,
|
||||
run_id=current_trace.run_id if current_trace else str(self.flow_id),
|
||||
parent_trace_id=current_trace.trace_id if current_trace else None,
|
||||
crew_type=CrewType.FLOW,
|
||||
run_type=RunType.KICKOFF,
|
||||
context={
|
||||
"crew_name": self.__class__.__name__,
|
||||
"inputs": kwargs.get("inputs", {}),
|
||||
"agents": [],
|
||||
"tasks": [],
|
||||
},
|
||||
)
|
||||
return trace
|
||||
|
||||
|
||||
# Flow step trace
|
||||
def trace_flow_step(
|
||||
func: Callable[..., Awaitable[Any]],
|
||||
) -> Callable[..., Awaitable[Any]]:
|
||||
"""Decorator to trace individual flow step executions.
|
||||
|
||||
Args:
|
||||
func: The async flow step function to be traced.
|
||||
|
||||
Returns:
|
||||
Wrapped async function that creates and manages the flow step trace context.
|
||||
"""
|
||||
|
||||
@wraps(func)
|
||||
async def wrapper(
|
||||
self: Any,
|
||||
method_name: str,
|
||||
method: Callable[..., Any],
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
if not should_trace():
|
||||
return await func(self, method_name, method, *args, **kwargs)
|
||||
|
||||
trace = build_flow_step_trace(self, method_name, method, *args, **kwargs)
|
||||
with TraceContext.set_current(trace):
|
||||
trace.start_trace()
|
||||
try:
|
||||
result = await func(self, method_name, method, *args, **kwargs)
|
||||
trace.end_trace(result=result)
|
||||
return result
|
||||
except Exception as e:
|
||||
trace.end_trace(error=str(e))
|
||||
raise
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def build_flow_step_trace(
|
||||
self: Any, method_name: str, method: Callable[..., Any], *args: Any, **kwargs: Any
|
||||
) -> "UnifiedTraceController":
|
||||
"""Build a trace controller for an individual flow step.
|
||||
|
||||
Args:
|
||||
self: The flow instance.
|
||||
method_name: Name of the method being executed.
|
||||
method: The actual method being executed.
|
||||
*args: Variable positional arguments.
|
||||
**kwargs: Variable keyword arguments.
|
||||
|
||||
Returns:
|
||||
UnifiedTraceController: The configured trace controller for the flow step.
|
||||
"""
|
||||
current_trace = TraceContext.get_current()
|
||||
|
||||
# Get method signature
|
||||
sig = inspect.signature(method)
|
||||
params = list(sig.parameters.values())
|
||||
|
||||
# Create inputs dictionary mapping parameter names to values
|
||||
method_params = [p for p in params if p.name != "self"]
|
||||
inputs: Dict[str, Any] = {}
|
||||
|
||||
# Map positional args to their parameter names
|
||||
for i, param in enumerate(method_params):
|
||||
if i < len(args):
|
||||
inputs[param.name] = args[i]
|
||||
|
||||
# Add keyword arguments
|
||||
inputs.update(kwargs)
|
||||
|
||||
trace = UnifiedTraceController(
|
||||
trace_type=TraceType.FLOW_STEP,
|
||||
run_type=current_trace.run_type if current_trace else RunType.KICKOFF,
|
||||
crew_type=current_trace.crew_type if current_trace else CrewType.FLOW,
|
||||
run_id=current_trace.run_id if current_trace else str(self.flow_id),
|
||||
parent_trace_id=current_trace.trace_id if current_trace else None,
|
||||
flow_step={
|
||||
"function_name": method_name,
|
||||
"inputs": inputs,
|
||||
"metadata": {
|
||||
"crew_name": self.__class__.__name__,
|
||||
},
|
||||
},
|
||||
)
|
||||
return trace
|
||||
|
||||
|
||||
# LLM trace
|
||||
def trace_llm_call(func: Callable[..., Any]) -> Callable[..., Any]:
|
||||
"""Decorator to trace LLM calls.
|
||||
|
||||
Args:
|
||||
func: The function to trace.
|
||||
|
||||
Returns:
|
||||
Wrapped function that creates and manages the LLM call trace context.
|
||||
"""
|
||||
|
||||
@wraps(func)
|
||||
def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
if not should_trace():
|
||||
return func(self, *args, **kwargs)
|
||||
|
||||
trace = build_llm_trace(self, *args, **kwargs)
|
||||
with TraceContext.set_current(trace):
|
||||
trace.start_trace()
|
||||
try:
|
||||
response = func(self, *args, **kwargs)
|
||||
# Extract relevant data from response
|
||||
trace_response = {
|
||||
"content": response["choices"][0]["message"]["content"],
|
||||
"finish_reason": response["choices"][0].get("finish_reason"),
|
||||
}
|
||||
|
||||
# Add usage metrics to context
|
||||
if "usage" in response:
|
||||
trace.context["tokens_used"] = response["usage"].get(
|
||||
"total_tokens", 0
|
||||
)
|
||||
trace.context["prompt_tokens"] = response["usage"].get(
|
||||
"prompt_tokens", 0
|
||||
)
|
||||
trace.context["completion_tokens"] = response["usage"].get(
|
||||
"completion_tokens", 0
|
||||
)
|
||||
|
||||
trace.end_trace(trace_response)
|
||||
return response
|
||||
except Exception as e:
|
||||
trace.end_trace(error=str(e))
|
||||
raise
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def build_llm_trace(
|
||||
self: Any, params: Dict[str, Any], *args: Any, **kwargs: Any
|
||||
) -> Any:
|
||||
"""Build a trace controller for an LLM call.
|
||||
|
||||
Args:
|
||||
self: The LLM instance.
|
||||
params: The parameters for the LLM call.
|
||||
*args: Variable positional arguments.
|
||||
**kwargs: Variable keyword arguments.
|
||||
|
||||
Returns:
|
||||
UnifiedTraceController: The configured trace controller for the LLM call.
|
||||
"""
|
||||
current_trace = TraceContext.get_current()
|
||||
agent, task = self._get_execution_context()
|
||||
|
||||
# Get new messages and tool results
|
||||
new_messages = self._get_new_messages(params.get("messages", []))
|
||||
new_tool_results = self._get_new_tool_results(agent)
|
||||
|
||||
# Create trace context
|
||||
trace = UnifiedTraceController(
|
||||
trace_type=TraceType.TOOL_CALL if new_tool_results else TraceType.LLM_CALL,
|
||||
crew_type=current_trace.crew_type if current_trace else CrewType.CREW,
|
||||
run_type=current_trace.run_type if current_trace else RunType.KICKOFF,
|
||||
run_id=current_trace.run_id if current_trace else str(uuid4()),
|
||||
parent_trace_id=current_trace.trace_id if current_trace else None,
|
||||
agent_role=agent.role if agent else "unknown",
|
||||
task_id=str(task.id) if task else None,
|
||||
task_name=task.name if task else None,
|
||||
task_description=task.description if task else None,
|
||||
model=self.model,
|
||||
messages=new_messages,
|
||||
temperature=self.temperature,
|
||||
max_tokens=self.max_tokens,
|
||||
stop_sequences=self.stop,
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
name=result["tool_name"],
|
||||
arguments=result["tool_args"],
|
||||
output=str(result["result"]),
|
||||
start_time=result.get("start_time", ""),
|
||||
end_time=datetime.now(UTC),
|
||||
)
|
||||
for result in new_tool_results
|
||||
],
|
||||
)
|
||||
return trace
|
||||
@@ -39,8 +39,8 @@
|
||||
"validation_error": "### Previous attempt failed validation: {guardrail_result_error}\n\n\n### Previous result:\n{task_output}\n\n\nTry again, making sure to address the validation error."
|
||||
},
|
||||
"tools": {
|
||||
"delegate_work": "Delegate a specific task to one of the following coworkers: {coworkers}\nThe input to this tool should be the coworker, the task you want them to do, and ALL necessary context to execute the task, they know nothing about the task, so share absolutely everything you know, don't reference things but instead explain them.",
|
||||
"ask_question": "Ask a specific question to one of the following coworkers: {coworkers}\nThe input to this tool should be the coworker, the question you have for them, and ALL necessary context to ask the question properly, they know nothing about the question, so share absolutely everything you know, don't reference things but instead explain them.",
|
||||
"delegate_work": "Delegate a specific task to one of the following coworkers: {coworkers}\nThe input to this tool should be the coworker, the task you want them to do, and ALL necessary context to execute the task, they know nothing about the task, so share absolute everything you know, don't reference things but instead explain them.",
|
||||
"ask_question": "Ask a specific question to one of the following coworkers: {coworkers}\nThe input to this tool should be the coworker, the question you have for them, and ALL necessary context to ask the question properly, they know nothing about the question, so share absolute everything you know, don't reference things but instead explain them.",
|
||||
"add_image": {
|
||||
"name": "Add image to content",
|
||||
"description": "See image to understand its content, you can optionally ask a question about the image",
|
||||
|
||||
@@ -44,7 +44,6 @@ def create_llm(
|
||||
# Extract attributes with explicit types
|
||||
model = (
|
||||
getattr(llm_value, "model_name", None)
|
||||
or getattr(llm_value, "model", None)
|
||||
or getattr(llm_value, "deployment_name", None)
|
||||
or str(llm_value)
|
||||
)
|
||||
|
||||
12
src/crewai/utilities/protocols.py
Normal file
12
src/crewai/utilities/protocols.py
Normal file
@@ -0,0 +1,12 @@
|
||||
from typing import Any, Protocol, runtime_checkable
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class AgentExecutorProtocol(Protocol):
|
||||
"""Protocol defining the expected interface for an agent executor."""
|
||||
|
||||
@property
|
||||
def agent(self) -> Any: ...
|
||||
|
||||
@property
|
||||
def task(self) -> Any: ...
|
||||
@@ -30,14 +30,8 @@ class TokenCalcHandler(CustomLogger):
|
||||
if hasattr(usage, "prompt_tokens"):
|
||||
self.token_cost_process.sum_prompt_tokens(usage.prompt_tokens)
|
||||
if hasattr(usage, "completion_tokens"):
|
||||
self.token_cost_process.sum_completion_tokens(
|
||||
usage.completion_tokens
|
||||
)
|
||||
if (
|
||||
hasattr(usage, "prompt_tokens_details")
|
||||
and usage.prompt_tokens_details
|
||||
and usage.prompt_tokens_details.cached_tokens
|
||||
):
|
||||
self.token_cost_process.sum_completion_tokens(usage.completion_tokens)
|
||||
if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
|
||||
self.token_cost_process.sum_cached_prompt_tokens(
|
||||
usage.prompt_tokens_details.cached_tokens
|
||||
)
|
||||
|
||||
@@ -915,6 +915,8 @@ def test_tool_result_as_answer_is_the_final_answer_for_the_agent():
|
||||
|
||||
@pytest.mark.vcr(filter_headers=["authorization"])
|
||||
def test_tool_usage_information_is_appended_to_agent():
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from crewai.tools import BaseTool
|
||||
|
||||
class MyCustomTool(BaseTool):
|
||||
@@ -924,30 +926,36 @@ def test_tool_usage_information_is_appended_to_agent():
|
||||
def _run(self) -> str:
|
||||
return "Howdy!"
|
||||
|
||||
agent1 = Agent(
|
||||
role="Friendly Neighbor",
|
||||
goal="Make everyone feel welcome",
|
||||
backstory="You are the friendly neighbor",
|
||||
tools=[MyCustomTool(result_as_answer=True)],
|
||||
)
|
||||
fixed_datetime = datetime(2025, 2, 10, 12, 0, 0, tzinfo=UTC)
|
||||
with patch("datetime.datetime") as mock_datetime:
|
||||
mock_datetime.now.return_value = fixed_datetime
|
||||
mock_datetime.side_effect = lambda *args, **kw: datetime(*args, **kw)
|
||||
|
||||
greeting = Task(
|
||||
description="Say an appropriate greeting.",
|
||||
expected_output="The greeting.",
|
||||
agent=agent1,
|
||||
)
|
||||
tasks = [greeting]
|
||||
crew = Crew(agents=[agent1], tasks=tasks)
|
||||
agent1 = Agent(
|
||||
role="Friendly Neighbor",
|
||||
goal="Make everyone feel welcome",
|
||||
backstory="You are the friendly neighbor",
|
||||
tools=[MyCustomTool(result_as_answer=True)],
|
||||
)
|
||||
|
||||
crew.kickoff()
|
||||
assert agent1.tools_results == [
|
||||
{
|
||||
"result": "Howdy!",
|
||||
"tool_name": "Decide Greetings",
|
||||
"tool_args": {},
|
||||
"result_as_answer": True,
|
||||
}
|
||||
]
|
||||
greeting = Task(
|
||||
description="Say an appropriate greeting.",
|
||||
expected_output="The greeting.",
|
||||
agent=agent1,
|
||||
)
|
||||
tasks = [greeting]
|
||||
crew = Crew(agents=[agent1], tasks=tasks)
|
||||
|
||||
crew.kickoff()
|
||||
assert agent1.tools_results == [
|
||||
{
|
||||
"result": "Howdy!",
|
||||
"tool_name": "Decide Greetings",
|
||||
"tool_args": {},
|
||||
"result_as_answer": True,
|
||||
"start_time": fixed_datetime,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_agent_definition_based_on_dict():
|
||||
|
||||
@@ -654,104 +654,3 @@ def test_flow_plotting():
|
||||
assert isinstance(received_events[0], FlowPlotEvent)
|
||||
assert received_events[0].flow_name == "StatelessFlow"
|
||||
assert isinstance(received_events[0].timestamp, datetime)
|
||||
|
||||
|
||||
def test_multiple_routers_from_same_trigger():
|
||||
"""Test that multiple routers triggered by the same method all activate their listeners."""
|
||||
execution_order = []
|
||||
|
||||
class MultiRouterFlow(Flow):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# Set diagnosed conditions to trigger all routers
|
||||
self.state["diagnosed_conditions"] = "DHA" # Contains D, H, and A
|
||||
|
||||
@start()
|
||||
def scan_medical(self):
|
||||
execution_order.append("scan_medical")
|
||||
return "scan_complete"
|
||||
|
||||
@router(scan_medical)
|
||||
def diagnose_conditions(self):
|
||||
execution_order.append("diagnose_conditions")
|
||||
return "diagnosis_complete"
|
||||
|
||||
@router(diagnose_conditions)
|
||||
def diabetes_router(self):
|
||||
execution_order.append("diabetes_router")
|
||||
if "D" in self.state["diagnosed_conditions"]:
|
||||
return "diabetes"
|
||||
return None
|
||||
|
||||
@listen("diabetes")
|
||||
def diabetes_analysis(self):
|
||||
execution_order.append("diabetes_analysis")
|
||||
return "diabetes_analysis_complete"
|
||||
|
||||
@router(diagnose_conditions)
|
||||
def hypertension_router(self):
|
||||
execution_order.append("hypertension_router")
|
||||
if "H" in self.state["diagnosed_conditions"]:
|
||||
return "hypertension"
|
||||
return None
|
||||
|
||||
@listen("hypertension")
|
||||
def hypertension_analysis(self):
|
||||
execution_order.append("hypertension_analysis")
|
||||
return "hypertension_analysis_complete"
|
||||
|
||||
@router(diagnose_conditions)
|
||||
def anemia_router(self):
|
||||
execution_order.append("anemia_router")
|
||||
if "A" in self.state["diagnosed_conditions"]:
|
||||
return "anemia"
|
||||
return None
|
||||
|
||||
@listen("anemia")
|
||||
def anemia_analysis(self):
|
||||
execution_order.append("anemia_analysis")
|
||||
return "anemia_analysis_complete"
|
||||
|
||||
flow = MultiRouterFlow()
|
||||
flow.kickoff()
|
||||
|
||||
# Verify all methods were called
|
||||
assert "scan_medical" in execution_order
|
||||
assert "diagnose_conditions" in execution_order
|
||||
|
||||
# Verify all routers were called
|
||||
assert "diabetes_router" in execution_order
|
||||
assert "hypertension_router" in execution_order
|
||||
assert "anemia_router" in execution_order
|
||||
|
||||
# Verify all listeners were called - this is the key test for the fix
|
||||
assert "diabetes_analysis" in execution_order
|
||||
assert "hypertension_analysis" in execution_order
|
||||
assert "anemia_analysis" in execution_order
|
||||
|
||||
# Verify execution order constraints
|
||||
assert execution_order.index("diagnose_conditions") > execution_order.index(
|
||||
"scan_medical"
|
||||
)
|
||||
|
||||
# All routers should execute after diagnose_conditions
|
||||
assert execution_order.index("diabetes_router") > execution_order.index(
|
||||
"diagnose_conditions"
|
||||
)
|
||||
assert execution_order.index("hypertension_router") > execution_order.index(
|
||||
"diagnose_conditions"
|
||||
)
|
||||
assert execution_order.index("anemia_router") > execution_order.index(
|
||||
"diagnose_conditions"
|
||||
)
|
||||
|
||||
# All analyses should execute after their respective routers
|
||||
assert execution_order.index("diabetes_analysis") > execution_order.index(
|
||||
"diabetes_router"
|
||||
)
|
||||
assert execution_order.index("hypertension_analysis") > execution_order.index(
|
||||
"hypertension_router"
|
||||
)
|
||||
assert execution_order.index("anemia_analysis") > execution_order.index(
|
||||
"anemia_router"
|
||||
)
|
||||
|
||||
@@ -6,7 +6,7 @@ import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from crewai.agents.agent_builder.utilities.base_token_process import TokenProcess
|
||||
from crewai.llm import CONTEXT_WINDOW_USAGE_RATIO, LLM
|
||||
from crewai.llm import LLM
|
||||
from crewai.utilities.events import crewai_event_bus
|
||||
from crewai.utilities.events.tool_usage_events import ToolExecutionErrorEvent
|
||||
from crewai.utilities.token_counter_callback import TokenCalcHandler
|
||||
@@ -285,23 +285,6 @@ def test_o3_mini_reasoning_effort_medium():
|
||||
assert isinstance(result, str)
|
||||
assert "Paris" in result
|
||||
|
||||
def test_context_window_validation():
|
||||
"""Test that context window validation works correctly."""
|
||||
# Test valid window size
|
||||
llm = LLM(model="o3-mini")
|
||||
assert llm.get_context_window_size() == int(200000 * CONTEXT_WINDOW_USAGE_RATIO)
|
||||
|
||||
# Test invalid window size
|
||||
with pytest.raises(ValueError) as excinfo:
|
||||
with patch.dict(
|
||||
"crewai.llm.LLM_CONTEXT_WINDOW_SIZES",
|
||||
{"test-model": 500}, # Below minimum
|
||||
clear=True,
|
||||
):
|
||||
llm = LLM(model="test-model")
|
||||
llm.get_context_window_size()
|
||||
assert "must be between 1024 and 2097152" in str(excinfo.value)
|
||||
|
||||
|
||||
@pytest.mark.vcr(filter_headers=["authorization"])
|
||||
@pytest.fixture
|
||||
|
||||
@@ -1,130 +0,0 @@
|
||||
import time
|
||||
from datetime import datetime, timedelta
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from crewai.agent import Agent
|
||||
from crewai.crew import Crew
|
||||
from crewai.memory.short_term.short_term_memory import ShortTermMemory
|
||||
from crewai.memory.short_term.short_term_memory_item import ShortTermMemoryItem
|
||||
from crewai.memory.storage.rag_storage import RAGStorage
|
||||
from crewai.task import Task
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def short_term_memory():
|
||||
"""Fixture to create a ShortTermMemory instance"""
|
||||
agent = Agent(
|
||||
role="Tutor",
|
||||
goal="Teach programming concepts",
|
||||
backstory="You are a programming tutor helping students learn.",
|
||||
tools=[],
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
task = Task(
|
||||
description="Explain programming concepts to students.",
|
||||
expected_output="Clear explanations of programming concepts.",
|
||||
agent=agent,
|
||||
)
|
||||
return ShortTermMemory(crew=Crew(agents=[agent], tasks=[task]))
|
||||
|
||||
|
||||
def test_memory_prioritizes_recent_topic(short_term_memory):
|
||||
"""Test that memory retrieval prioritizes the most recent topic in a conversation."""
|
||||
# First topic: Python variables
|
||||
topic1_data = "Variables in Python are dynamically typed. You can assign any value to a variable without declaring its type."
|
||||
topic1_timestamp = datetime.now() - timedelta(minutes=10) # Older memory
|
||||
|
||||
# Second topic: Python abstract classes
|
||||
topic2_data = "Abstract classes in Python are created using the ABC module. They cannot be instantiated and are used as a blueprint for other classes."
|
||||
topic2_timestamp = datetime.now() # More recent memory
|
||||
|
||||
# Mock search results to simulate what would be returned by RAGStorage
|
||||
mock_results = [
|
||||
{
|
||||
"id": "2",
|
||||
"metadata": {
|
||||
"agent": "Tutor",
|
||||
"topic": "python_abstract_classes",
|
||||
"timestamp": topic2_timestamp.isoformat()
|
||||
},
|
||||
"context": topic2_data,
|
||||
"score": 0.85, # Higher score due to recency boost
|
||||
},
|
||||
{
|
||||
"id": "1",
|
||||
"metadata": {
|
||||
"agent": "Tutor",
|
||||
"topic": "python_variables",
|
||||
"timestamp": topic1_timestamp.isoformat()
|
||||
},
|
||||
"context": topic1_data,
|
||||
"score": 0.75, # Lower score due to being older
|
||||
}
|
||||
]
|
||||
|
||||
# Mock the search method to return our predefined results
|
||||
with patch.object(RAGStorage, 'search', return_value=mock_results):
|
||||
# Query that could match both topics but should prioritize the more recent one
|
||||
query = "Can you give me another example of that?"
|
||||
|
||||
# Search with recency consideration
|
||||
results = short_term_memory.search(query)
|
||||
|
||||
# Verify that the most recent topic (abstract classes) is prioritized
|
||||
assert len(results) > 0, "No search results returned"
|
||||
|
||||
# The first result should be about abstract classes (the more recent topic)
|
||||
assert "abstract classes" in results[0]["context"].lower(), "Recent topic (abstract classes) not prioritized"
|
||||
|
||||
# If there are multiple results, check if the older topic is also returned but with lower priority
|
||||
if len(results) > 1:
|
||||
assert "variables" in results[1]["context"].lower(), "Older topic should be second"
|
||||
|
||||
# Verify that the scores reflect the recency prioritization
|
||||
assert results[0]["score"] > results[1]["score"], "Recent topic should have higher score"
|
||||
|
||||
|
||||
def test_future_timestamp_validation():
|
||||
"""Test that ShortTermMemoryItem raises ValueError for future timestamps."""
|
||||
# Setup agent and task for memory
|
||||
agent = Agent(
|
||||
role="Tutor",
|
||||
goal="Teach programming concepts",
|
||||
backstory="You are a programming tutor helping students learn.",
|
||||
tools=[],
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
task = Task(
|
||||
description="Explain programming concepts to students.",
|
||||
expected_output="Clear explanations of programming concepts.",
|
||||
agent=agent,
|
||||
)
|
||||
|
||||
# Create a future timestamp
|
||||
future_timestamp = datetime.now() + timedelta(days=1)
|
||||
|
||||
# Test constructor validation
|
||||
with pytest.raises(ValueError, match="Timestamp cannot be in the future"):
|
||||
ShortTermMemoryItem(data="Test data", timestamp=future_timestamp)
|
||||
|
||||
# Test save method validation
|
||||
memory = ShortTermMemory(crew=Crew(agents=[agent], tasks=[task]))
|
||||
|
||||
# Create a memory item with a future timestamp
|
||||
future_data = "Test data with future timestamp"
|
||||
|
||||
# We need to pass the data directly to the save method
|
||||
# The save method will create a ShortTermMemoryItem internally
|
||||
# and then we'll modify its timestamp before it's saved
|
||||
|
||||
# Mock datetime.now to return a fixed time
|
||||
with patch('crewai.memory.short_term.short_term_memory_item.datetime') as mock_datetime:
|
||||
# Set up the mock to return our future timestamp when now() is called
|
||||
mock_datetime.now.return_value = future_timestamp
|
||||
|
||||
with pytest.raises(ValueError, match="Cannot save memory item with future timestamp"):
|
||||
memory.save(value=future_data)
|
||||
360
tests/traces/test_unified_trace_controller.py
Normal file
360
tests/traces/test_unified_trace_controller.py
Normal file
@@ -0,0 +1,360 @@
|
||||
import os
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import MagicMock, patch
|
||||
from uuid import UUID
|
||||
|
||||
import pytest
|
||||
|
||||
from crewai.traces.context import TraceContext
|
||||
from crewai.traces.enums import CrewType, RunType, TraceType
|
||||
from crewai.traces.models import (
|
||||
CrewTrace,
|
||||
FlowStepIO,
|
||||
LLMRequest,
|
||||
LLMResponse,
|
||||
)
|
||||
from crewai.traces.unified_trace_controller import (
|
||||
UnifiedTraceController,
|
||||
init_crew_main_trace,
|
||||
init_flow_main_trace,
|
||||
should_trace,
|
||||
trace_flow_step,
|
||||
trace_llm_call,
|
||||
)
|
||||
|
||||
|
||||
class TestUnifiedTraceController:
|
||||
@pytest.fixture
|
||||
def basic_trace_controller(self):
|
||||
return UnifiedTraceController(
|
||||
trace_type=TraceType.LLM_CALL,
|
||||
run_type=RunType.KICKOFF,
|
||||
crew_type=CrewType.CREW,
|
||||
run_id="test-run-id",
|
||||
agent_role="test-agent",
|
||||
task_name="test-task",
|
||||
task_description="test description",
|
||||
task_id="test-task-id",
|
||||
)
|
||||
|
||||
def test_initialization(self, basic_trace_controller):
|
||||
"""Test basic initialization of UnifiedTraceController"""
|
||||
assert basic_trace_controller.trace_type == TraceType.LLM_CALL
|
||||
assert basic_trace_controller.run_type == RunType.KICKOFF
|
||||
assert basic_trace_controller.crew_type == CrewType.CREW
|
||||
assert basic_trace_controller.run_id == "test-run-id"
|
||||
assert basic_trace_controller.agent_role == "test-agent"
|
||||
assert basic_trace_controller.task_name == "test-task"
|
||||
assert basic_trace_controller.task_description == "test description"
|
||||
assert basic_trace_controller.task_id == "test-task-id"
|
||||
assert basic_trace_controller.status == "running"
|
||||
assert isinstance(UUID(basic_trace_controller.trace_id), UUID)
|
||||
|
||||
def test_start_trace(self, basic_trace_controller):
|
||||
"""Test starting a trace"""
|
||||
result = basic_trace_controller.start_trace()
|
||||
assert result == basic_trace_controller
|
||||
assert basic_trace_controller.start_time is not None
|
||||
assert isinstance(basic_trace_controller.start_time, datetime)
|
||||
|
||||
def test_end_trace_success(self, basic_trace_controller):
|
||||
"""Test ending a trace successfully"""
|
||||
basic_trace_controller.start_trace()
|
||||
basic_trace_controller.end_trace(result={"test": "result"})
|
||||
|
||||
assert basic_trace_controller.end_time is not None
|
||||
assert basic_trace_controller.status == "completed"
|
||||
assert basic_trace_controller.error is None
|
||||
assert basic_trace_controller.context.get("response") == {"test": "result"}
|
||||
|
||||
def test_end_trace_with_error(self, basic_trace_controller):
|
||||
"""Test ending a trace with an error"""
|
||||
basic_trace_controller.start_trace()
|
||||
basic_trace_controller.end_trace(error="Test error occurred")
|
||||
|
||||
assert basic_trace_controller.end_time is not None
|
||||
assert basic_trace_controller.status == "error"
|
||||
assert basic_trace_controller.error == "Test error occurred"
|
||||
|
||||
def test_add_child_trace(self, basic_trace_controller):
|
||||
"""Test adding a child trace"""
|
||||
child_trace = {"id": "child-1", "type": "test"}
|
||||
basic_trace_controller.add_child_trace(child_trace)
|
||||
assert len(basic_trace_controller.children) == 1
|
||||
assert basic_trace_controller.children[0] == child_trace
|
||||
|
||||
def test_to_crew_trace_llm_call(self):
|
||||
"""Test converting to CrewTrace for LLM call"""
|
||||
test_messages = [{"role": "user", "content": "test"}]
|
||||
test_response = {
|
||||
"content": "test response",
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
|
||||
controller = UnifiedTraceController(
|
||||
trace_type=TraceType.LLM_CALL,
|
||||
run_type=RunType.KICKOFF,
|
||||
crew_type=CrewType.CREW,
|
||||
run_id="test-run-id",
|
||||
context={
|
||||
"messages": test_messages,
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 100,
|
||||
},
|
||||
)
|
||||
|
||||
# Set model and messages in the context
|
||||
controller.context["model"] = "gpt-4"
|
||||
controller.context["messages"] = test_messages
|
||||
|
||||
controller.start_trace()
|
||||
controller.end_trace(result=test_response)
|
||||
|
||||
crew_trace = controller.to_crew_trace()
|
||||
assert isinstance(crew_trace, CrewTrace)
|
||||
assert isinstance(crew_trace.request, LLMRequest)
|
||||
assert isinstance(crew_trace.response, LLMResponse)
|
||||
assert crew_trace.request.model == "gpt-4"
|
||||
assert crew_trace.request.messages == test_messages
|
||||
assert crew_trace.response.content == test_response["content"]
|
||||
assert crew_trace.response.finish_reason == test_response["finish_reason"]
|
||||
|
||||
def test_to_crew_trace_flow_step(self):
|
||||
"""Test converting to CrewTrace for flow step"""
|
||||
flow_step_data = {
|
||||
"function_name": "test_function",
|
||||
"inputs": {"param1": "value1"},
|
||||
"metadata": {"meta": "data"},
|
||||
}
|
||||
|
||||
controller = UnifiedTraceController(
|
||||
trace_type=TraceType.FLOW_STEP,
|
||||
run_type=RunType.KICKOFF,
|
||||
crew_type=CrewType.FLOW,
|
||||
run_id="test-run-id",
|
||||
flow_step=flow_step_data,
|
||||
)
|
||||
|
||||
controller.start_trace()
|
||||
controller.end_trace(result="test result")
|
||||
|
||||
crew_trace = controller.to_crew_trace()
|
||||
assert isinstance(crew_trace, CrewTrace)
|
||||
assert isinstance(crew_trace.flow_step, FlowStepIO)
|
||||
assert crew_trace.flow_step.function_name == "test_function"
|
||||
assert crew_trace.flow_step.inputs == {"param1": "value1"}
|
||||
assert crew_trace.flow_step.outputs == {"result": "test result"}
|
||||
|
||||
def test_should_trace(self):
|
||||
"""Test should_trace function"""
|
||||
with patch.dict(os.environ, {"CREWAI_ENABLE_TRACING": "true"}):
|
||||
assert should_trace() is True
|
||||
|
||||
with patch.dict(os.environ, {"CREWAI_ENABLE_TRACING": "false"}):
|
||||
assert should_trace() is False
|
||||
|
||||
with patch.dict(os.environ, clear=True):
|
||||
assert should_trace() is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_flow_step_decorator(self):
|
||||
"""Test trace_flow_step decorator"""
|
||||
|
||||
class TestFlow:
|
||||
flow_id = "test-flow-id"
|
||||
|
||||
@trace_flow_step
|
||||
async def test_method(self, method_name, method, *args, **kwargs):
|
||||
return "test result"
|
||||
|
||||
with patch.dict(os.environ, {"CREWAI_ENABLE_TRACING": "true"}):
|
||||
flow = TestFlow()
|
||||
result = await flow.test_method("test_method", lambda x: x, arg1="value1")
|
||||
assert result == "test result"
|
||||
|
||||
def test_trace_llm_call_decorator(self):
|
||||
"""Test trace_llm_call decorator"""
|
||||
|
||||
class TestLLM:
|
||||
model = "gpt-4"
|
||||
temperature = 0.7
|
||||
max_tokens = 100
|
||||
stop = None
|
||||
|
||||
def _get_execution_context(self):
|
||||
return MagicMock(), MagicMock()
|
||||
|
||||
def _get_new_messages(self, messages):
|
||||
return messages
|
||||
|
||||
def _get_new_tool_results(self, agent):
|
||||
return []
|
||||
|
||||
@trace_llm_call
|
||||
def test_method(self, params):
|
||||
return {
|
||||
"choices": [
|
||||
{
|
||||
"message": {"content": "test response"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"total_tokens": 50,
|
||||
"prompt_tokens": 20,
|
||||
"completion_tokens": 30,
|
||||
},
|
||||
}
|
||||
|
||||
with patch.dict(os.environ, {"CREWAI_ENABLE_TRACING": "true"}):
|
||||
llm = TestLLM()
|
||||
result = llm.test_method({"messages": []})
|
||||
assert result["choices"][0]["message"]["content"] == "test response"
|
||||
|
||||
def test_init_crew_main_trace_kickoff(self):
|
||||
"""Test init_crew_main_trace in kickoff mode"""
|
||||
trace_context = None
|
||||
|
||||
class TestCrew:
|
||||
id = "test-crew-id"
|
||||
_test = False
|
||||
_train = False
|
||||
|
||||
@init_crew_main_trace
|
||||
def test_method(self):
|
||||
nonlocal trace_context
|
||||
trace_context = TraceContext.get_current()
|
||||
return "test result"
|
||||
|
||||
with patch.dict(os.environ, {"CREWAI_ENABLE_TRACING": "true"}):
|
||||
crew = TestCrew()
|
||||
result = test_method(crew)
|
||||
assert result == "test result"
|
||||
assert trace_context is not None
|
||||
assert trace_context.trace_type == TraceType.LLM_CALL
|
||||
assert trace_context.run_type == RunType.KICKOFF
|
||||
assert trace_context.crew_type == CrewType.CREW
|
||||
assert trace_context.run_id == str(crew.id)
|
||||
|
||||
def test_init_crew_main_trace_test_mode(self):
|
||||
"""Test init_crew_main_trace in test mode"""
|
||||
trace_context = None
|
||||
|
||||
class TestCrew:
|
||||
id = "test-crew-id"
|
||||
_test = True
|
||||
_train = False
|
||||
|
||||
@init_crew_main_trace
|
||||
def test_method(self):
|
||||
nonlocal trace_context
|
||||
trace_context = TraceContext.get_current()
|
||||
return "test result"
|
||||
|
||||
with patch.dict(os.environ, {"CREWAI_ENABLE_TRACING": "true"}):
|
||||
crew = TestCrew()
|
||||
result = test_method(crew)
|
||||
assert result == "test result"
|
||||
assert trace_context is not None
|
||||
assert trace_context.run_type == RunType.TEST
|
||||
|
||||
def test_init_crew_main_trace_train_mode(self):
|
||||
"""Test init_crew_main_trace in train mode"""
|
||||
trace_context = None
|
||||
|
||||
class TestCrew:
|
||||
id = "test-crew-id"
|
||||
_test = False
|
||||
_train = True
|
||||
|
||||
@init_crew_main_trace
|
||||
def test_method(self):
|
||||
nonlocal trace_context
|
||||
trace_context = TraceContext.get_current()
|
||||
return "test result"
|
||||
|
||||
with patch.dict(os.environ, {"CREWAI_ENABLE_TRACING": "true"}):
|
||||
crew = TestCrew()
|
||||
result = test_method(crew)
|
||||
assert result == "test result"
|
||||
assert trace_context is not None
|
||||
assert trace_context.run_type == RunType.TRAIN
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_flow_main_trace(self):
|
||||
"""Test init_flow_main_trace decorator"""
|
||||
trace_context = None
|
||||
test_inputs = {"test": "input"}
|
||||
|
||||
class TestFlow:
|
||||
flow_id = "test-flow-id"
|
||||
|
||||
@init_flow_main_trace
|
||||
async def test_method(self, **kwargs):
|
||||
nonlocal trace_context
|
||||
trace_context = TraceContext.get_current()
|
||||
# Verify the context is set during execution
|
||||
assert trace_context.context["context"]["inputs"] == test_inputs
|
||||
return "test result"
|
||||
|
||||
with patch.dict(os.environ, {"CREWAI_ENABLE_TRACING": "true"}):
|
||||
flow = TestFlow()
|
||||
result = await flow.test_method(inputs=test_inputs)
|
||||
assert result == "test result"
|
||||
assert trace_context is not None
|
||||
assert trace_context.trace_type == TraceType.FLOW_STEP
|
||||
assert trace_context.crew_type == CrewType.FLOW
|
||||
assert trace_context.run_type == RunType.KICKOFF
|
||||
assert trace_context.run_id == str(flow.flow_id)
|
||||
assert trace_context.context["context"]["inputs"] == test_inputs
|
||||
|
||||
def test_trace_context_management(self):
|
||||
"""Test TraceContext management"""
|
||||
trace1 = UnifiedTraceController(
|
||||
trace_type=TraceType.LLM_CALL,
|
||||
run_type=RunType.KICKOFF,
|
||||
crew_type=CrewType.CREW,
|
||||
run_id="test-run-1",
|
||||
)
|
||||
|
||||
trace2 = UnifiedTraceController(
|
||||
trace_type=TraceType.FLOW_STEP,
|
||||
run_type=RunType.TEST,
|
||||
crew_type=CrewType.FLOW,
|
||||
run_id="test-run-2",
|
||||
)
|
||||
|
||||
# Test that context is initially empty
|
||||
assert TraceContext.get_current() is None
|
||||
|
||||
# Test setting and getting context
|
||||
with TraceContext.set_current(trace1):
|
||||
assert TraceContext.get_current() == trace1
|
||||
|
||||
# Test nested context
|
||||
with TraceContext.set_current(trace2):
|
||||
assert TraceContext.get_current() == trace2
|
||||
|
||||
# Test context restoration after nested block
|
||||
assert TraceContext.get_current() == trace1
|
||||
|
||||
# Test context cleanup after with block
|
||||
assert TraceContext.get_current() is None
|
||||
|
||||
def test_trace_context_error_handling(self):
|
||||
"""Test TraceContext error handling"""
|
||||
trace = UnifiedTraceController(
|
||||
trace_type=TraceType.LLM_CALL,
|
||||
run_type=RunType.KICKOFF,
|
||||
crew_type=CrewType.CREW,
|
||||
run_id="test-run",
|
||||
)
|
||||
|
||||
# Test that context is properly cleaned up even if an error occurs
|
||||
try:
|
||||
with TraceContext.set_current(trace):
|
||||
raise ValueError("Test error")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
assert TraceContext.get_current() is None
|
||||
@@ -606,7 +606,7 @@ def test_llm_emits_call_failed_event():
|
||||
received_events.append(event)
|
||||
|
||||
error_message = "Simulated LLM call failure"
|
||||
with patch("crewai.llm.litellm.completion", side_effect=Exception(error_message)):
|
||||
with patch.object(LLM, "_call_llm", side_effect=Exception(error_message)):
|
||||
llm = LLM(model="gpt-4o-mini")
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
llm.call("Hello, how are you?")
|
||||
|
||||
Reference in New Issue
Block a user