Spaces:

NYSERDA-CRE-Working-Group
/

Updated_code_complaince

Sleeping

Upload 70 files

e1ced8e verified 8 months ago

4.61 kB

"""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

	"""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