Files
loki/examples/langchain-sisyphus/sisyphus_langchain/state.py

101 lines
3.8 KiB
Python

"""
Shared state schema for the Sisyphus orchestrator graph.
In LangGraph, state is the single source of truth that flows through every node.
This is analogous to Loki's per-agent session context, but unified into one typed
dictionary that the entire graph shares.
Loki Concept Mapping:
- Loki session context → SisyphusState (TypedDict)
- Loki todo__init / todo__add → SisyphusState.todos list
- Loki agent__spawn outputs → SisyphusState.agent_outputs dict
- Loki intent classification → SisyphusState.intent field
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Annotated, Literal
from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages
from typing_extensions import TypedDict
# ---------------------------------------------------------------------------
# Intent types — mirrors Loki's Sisyphus classification table
# ---------------------------------------------------------------------------
IntentType = Literal[
"trivial", # Single file, known location, typo fix → handle yourself
"exploration", # "Find X", "Where is Y" → spawn explore
"implementation", # "Add feature", "Fix bug" → spawn coder
"architecture", # Design questions, oracle triggers → spawn oracle
"ambiguous", # Unclear scope → ask user
]
# ---------------------------------------------------------------------------
# Todo item — mirrors Loki's built-in todo system
# ---------------------------------------------------------------------------
@dataclass
class TodoItem:
"""A single task in the orchestrator's todo list."""
id: int
task: str
done: bool = False
def _merge_todos(existing: list[TodoItem], new: list[TodoItem]) -> list[TodoItem]:
"""
Reducer for the todos field.
LangGraph requires a reducer for any state field that can be written by
multiple nodes. This merges by id: if a todo with the same id already
exists, the incoming version wins (allows marking done).
"""
by_id = {t.id: t for t in existing}
for t in new:
by_id[t.id] = t
return list(by_id.values())
# ---------------------------------------------------------------------------
# Core graph state
# ---------------------------------------------------------------------------
class SisyphusState(TypedDict):
"""
The shared state that flows through every node in the Sisyphus graph.
Annotated fields use *reducers* — functions that merge concurrent writes.
Without reducers, parallel node outputs would overwrite each other.
"""
# Conversation history — the `add_messages` reducer appends new messages
# instead of replacing the list. This is critical: every node adds its
# response here, and downstream nodes see the full history.
#
# Loki equivalent: each agent's chat session accumulates messages the same
# way, but messages are scoped per-agent. In LangGraph the shared message
# list IS the inter-agent communication channel.
messages: Annotated[list[BaseMessage], add_messages]
# Classified intent for the current request
intent: IntentType
# Which agent the supervisor routed to last
next_agent: str
# Iteration counter — safety valve analogous to Loki's max_auto_continues
iteration_count: int
# Todo list for multi-step tracking (mirrors Loki's todo__* tools)
todos: Annotated[list[TodoItem], _merge_todos]
# Accumulated outputs from sub-agent nodes, keyed by agent name.
# The supervisor reads these to decide what to do next.
agent_outputs: dict[str, str]
# Final synthesized answer to return to the user
final_output: str
# The working directory / project path (mirrors Loki's project_dir variable)
project_dir: str