155 lines
4.7 KiB
Python
155 lines
4.7 KiB
Python
# Copyright (c) 2025 Patrick Motsch
|
|
# All rights reserved.
|
|
"""Data models for the Agent service."""
|
|
|
|
from typing import List, Dict, Any, Optional
|
|
from enum import Enum
|
|
from pydantic import BaseModel, Field
|
|
from modules.shared.timeUtils import getUtcTimestamp
|
|
import uuid
|
|
|
|
|
|
class AgentStatusEnum(str, Enum):
|
|
RUNNING = "running"
|
|
COMPLETED = "completed"
|
|
MAX_ROUNDS_REACHED = "maxRoundsReached"
|
|
BUDGET_EXCEEDED = "budgetExceeded"
|
|
ERROR = "error"
|
|
STOPPED = "stopped"
|
|
|
|
|
|
class AgentEventTypeEnum(str, Enum):
|
|
MESSAGE = "message"
|
|
CHUNK = "chunk"
|
|
TOOL_CALL = "toolCall"
|
|
TOOL_RESULT = "toolResult"
|
|
AGENT_PROGRESS = "agentProgress"
|
|
AGENT_SUMMARY = "agentSummary"
|
|
FILE_CREATED = "fileCreated"
|
|
FILE_UPDATED = "fileUpdated"
|
|
FILE_EDIT_PROPOSAL = "fileEditProposal"
|
|
FILE_VERSION = "fileVersion"
|
|
FILE_EDIT_REJECTED = "fileEditRejected"
|
|
DATA_SOURCE_ACCESS = "dataSourceAccess"
|
|
VOICE_RESPONSE = "voiceResponse"
|
|
FINAL = "final"
|
|
ERROR = "error"
|
|
|
|
|
|
class ToolDefinition(BaseModel):
|
|
"""Schema for a tool available to the agent."""
|
|
name: str = Field(description="Unique tool name")
|
|
description: str = Field(description="What this tool does")
|
|
parameters: Dict[str, Any] = Field(
|
|
default_factory=dict,
|
|
description="JSON Schema for tool parameters"
|
|
)
|
|
readOnly: bool = Field(
|
|
default=False,
|
|
description="If True, tool can run in parallel with other readOnly tools"
|
|
)
|
|
featureType: Optional[str] = Field(
|
|
default=None,
|
|
description="Feature scope for this tool (None = available to all)"
|
|
)
|
|
toolSet: Optional[str] = Field(
|
|
default=None,
|
|
description="Tool-set scope (None = available to all sets, e.g. 'core', 'workspace')"
|
|
)
|
|
|
|
|
|
class ToolCallRequest(BaseModel):
|
|
"""A tool call requested by the AI model."""
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()))
|
|
name: str
|
|
args: Dict[str, Any] = Field(default_factory=dict)
|
|
|
|
|
|
class ToolResult(BaseModel):
|
|
"""Result from executing a tool."""
|
|
toolCallId: str
|
|
toolName: str
|
|
success: bool = True
|
|
data: str = ""
|
|
error: Optional[str] = None
|
|
durationMs: int = 0
|
|
sideEvents: Optional[List[Dict[str, Any]]] = None
|
|
|
|
|
|
class AgentEvent(BaseModel):
|
|
"""Event emitted during agent execution for SSE streaming."""
|
|
type: AgentEventTypeEnum
|
|
content: Optional[str] = None
|
|
data: Optional[Dict[str, Any]] = None
|
|
|
|
|
|
class AgentConfig(BaseModel):
|
|
"""Configuration for an agent run."""
|
|
maxRounds: int = Field(default=25, ge=1, le=100)
|
|
maxCostCHF: Optional[float] = Field(default=None, ge=0.0)
|
|
toolSet: str = Field(default="core")
|
|
initialToolboxes: List[str] = Field(default_factory=lambda: ["core"])
|
|
availableToolboxes: List[str] = Field(default_factory=list)
|
|
temperature: Optional[float] = Field(default=None, ge=0.0, le=2.0)
|
|
|
|
|
|
class AgentState(BaseModel):
|
|
"""Tracks state across an agent loop execution."""
|
|
workflowId: str
|
|
currentRound: int = 0
|
|
maxRounds: int = 25
|
|
totalAiCalls: int = 0
|
|
totalToolCalls: int = 0
|
|
totalCostCHF: float = 0.0
|
|
totalProcessingTime: float = 0.0
|
|
status: AgentStatusEnum = AgentStatusEnum.RUNNING
|
|
abortReason: Optional[str] = None
|
|
|
|
|
|
class ToolCallLog(BaseModel):
|
|
"""Log of a single tool call for observability."""
|
|
toolName: str
|
|
args: Dict[str, Any] = Field(default_factory=dict)
|
|
success: bool = True
|
|
durationMs: int = 0
|
|
error: Optional[str] = None
|
|
resultData: str = Field(default="", description="Short result summary for artifact tracking")
|
|
|
|
|
|
class AgentRoundLog(BaseModel):
|
|
"""Log of a single agent round for observability."""
|
|
roundNumber: int
|
|
aiModel: str = ""
|
|
inputTokens: int = 0
|
|
outputTokens: int = 0
|
|
costCHF: float = 0.0
|
|
toolCalls: List[ToolCallLog] = Field(default_factory=list)
|
|
durationMs: int = 0
|
|
|
|
|
|
class AgentTrace(BaseModel):
|
|
"""Full trace of an agent workflow for observability."""
|
|
workflowId: str
|
|
userId: str = ""
|
|
featureInstanceId: str = ""
|
|
startedAt: float = Field(default_factory=getUtcTimestamp)
|
|
completedAt: Optional[float] = None
|
|
status: AgentStatusEnum = AgentStatusEnum.RUNNING
|
|
totalRounds: int = 0
|
|
totalToolCalls: int = 0
|
|
totalCostCHF: float = 0.0
|
|
abortReason: Optional[str] = None
|
|
rounds: List[AgentRoundLog] = Field(default_factory=list)
|
|
|
|
|
|
class PendingFileEdit(BaseModel):
|
|
"""A proposed file edit awaiting user approval."""
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()))
|
|
fileId: str
|
|
fileName: str
|
|
mimeType: str = ""
|
|
oldContent: str = ""
|
|
newContent: str = ""
|
|
status: str = Field(default="pending", description="pending | accepted | rejected")
|
|
toolCallId: str = ""
|
|
workflowId: str = ""
|