qzl
178 lines
4.6 KiB
Python
178 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
from enum import Enum
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
|
|
|
from schemas.agent.ui_hints import UiHintsPayload
|
|
|
|
|
|
class TaskType(str, Enum):
|
|
KNOWLEDGE = "knowledge"
|
|
RECOMMENDATION = "recommendation"
|
|
PLANNING = "planning"
|
|
SCHEDULING = "scheduling"
|
|
REMINDER_MANAGEMENT = "reminder_management"
|
|
TODO_MANAGEMENT = "todo_management"
|
|
COMMUNICATION_DRAFTING = "communication_drafting"
|
|
INFORMATION_ORGANIZATION = "information_organization"
|
|
STATUS_TRACKING = "status_tracking"
|
|
TRANSACTION_ASSIST = "transaction_assist"
|
|
ACTION_EXECUTION = "action_execution"
|
|
TROUBLESHOOTING = "troubleshooting"
|
|
UNKNOWN = "unknown"
|
|
|
|
|
|
class ResultType(str, Enum):
|
|
DIRECT_ANSWER = "direct_answer"
|
|
OPTIONS_WITH_RECOMMENDATION = "options_with_recommendation"
|
|
ACTION_PLAN = "action_plan"
|
|
SCHEDULE_PROPOSAL = "schedule_proposal"
|
|
TODO_LIST = "todo_list"
|
|
DRAFT_MESSAGE = "draft_message"
|
|
SUMMARY = "summary"
|
|
PROGRESS_SUMMARY = "progress_summary"
|
|
DIAGNOSIS_REPORT = "diagnosis_report"
|
|
STRUCTURED_PAYLOAD = "structured_payload"
|
|
EXECUTION_REPORT = "execution_report"
|
|
CLARIFICATION_REQUEST = "clarification_request"
|
|
SAFETY_BLOCK = "safety_block"
|
|
ERROR_REPORT = "error_report"
|
|
UNKNOWN = "unknown"
|
|
|
|
|
|
class TaskTyping(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
primary: TaskType
|
|
secondary: list[TaskType] = Field(default_factory=list, max_length=3)
|
|
|
|
|
|
class ResultTyping(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
primary: ResultType
|
|
secondary: list[ResultType] = Field(default_factory=list, max_length=3)
|
|
|
|
|
|
class ExecutionMode(str, Enum):
|
|
ONESTEP = "onestep"
|
|
TOOL_ASSISTED = "tool_assisted"
|
|
MULTISTEP = "multistep"
|
|
|
|
|
|
class RunStatus(str, Enum):
|
|
SUCCESS = "success"
|
|
PARTIAL_SUCCESS = "partial_success"
|
|
FAILED = "failed"
|
|
|
|
|
|
class ToolStatus(str, Enum):
|
|
SUCCESS = "success"
|
|
FAILURE = "failure"
|
|
PARTIAL = "partial"
|
|
|
|
|
|
class KeyEntity(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
name: str
|
|
type: str
|
|
value: str | None = None
|
|
|
|
@field_validator("value", mode="before")
|
|
@classmethod
|
|
def normalize_value(cls, value: object) -> object:
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, str):
|
|
return value
|
|
if isinstance(value, bool | int | float):
|
|
return str(value)
|
|
return value
|
|
|
|
|
|
class ConstraintItem(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
key: str
|
|
value: str
|
|
required: bool = True
|
|
|
|
@field_validator("value", mode="before")
|
|
@classmethod
|
|
def normalize_value(cls, value: object) -> object:
|
|
if isinstance(value, bool | int | float):
|
|
return str(value)
|
|
return value
|
|
|
|
|
|
class NormalizedTaskInput(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
user_text: str
|
|
multimodal_summary: list[str] = Field(default_factory=list)
|
|
context_summary: str = Field(default="", max_length=2000)
|
|
|
|
|
|
class RouterAgentOutput(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
normalized_task_input: NormalizedTaskInput
|
|
key_entities: list[KeyEntity] = Field(default_factory=list)
|
|
constraints: list[ConstraintItem] = Field(default_factory=list)
|
|
task_typing: TaskTyping
|
|
execution_mode: ExecutionMode
|
|
result_typing: ResultTyping
|
|
|
|
|
|
class ErrorInfo(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
code: str
|
|
message: str
|
|
retryable: bool = False
|
|
details: dict[str, Any] | None = None
|
|
|
|
|
|
class ToolAgentOutput(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
tool_name: str
|
|
tool_call_id: str
|
|
tool_call_args: dict[str, Any] | None = None
|
|
status: ToolStatus
|
|
result: str
|
|
error: ErrorInfo | None = None
|
|
|
|
|
|
class WorkerAgentOutputLite(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
status: RunStatus = RunStatus.SUCCESS
|
|
answer: str
|
|
key_points: list[str] = Field(default_factory=list)
|
|
result_type: ResultType = ResultType.UNKNOWN
|
|
suggested_actions: list[str] = Field(default_factory=list)
|
|
error: ErrorInfo | None = None
|
|
|
|
|
|
class WorkerAgentOutputRich(WorkerAgentOutputLite):
|
|
ui_hints: UiHintsPayload | None = None
|
|
|
|
|
|
class AgentOutput(WorkerAgentOutputRich):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
|
|
WorkerAgentOutput = WorkerAgentOutputLite | WorkerAgentOutputRich
|
|
|
|
|
|
def resolve_worker_output_model(
|
|
execution_mode: ExecutionMode,
|
|
) -> type[WorkerAgentOutputLite]:
|
|
if execution_mode == ExecutionMode.ONESTEP:
|
|
return WorkerAgentOutputLite
|
|
return WorkerAgentOutputRich
|