You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
90 lines
4.0 KiB
90 lines
4.0 KiB
from typing import Any, Dict, List, Optional
|
|
from pydantic import BaseModel, Field
|
|
import time
|
|
import uuid
|
|
from swarms.utils.litellm_tokenizer import count_tokens
|
|
|
|
class Step(BaseModel):
|
|
"""
|
|
Represents a single step in an agent's task execution.
|
|
"""
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()))
|
|
name: str = Field(..., description="Name of the agent")
|
|
task: Optional[str] = Field(None, description="Task given to the agent at this step")
|
|
input: Optional[str] = Field(None, description="Input provided to the agent at this step")
|
|
output: Optional[str] = Field(None, description="Output generated by the agent at this step")
|
|
error: Optional[str] = Field(None, description="Error message if any error occurred during the step")
|
|
start_time: str = Field(default_factory=lambda: time.strftime("%Y-%m-%d %H:%M:%S"))
|
|
end_time: Optional[str] = Field(None, description="End time of the step")
|
|
runtime: Optional[float] = Field(None, description="Runtime of the step in seconds")
|
|
tokens_used: Optional[int] = Field(None, description="Number of tokens used in this step")
|
|
cost: Optional[float] = Field(None, description="Cost of the step")
|
|
metadata: Optional[Dict[str, Any]] = Field(
|
|
None, description="Additional metadata about the step"
|
|
)
|
|
|
|
def calculate_tokens(self, model: str = "gpt-4o") -> int:
|
|
"""Calculate total tokens used in this step"""
|
|
total = 0
|
|
if self.input:
|
|
total += count_tokens(self.input, model)
|
|
if self.output:
|
|
total += count_tokens(self.output, model)
|
|
self.tokens_used = total
|
|
return total
|
|
|
|
class AgentTaskOutput(BaseModel):
|
|
"""
|
|
Represents the output of an agent's execution.
|
|
"""
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()))
|
|
agent_name: str = Field(..., description="Name of the agent")
|
|
task: str = Field(..., description="The task agent was asked to perform")
|
|
steps: List[Step] = Field(..., description="List of steps taken by the agent")
|
|
start_time: str = Field(default_factory=lambda: time.strftime("%Y-%m-%d %H:%M:%S"))
|
|
end_time: Optional[str] = Field(None, description="End time of the agent's execution")
|
|
total_tokens: Optional[int] = Field(None, description="Total tokens used by the agent")
|
|
cost: Optional[float] = Field(None, description="Total cost of the agent execution")
|
|
# Add any other fields from ManySteps that are relevant, like full_history
|
|
|
|
def calculate_total_tokens(self, model: str = "gpt-4o") -> int:
|
|
"""Calculate total tokens across all steps"""
|
|
total = 0
|
|
for step in self.steps:
|
|
total += step.calculate_tokens(model)
|
|
self.total_tokens = total
|
|
return total
|
|
|
|
class OutputSchema(BaseModel):
|
|
"""
|
|
Unified output schema for all swarm types.
|
|
"""
|
|
swarm_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
|
|
swarm_type: str = Field(..., description="Type of the swarm")
|
|
task: str = Field(..., description="The task given to the swarm")
|
|
agent_outputs: List[AgentTaskOutput] = Field(..., description="List of agent outputs")
|
|
timestamp: str = Field(default_factory=lambda: time.strftime("%Y-%m-%d %H:%M:%S"))
|
|
swarm_specific_output: Optional[Dict] = Field(None, description="Additional data specific to the swarm type")
|
|
|
|
class SwarmOutputFormatter:
|
|
"""
|
|
Formatter class to transform raw swarm output into the OutputSchema format.
|
|
"""
|
|
|
|
@staticmethod
|
|
def format_output(
|
|
swarm_id: str,
|
|
swarm_type: str,
|
|
task: str,
|
|
agent_outputs: List[AgentTaskOutput],
|
|
swarm_specific_output: Optional[Dict] = None,
|
|
) -> str:
|
|
"""Formats the output into a standardized JSON string."""
|
|
output = OutputSchema(
|
|
swarm_id=swarm_id,
|
|
swarm_type=swarm_type,
|
|
task=task,
|
|
agent_outputs=agent_outputs,
|
|
swarm_specific_output=swarm_specific_output,
|
|
)
|
|
return output.model_dump_json(indent=4) |