82 lines
2.1 KiB
Python
82 lines
2.1 KiB
Python
from dataclasses import dataclass, field
|
|
from typing import TypedDict, Annotated, Sequence
|
|
from enum import Enum
|
|
|
|
from langchain_core.messages import BaseMessage
|
|
from langgraph.graph.message import add_messages
|
|
|
|
|
|
class AgentRole(str, Enum):
|
|
MASTER = "master"
|
|
SCHEDULE_PLANNER = "schedule_planner"
|
|
EXECUTOR = "executor"
|
|
LIBRARIAN = "librarian"
|
|
ANALYST = "analyst"
|
|
|
|
|
|
@dataclass
|
|
class ConversationTurn:
|
|
role: str # "user" | "assistant"
|
|
content: str
|
|
agent: AgentRole | None = None
|
|
model: str | None = None
|
|
|
|
|
|
class AgentState(TypedDict):
|
|
# Core message history with add_messages reducer
|
|
messages: Annotated[list[BaseMessage], add_messages]
|
|
|
|
# Session identifiers
|
|
user_id: str
|
|
conversation_id: str
|
|
|
|
# Agent routing state
|
|
current_agent: str | None
|
|
next_step: str | None # For explicit graph routing
|
|
|
|
# Traceability
|
|
agent_trace: list[str]
|
|
|
|
# Task & Entity Tracking (Business Logic)
|
|
pending_tasks: list[dict]
|
|
completed_tasks: list[dict]
|
|
created_entities: list[dict]
|
|
|
|
# Context summaries (for long-term or cross-agent context)
|
|
knowledge_context: str | None
|
|
schedule_context_summary: str | None
|
|
analysis_report: str | None
|
|
|
|
# Output control
|
|
final_response: str | None
|
|
|
|
# Memory & Environment
|
|
memory_context: str | None
|
|
current_datetime_context: str | None
|
|
|
|
# Configuration
|
|
user_llm_config: dict | None
|
|
provider_capabilities: dict | 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.value,
|
|
next_step=None,
|
|
agent_trace=[AgentRole.MASTER.value],
|
|
pending_tasks=[],
|
|
completed_tasks=[],
|
|
created_entities=[],
|
|
knowledge_context=None,
|
|
schedule_context_summary=None,
|
|
analysis_report=None,
|
|
final_response=None,
|
|
memory_context=None,
|
|
current_datetime_context=None,
|
|
user_llm_config=None,
|
|
provider_capabilities=None,
|
|
)
|