backend/state/graph_state.py

"""
backend/state/graph_state.py

The WorkflowState is the single shared data structure
that flows through every agent node in the LangGraph.

Design principles:
  - Immutable history: agents append, never mutate past entries
  - Full audit trail: every decision, tool call, error recorded
  - TypedDict for LangGraph compatibility
  - Rich metadata for debugging and UI rendering
"""
from __future__ import annotations
import uuid
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Annotated
from typing_extensions import TypedDict
from langgraph.graph.message import add_messages
from langchain_core.messages import BaseMessage


class TaskStatus(str, Enum):
    PENDING   = "pending"
    PLANNING  = "planning"
    EXECUTING = "executing"
    REFLECTING = "reflecting"
    COMPLETED  = "completed"
    FAILED     = "failed"
    CANCELLED  = "cancelled"


class StepStatus(str, Enum):
    PENDING  = "pending"
    RUNNING  = "running"
    DONE     = "done"
    FAILED   = "failed"
    SKIPPED  = "skipped"


class AgentRole(str, Enum):
    PLANNER  = "planner"
    EXECUTOR = "executor"
    CRITIC   = "critic"
    MEMORY   = "memory"
    SYSTEM   = "system"


# ── Sub-models (plain dicts for TypedDict compat) ─────────────────────────────

def make_plan_step(
    step_id: str,
    title: str,
    description: str,
    tool: str | None = None,
    depends_on: list[str] | None = None,
) -> dict:
    return {
        "step_id": step_id,
        "title": title,
        "description": description,
        "tool": tool,
        "depends_on": depends_on or [],
        "status": StepStatus.PENDING,
        "result": None,
        "error": None,
        "attempts": 0,
        "started_at": None,
        "finished_at": None,
    }


def make_agent_event(
    agent: AgentRole,
    event_type: str,
    content: str,
    metadata: dict | None = None,
) -> dict:
    return {
        "event_id": str(uuid.uuid4())[:8],
        "agent": agent,
        "event_type": event_type,       # "thought" | "tool_call" | "tool_result" | "decision" | "error"
        "content": content,
        "metadata": metadata or {},
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }


def make_memory_entry(
    content: str,
    memory_type: str = "episodic",     # "episodic" | "semantic" | "procedural"
    importance: float = 0.5,           # 0–1
    tags: list[str] | None = None,
) -> dict:
    return {
        "memory_id": str(uuid.uuid4())[:8],
        "content": content,
        "memory_type": memory_type,
        "importance": importance,
        "tags": tags or [],
        "created_at": datetime.now(timezone.utc).isoformat(),
        "access_count": 0,
    }


# ── Main WorkflowState ────────────────────────────────────────────────────────

class WorkflowState(TypedDict):
    # Identity
    task_id: str
    task: str                           # Original user task
    status: TaskStatus

    # Messages (LangGraph built-in — auto-appended)
    messages: Annotated[list[BaseMessage], add_messages]

    # Planning
    plan: list[dict]                    # List of plan steps
    current_step_index: int
    iteration: int                      # How many planner loops

    # Execution
    step_results: dict[str, Any]        # step_id → result
    tool_calls_log: list[dict]          # All tool calls made

    # Agent events (audit trail)
    events: list[dict]                  # All agent events

    # Critic
    critique: str | None                # Latest critique
    quality_score: float | None         # 0–100
    needs_replanning: bool              # Critic flagged replan needed

    # Memory
    memories: list[dict]                # Retrieved relevant memories
    new_memories: list[dict]            # Memories to store

    # Output
    final_output: str | None
    error_message: str | None

    # Metadata
    created_at: str
    updated_at: str
    total_tokens: int


def create_initial_state(task: str, task_id: str | None = None) -> WorkflowState:
    """Create a fresh WorkflowState for a new task."""
    now = datetime.now(timezone.utc).isoformat()
    return WorkflowState(
        task_id=task_id or str(uuid.uuid4()),
        task=task,
        status=TaskStatus.PENDING,
        messages=[],
        plan=[],
        current_step_index=0,
        iteration=0,
        step_results={},
        tool_calls_log=[],
        events=[make_agent_event(AgentRole.SYSTEM, "task_created", f"Task created: {task}")],
        critique=None,
        quality_score=None,
        needs_replanning=False,
        memories=[],
        new_memories=[],
        final_output=None,
        error_message=None,
        created_at=now,
        updated_at=now,
        total_tokens=0,
    )