| from __future__ import annotations
|
|
|
| from typing import List, Optional, Any, Dict
|
| from pydantic import BaseModel, Field
|
|
|
| from ai_agent.generator.schema import ToolSelection, CandidateDoc
|
|
|
|
|
| class ToolRunLog(BaseModel):
|
| tool: str
|
| inputs: Dict[str, Any] = Field(default_factory=dict)
|
| error: Optional[str] = None
|
| timestamp: Optional[str] = None
|
|
|
|
|
| class UsageStats(BaseModel):
|
| """Token usage statistics from the agent."""
|
|
|
| total_tokens: int = 0
|
| input_tokens: int = 0
|
| output_tokens: int = 0
|
|
|
|
|
| class AgentToolSelection(ToolSelection):
|
| tool_calls: List[ToolRunLog] = Field(default_factory=list)
|
| usage: Optional[UsageStats] = None
|
|
|
| def to_legacy_dict(self) -> Dict[str, Any]:
|
|
|
| return {
|
| "conversation": self.conversation.model_dump(mode="python"),
|
| "choices": [c.model_dump(mode="python") for c in self.choices],
|
| "reason": self.reason,
|
| "explanation": self.explanation,
|
| "tool_calls": [c.model_dump(mode="python") for c in self.tool_calls],
|
| "usage": self.usage.model_dump(mode="python") if self.usage else None,
|
| }
|
|
|
|
|
| __all__ = [
|
| "AgentToolSelection",
|
| "ToolRunLog",
|
| "UsageStats",
|
| "CandidateDoc",
|
| ]
|
|
|
