state.py

"""State schema for the Code Compliance Bot — extends drawing_reader patterns."""
from __future__ import annotations

import operator
from typing import Annotated, TypedDict

from langgraph.graph import add_messages

# ---------------------------------------------------------------------------
# Reusable types from drawing_reader (re-declared here to stay self-contained)
# ---------------------------------------------------------------------------


class ImageRef(TypedDict):
    """Lightweight reference to an image stored on disk."""
    id: str
    path: str
    label: str
    page_num: int
    crop_type: str       # "full_page" | "crop" | "annotated"
    width: int
    height: int


class CropTask(TypedDict):
    """A single crop+annotate instruction."""
    page_num: int
    crop_instruction: str
    annotate: bool
    annotation_prompt: str
    label: str
    priority: int


class PageMetadata(TypedDict):
    """Structured description of a single PDF page."""
    page_num: int
    sheet_id: str
    sheet_title: str
    discipline: str
    page_type: str
    description: str
    key_elements: list[str]
    spatial_coverage: str


# ---------------------------------------------------------------------------
# New types for code compliance
# ---------------------------------------------------------------------------


class CodeQuery(TypedDict):
    """A single code lookup request produced by the planner or analyst."""
    query: str              # Natural language query for ChromaDB semantic search
    focus_area: str         # e.g., "egress", "fire_protection", "occupancy"
    context: str            # Why this lookup is needed
    priority: int           # 0 = critical, 1 = supplementary


class CodeSection(TypedDict):
    """A retrieved code section from ChromaDB."""
    section_full: str       # e.g., "1005.1"
    code_type: str          # e.g., "Building"
    parent_major: str       # e.g., "10"
    text: str               # Full section text
    relevance: str          # Why this section was retrieved


class AgentMessage(TypedDict):
    """A single message in the inter-agent discussion log."""
    timestamp: str
    agent: str              # "planner" | "code_analyst" | "compliance_analyst" | "reviewer"
    action: str             # "plan" | "search_code" | "analyze" | "request_more" | "review" | "verdict"
    summary: str            # One-line summary for collapsed UI view
    detail: str             # Full content for expanded view
    evidence_refs: list[str]  # Image crop IDs or code section IDs referenced


# ---------------------------------------------------------------------------
# Main state
# ---------------------------------------------------------------------------


class ComplianceState(TypedDict):
    """Full state for the code compliance LangGraph workflow."""

    # -- Conversation --
    messages: Annotated[list, add_messages]
    question: str

    # -- Document (set at ingest time) --
    pdf_path: str
    page_image_dir: str
    num_pages: int

    # -- Page Metadata --
    page_metadata_json: str       # JSON string of per-page metadata

    # -- Planning (from compliance_planner) --
    legend_pages: list[int]
    target_pages: list[int]
    crop_tasks: list[CropTask]
    code_queries: list[CodeQuery]

    # -- Images (additive reducer — new refs are appended, never replaced) --
    image_refs: Annotated[list[ImageRef], operator.add]

    # -- Code Lookup Results (additive reducer) --
    code_sections: Annotated[list[CodeSection], operator.add]
    code_report: str
    code_chapters_fetched: list[str]

    # -- Analysis --
    compliance_analysis: str
    reviewer_analysis: str
    final_verdict: str

    # -- Discussion Log (additive reducer — the "agents talking" feature) --
    discussion_log: Annotated[list[AgentMessage], operator.add]

    # -- Investigation Loop Control --
    additional_crop_tasks: list[CropTask]
    additional_code_queries: list[CodeQuery]
    needs_more_investigation: bool
    investigation_round: int
    max_rounds: int

    # -- Settings --
    enable_consensus: bool
    enable_annotation: bool

    # -- Status (additive reducer — parallel nodes can both write status updates) --
    status_message: Annotated[list[str], operator.add]


# Alias so that copied drawing_reader nodes (cropper, annotator) which import
# ``DrawingReaderState`` keep working — ComplianceState is a superset.
DrawingReaderState = ComplianceState