src/reasoning/state.py

"""
LangGraph State Definition
Defines the RAGState schema used by all reasoning nodes.
"""

from typing import Any, TypedDict


class RAGState(TypedDict):
    """
    State schema for the RAG Reasoning Engine.
    Tracks query, context, intermediate steps, and performance metrics.
    """

    # --- Core Input/Output ---
    query: str
    generated_answer: str

    # --- Tenant Isolation ---
    tenant_id: str

    # --- Intermediate Data ---
    sub_tasks: list[str]  # Generated by Planner
    retrieved_context: list[dict[str, Any]]  # Fetched by Retrieval Agent
    current_node: str  # Tracks graph progress for debugging

    # --- Validation & Error Handling ---
    validation_passed: bool
    error_message: str | None

    # --- Performance Tracking (Strategic Requirement #2) ---
    node_latency_ms: dict[str, float]
    total_latency_ms: float

    # --- LLM Configuration ---
    llm_api_key: str | None  # User-provided key for this request (optional)

    # --- Guardrails ---
    pii_redacted_query: str | None  # Query with PII redacted (if PII was detected)
    total_tokens_used: int  # Total tokens consumed across all LLM calls
    source_files: list[str]  # Unique source filenames cited in the answer