106 lines
2.3 KiB
Python
106 lines
2.3 KiB
Python
from dataclasses import dataclass
|
|
from typing import TypedDict, Annotated
|
|
from enum import Enum
|
|
|
|
|
|
class AgentRole(str, Enum):
|
|
MASTER = "master"
|
|
PLANNER = "planner"
|
|
EXECUTOR = "executor"
|
|
LIBRARIAN = "librarian"
|
|
ANALYST = "analyst"
|
|
|
|
|
|
@dataclass
|
|
class AgentInfo:
|
|
name: str
|
|
role: AgentRole
|
|
description: str
|
|
|
|
|
|
@dataclass
|
|
class ToolCall:
|
|
tool: str
|
|
args: dict
|
|
result: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class ConversationTurn:
|
|
role: str # "user" | "assistant"
|
|
content: str
|
|
agent: AgentRole | None = None
|
|
model: str | None = None
|
|
|
|
|
|
def turn_to_message(turn: ConversationTurn) -> HumanMessage:
|
|
return HumanMessage(content=turn.content)
|
|
|
|
|
|
def message_to_turn(msg, agent: AgentRole | None = None) -> ConversationTurn:
|
|
msg_type = getattr(msg, "type", None) or getattr(msg, "role", "assistant")
|
|
return ConversationTurn(
|
|
role="user" if msg_type in ("human", "user") else "assistant",
|
|
content=msg.content,
|
|
agent=agent,
|
|
model=getattr(msg, "model", None),
|
|
)
|
|
|
|
|
|
class AgentState(TypedDict):
|
|
messages: Annotated[list, None]
|
|
user_id: str
|
|
conversation_id: str
|
|
|
|
# Agent routing
|
|
current_agent: AgentRole
|
|
active_agents: list[AgentRole]
|
|
|
|
# Task tracking
|
|
pending_tasks: list[dict]
|
|
completed_tasks: list[dict]
|
|
|
|
# Tool usage
|
|
tool_calls: list[ToolCall]
|
|
last_tool_result: str | None
|
|
|
|
# Knowledge context
|
|
knowledge_context: str | None
|
|
graph_context: str | None
|
|
|
|
# Planning
|
|
plan: str | None
|
|
plan_steps: list[dict]
|
|
|
|
# Analysis
|
|
analysis_report: str | None
|
|
|
|
# Output control
|
|
final_response: str | None
|
|
should_respond: bool
|
|
|
|
# Memory context (injected at start of each conversation)
|
|
memory_context: str | None
|
|
|
|
|
|
def initial_state(user_id: str, conversation_id: str) -> AgentState:
|
|
return AgentState(
|
|
messages=[],
|
|
user_id=user_id,
|
|
conversation_id=conversation_id,
|
|
current_agent=AgentRole.MASTER,
|
|
active_agents=[AgentRole.MASTER],
|
|
pending_tasks=[],
|
|
completed_tasks=[],
|
|
tool_calls=[],
|
|
last_tool_result=None,
|
|
knowledge_context=None,
|
|
graph_context=None,
|
|
plan=None,
|
|
plan_steps=[],
|
|
analysis_report=None,
|
|
final_response=None,
|
|
should_respond=True,
|
|
memory_context=None,
|
|
)
|