Update app.py

app.py

@@ -1,1156 +1,1158 @@
-import gradio as gr
-import json
-import os
-import re
-from google import genai
-from google.genai import types
-import chromadb
-from chromadb.utils import embedding_functions
-from collections import Counter
-import base64
-import io
-from PIL import Image
-import matplotlib.pyplot as plt
-import openai
-from datetime import datetime
-import threading
-
-os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
-os.environ["GEMINI_API_KEY"] = os.getenv("GEMINI_API_KEY")
-
-# Global state for the interface
-class InterfaceState:
-    def __init__(self):
-        self.log_messages = []
-        self.analysis_messages = []
-        self.current_chapter = ""
-        self.current_images = []
-        self.staged_audit_images = []
-        self.final_answer = ""
-        self.done = False
-        self.lock = threading.Lock()
-    
-    def add_log(self, message):
-        timestamp = datetime.now().strftime("%H:%M:%S")
-        with self.lock:
-            self.log_messages.append(f"**[{timestamp}]** {message}")
-            return "\n\n".join(self.log_messages)
-
-    def add_analysis(self, message):
-        timestamp = datetime.now().strftime("%H:%M:%S")
-        with self.lock:
-            self.analysis_messages.append(f"**[{timestamp}]** {message}")
-            return "\n\n".join(self.analysis_messages)
-
-    def set_chapter(self, chapter_text):
-        with self.lock:
-            self.current_chapter = chapter_text
-            return chapter_text
-
-    def add_image(self, img_pil):
-        with self.lock:
-            self.current_images.append(img_pil)
-            return self.current_images.copy()
-        
-    def add_staged_image_part(self, image_part):
-        """Thread-safe method to stage images for the Gemini Audit."""
-        with self.lock:
-            self.staged_audit_images.append(image_part)
-            # Log it so we can verify it happened in the console
-            print(f"DEBUG: Staged image part. Total staged: {len(self.staged_audit_images)}")
-
-    def get_staged_images(self):
-        """Safely retrieve the staged images for the audit turn."""
-        with self.lock:
-            return list(self.staged_audit_images) # Return a copy to prevent mutation
-    
-    def clear(self):
-        with self.lock:
-            self.log_messages.clear()
-            self.analysis_messages.clear()
-            self.current_chapter = ""
-            self.current_images.clear()
-            self.final_answer = ""
-            self.done = False
-            
-    
-
-state = InterfaceState()
-
-# Load your data (same as original)
-with open('Preprocessed Files/page_metadata.json', 'r') as json_file:
-    page_metadata = json.load(json_file)
-page_metadata = {int(k): v for k, v in page_metadata.items()}
-
-with open('Preprocessed Files/text_list.json', 'r') as json_file:
-    text_list = json.load(json_file)
-    
-with open('Preprocessed Files/tile_metadata.json', 'r') as json_file:
-    tile_metadata = json.load(json_file)
-tile_metadata = {
-    int(outer_k): {
-        int(inner_k): inner_v
-        for inner_k, inner_v in outer_v.items()
-    }
-    for outer_k, outer_v in tile_metadata.items()
-}
-
-def load_fullpage_images(folder="Images"):
-    files = os.listdir(folder)
-    page_files = []
-    for f in files:
-        match = re.search(r"page_(\d+)_fullpage\.png", f)
-        if match:
-            page_num = int(match.group(1))
-            page_files.append((page_num, f))
-    page_files.sort(key=lambda x: x[0])
-    image_bytes_list = []
-    for page_num, filename in page_files:
-        path = os.path.join(folder, filename)
-        with open(path, "rb") as f:
-            img_bytes = f.read()
-        image_bytes_list.append(img_bytes)
-    return image_bytes_list
-
-def load_tile_images(page):
-    files = os.listdir('Tiles')
-    page_files = []
-    for f in files:
-        match = re.search(f"page_{page}_tile_(\d+)\.png", f)
-        if match:
-            page_num = int(match.group(1))
-            page_files.append((page_num, f))
-    page_files.sort(key=lambda x: x[0])
-    image_bytes_list = []
-    for page_num, filename in page_files:
-        path = os.path.join('Tiles', filename)
-        with open(path, "rb") as f:
-            img_bytes = f.read()
-        image_bytes_list.append(img_bytes)
-    return image_bytes_list
-
-image_bytes_list = load_fullpage_images()
-
-tile_bytes = {}
-for page in range(44):
-    tile_list = load_tile_images(page)
-    if tile_list:
-        tile_bytes[page] = load_tile_images(page)
-
-# Vector Code Base
-chroma_client = chromadb.PersistentClient(path="nyc_code_db")
-embedding_model = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
-collection = chroma_client.get_collection(name="nyc_building_codes", embedding_function=embedding_model)
-
-all_pending_images = []
-
-# Modified tool functions with Gradio updates
-def search_page_text(page_number: int, research_goal: str):
-    state.add_log(f'🔍 Searching page **{page_metadata[page_number]["sheet_title"]}** for details')
-    
-    state.add_analysis(
-        f'🔍 Searching page {page_metadata[page_number]["sheet_title"]} with prompt\n{research_goal}'
-    )
-    
-    raw_text = text_list[page_number]
-    
-    client = openai.OpenAI()
-    response = client.chat.completions.create(
-        model="gpt-5-mini", 
-        messages=[
-            {"role": "system", "content": """
-            You are a Fast NYC Plans Examiner Signal Agent.
-            
-            Your ONLY job is to extract **code-relevant signals** from the OCR text of a SINGLE drawing page.
-            You do NOT interpret the law and you do NOT summarize design intent.
-            
-            Your output will be used to CONSTRAIN a downstream legal research agent.
-            
-            ========================
-            WHAT TO EXTRACT
-            ========================
-            Look only for information that determines which parts of the NYC Code apply such as:
-            
-            - Occupancy classification (e.g., R-2, A-3, M, S, F, mixed-use)
-            - Building height (stories, feet, high-rise indicators)
-            - Construction type (I, II, III, IV, V)
-            - Fire protection systems (sprinklers, standpipes, fire alarm, smoke control)
-            - Means of egress references (stairs, exits, exit access, doors, corridors)
-            - Structural system hints (steel, concrete, load-bearing walls, columns, transfer girders)
-            - Mechanical / fuel / plumbing system mentions (boilers, gas piping, HVAC type, shafts)
-            - Zoning or special district references (if present)
-            - Scope flags (new building, alteration, addition, change of occupancy, retrofit)
-            
-            However only return relevant signals to the provided research goal.
-            
-            ========================
-            OUTPUT FORMAT (STRICT MARKDOWN)
-            ========================
-            Return ONLY the following sections:
-            
-            ### Code-Relevant Signals
-            - Bullet list of extracted facts
-            
-            ### Likely Governing Code Domains
-            - One-line list chosen from: Administrative, Building, Mechanical, FuelGas, Plumbing, Fire
-            
-            ### Text Evidence
-            - Short quoted snippets from the page that support each signal
-            
-            ========================
-            RULES
-            ========================
-            - Do NOT speculate
-            - If a signal is not present, omit it
-            - Prefer exact phrases over paraphrase
-            - Keep total length under 500 words
-            - No legal conclusions, no compliance advice
-            """},
-            {"role": "user", "content": f"PAGE TEXT:\n{raw_text}\n\nRESEARCH GOAL: {research_goal}\n\nReturn a breif but comprehensive Markdown summary of your findings and justification with text snippets."}
-        ]
-    )
-    
-    analysis_text = response.choices[0].message.content
-
-    state.add_analysis(
-        f"🟦 Text Analyst (Page {page_number})\n{analysis_text}"
-    )
-    
-    return {
-        "page": page_number,
-        "summary": analysis_text
-    }
-
-def discover_code_locations(query: str):
-    state.add_log(f'📚 Searching NYC Code for: **{query}**')
-    
-    results = collection.query(
-        query_texts=[query],
-        n_results=25,
-        include=["metadatas", "documents"]
-    )
-    
-    if not results['metadatas'][0]:
-        return "No results found. Try a different technical keyword."
-
-    metas = results['metadatas'][0]
-    docs = results['documents'][0]
-    
-    category_chapter_pairs = [f"{m['code_type']} | Ch. {m['parent_major']}" for m in metas]
-    counts = Counter(category_chapter_pairs)
-    chapter_summary = "\n".join([f"- {pair} ({count} hits)" for pair, count in counts.most_common(5)])
-
-    section_reports = []
-    for m, doc in zip(metas, docs):
-        report = (
-            f"ID: {m['section_full']} | Code: {m['code_type']} | Chapter: {m['parent_major']}\n"
-            f"Snippet: {doc}"
-        )
-        section_reports.append(report)
-
-    output = (
-        "### CODE DISCOVERY REPORT ###\n"
-        f"MOST RELEVANT CHAPTERS:\n{chapter_summary}\n\n"
-        "TOP RELEVANT SECTIONS:\n" + 
-        "\n---\n".join(section_reports)
-    )
-    
-    return output
-
-def fetch_full_chapter(code_type: str, chapter_id: str):
-    state.add_log(f'📖 Fetching Chapter **{chapter_id}** from **{code_type}** code')
-    
-    try:
-        chapter_data = collection.get(
-            where={
-                "$and": [
-                    {"code_type": {"$eq": code_type}},
-                    {"parent_major": {"$eq": chapter_id}}
-                ]
-            },
-            include=["documents", "metadatas"]
-        )
-
-        if not chapter_data['documents']:
-            return f"No documentation found for {code_type} Chapter {chapter_id}."
-
-        sections = sorted(zip(chapter_data['metadatas'], chapter_data['documents']), 
-                          key=lambda x: x[0]['section_full'])
-
-        full_text = f"## FULL LEGAL TEXT: {code_type.upper()} CODE - CHAPTER {chapter_id}\n\n"
-
-        for meta, doc in sections:
-            blocks = doc.split("[CONT.]:")
-            unique_blocks = []
-            for b in blocks:
-                clean_b = b.strip()
-                if clean_b and clean_b not in unique_blocks:
-                    unique_blocks.append(clean_b)
-            
-            clean_doc = " ".join(unique_blocks)
-            full_text += f"### SECTION {meta['section_full']}\n{clean_doc}\n\n---\n\n"
-
-        # Update the chapter display
-        state.set_chapter(full_text)
-        
-        return full_text
-
-    except Exception as e:
-        return f"Error retrieving chapter content: {str(e)}"
-
-def nyc_legal_sub_agent(research_goal: str):
-    state.add_log(f'⚖️ Investigating NYC Code for: **{research_goal}**')
-    
-    state.add_analysis(
-        f"⚖️ Legal Analyst is searching\n{research_goal}"
-    )
-    
-    client = openai.OpenAI()
-    
-    internal_tools = [
-        {
-            "type": "function",
-            "function": {
-                "name": "discover_code_locations",
-                "description": "Scans NYC code in a semantic vector database. Use this FIRST to find which chapters/sections are relevant.",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "query": {"type": "string", "description": "semantic search string for a vector database (Not a keyword search use a full sentence)"}
-                    },
-                    "required": ["query"]
-                }
-            }
-        },
-        {
-            "type": "function",
-            "function": {
-                "name": "fetch_full_chapter",
-                "description": "Retrieves the full legal text of a specific chapter for deep analysis.",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "code_type": {
-                            "type": "string", 
-                            "enum": ["Administrative", "Building", "FuelGas", "Mechanical", "Plumbing"],
-                            "description": "The specific NYC code volume to search."
-                        },
-                        "chapter_id": {"type": "string", "description": "The chapter number string"}
-                    },
-                    "required": ["code_type", "chapter_id"]
-                }
-            }
-        }
-    ]
-
-    messages = [
-        {"role": "system", "content": """
-        You are a Senior NYC Building Code Consultant and Legal Research Agent.
-        
-        Your task is to produce a **definitive, citation-backed legal report** that can be used directly by a downstream orchestration agent. 
-        Accuracy, traceability, and completeness matter more than brevity.
-        
-        ========================
-        PRIMARY OBJECTIVE
-        ========================
-        Given a research goal, identify and analyze relevant NYC Code provisions, including:
-        - Governing sections
-        - Exceptions
-        - Cross-references
-        - Related chapters that modify, limit, or expand the rule
-        
-        Every legal claim MUST be supported by a specific code citation.
-        
-        You are operating in FAST LEGAL MODE.
-
-        SEARCH BUDGET:
-        - Maximum of 2 calls to `discover_code_locations`
-        - Maximum of 2 calls to `fetch_full_chapter`
-        
-        STOP CONDITIONS:
-        - If the first chapter fetch contains governing text AND exceptions, STOP and synthesize.
-        - Only fetch a second chapter if the first chapter explicitly cross-references another chapter.
-        
-        PRIORITY ORDER:
-        1) Governing rule section
-        2) Exceptions
-        3) Cross-references that MODIFY the rule
-        Ignore definitions and administrative content unless directly referenced.
-        
-        GOOD ENOUGH STANDARD:
-        If you can identify:
-        - The governing section
-        - At least one exception or limitation
-        You must STOP and report.
-        
-        ========================
-        TOOL STRATEGY (MANDATORY)
-        ========================
-        This is a semantic vector database, NOT a keyword index. Always search in full English questions.
-        
-        1) FIRST — Call `discover_code_locations`
-           - Use a natural-language query describing the legal requirement you are trying to find
-           - Example: "What NYC Building Code sections regulate emergency egress width in residential buildings"
-           NEVER use a keyword search thi will not work you are searching a vector database. 
-           If you know what chaoter you need call the fetch_full_chapter tool instead.
-           If you perform TWO consecutive `discover_code_locations` calls
-            and both return no new relevant chapters or sections:
-            
-            You MUST stop searching and do one of the following:
-            - Conclude that the table/section does NOT exist as a standalone provision in the NYC Code corpus, OR
-            - Conclude that the requirement is embedded within the previously retrieved sections
-            
-            Then proceed to report findings using the closest governing section.
-            
-            DO NOT continue reformulating the same query.
-            You MUST NOT call `discover_code_locations` more than once for the same legal concept.
-            If a new query is semantically similar to a prior query, STOP and move forward with analysis.
-            
-        2) SECOND — Call `fetch_full_chapter`
-           - If multiple relevant sections appear in the same chapter
-           - OR if a section contains exceptions, references, or conditional language
-           - OR if you know what section of the code is relevant and want to see a full chapter
-        
-        3) THIRD — Follow Cross-References
-           - If a section says "See Section X", "As required by Chapter Y", or "Except as permitted in..."
-           - You MUST search and retrieve those sections as well
-        
-        4) STOP ONLY WHEN
-           - All exceptions are reviewed
-           - All cross-references are resolved
-           - No additional modifying sections remain
-        
-        ========================
-        OUTPUT FORMAT (STRICT)
-        ========================
-        Return a structured legal report in the following format:
-        
-        ### Legal Summary
-        Brief, plain-language explanation of what the code requires.
-        
-        ### Governing Code Sections
-        - **[Code Type] §[Section Number] — [Title]**
-          - Summary:
-          - Key Requirements:
-          - Applicability Conditions:
-          - Exceptions:
-        
-        ### Cross-References Analyzed
-        - **§[Section Number] — [Title]**
-          - Why It Matters:
-          - Impact on Main Rule:
-        
-        ### Edge Cases & Enforcement Notes
-        - Special conditions (building type, occupancy class, height, system type, jurisdictional notes)
-        - Common misinterpretations
-        - DOB or FDNY enforcement implications (if relevant)
-        
-        ### Compliance Checklist
-        - Bullet list of actionable compliance steps derived from the code
-        
-        ========================
-        QUALITY RULES
-        ========================
-        - NEVER summarize without citing
-        - NEVER assume jurisdiction, building type, or occupancy unless the code explicitly states it
-        - If legal text is ambiguous, flag it as **Interpretive**
-        - Prefer quoting short legal phrases when clarity matters
-        
-        ========================
-        TONE
-        ========================
-        Professional. Precise. Legal-research quality. No speculation.
-        """},
-        {"role": "user", "content": f"Analyze the NYC building code with this goal: {research_goal}"}
-    ]
-
-    for _ in range(20):
-        response = client.chat.completions.create(
-            model="gpt-5-mini",
-            messages=messages,
-            tools=internal_tools,
-            tool_choice="auto"
-        )
-        
-        msg = response.choices[0].message
-        messages.append(msg)
-
-        if not msg.tool_calls:
-            
-            break
-
-        for tool_call in msg.tool_calls:
-            func_name = tool_call.function.name
-            args = json.loads(tool_call.function.arguments)
-            
-            if func_name == "discover_code_locations":
-                result = discover_code_locations(args['query'])
-            elif func_name == "fetch_full_chapter":
-                result = fetch_full_chapter(args['code_type'], args['chapter_id'])
-            
-            messages.append({
-                "role": "tool",
-                "tool_call_id": tool_call.id,
-                "content": result
-            })
-    
-    state.add_analysis(
-        f"🟨 Legal Analyst\n{msg.content}"
-    )
-    
-    return msg.content
-
-def merge_tiles(tile_indexes: list[int], page_num: int):
-    state.add_log(f'🔬 Stitching tiles **{tile_indexes}** from page **{page_num}**')
-    
-    images = []
-    positions = []
-    
-    tiles = tile_bytes[page_num]
-    tiles_coords_dict = tile_metadata[page_num]
-
-    for index in tile_indexes:
-        if index < 0 or index >= len(tiles):
-            raise ValueError(f"Tile index {index} out of range")
-
-        img_bytes = tiles[index]
-        if img_bytes is None:
-            raise ValueError(f"No image bytes found for tile {index}")
-
-        img = Image.open(io.BytesIO(img_bytes)).convert('RGBA')
-        images.append(img)
-
-        x = tiles_coords_dict[index]['coords'][0]
-        y = tiles_coords_dict[index]['coords'][1]
-        positions.append((x, y))
-
-    if not images:
-        return None
-
-    min_x = min(x for x, y in positions)
-    min_y = min(y for x, y in positions)
-    normalized_positions = [(x - min_x, y - min_y) for x, y in positions]
-
-    total_width = max(pos[0] + img.width for pos, img in zip(normalized_positions, images))
-    total_height = max(pos[1] + img.height for pos, img in zip(normalized_positions, images))
-
-    stitched_image = Image.new('RGB', (total_width, total_height), (255, 255, 255))
-
-    for img, pos in zip(images, normalized_positions):
-        stitched_image.paste(img, pos)
-    
-    # Add to image gallery
-    
-    output_buffer = io.BytesIO()
-    stitched_image.save(output_buffer, format='PNG')
-    stitched_bytes = output_buffer.getvalue()
-
-    return stitched_bytes
-
-def extract_json(s: str):
-    s = s.strip()
-    start = s.find("{")
-    end = s.rfind("}")
-    if start == -1 or end == -1 or end < start:
-        raise ValueError("No JSON object found in model output:\n" + repr(s))
-    json_str = s[start:end+1]
-    return json.loads(json_str)
-
-def sanitize_tile_indices(data):
-    """
-    Forcefully converts various LLM outputs into a clean list of integers.
-    Handles: [1, 2], ["1", "2"], "1, 2, 3", "[1, 2, 3]", and None.
-    """
-    if not data:
-        return []
-    
-    # If it's already a list, ensure all elements are integers
-    if isinstance(data, list):
-        clean_list = []
-        for item in data:
-            try:
-                # This handles strings inside the list like ["1", "2"]
-                clean_list.append(int(str(item).strip()))
-            except (ValueError, TypeError):
-                continue
-        return clean_list
-
-    # If it's a string, use Regex to find all sequences of digits
-    if isinstance(data, str):
-        # findall returns all non-overlapping matches of the pattern
-        numbers = re.findall(r'\d+', data)
-        return [int(n) for n in numbers]
-
-    return []
-
-def execute_page_expert(expert_instructions: str, page_num: int):
-    state.add_log(f'👁️ Spawning Page Expert for page **{page_num}**')
-    state.add_analysis(f"👁️ Page Expert searching for {expert_instructions}")
-    state.add_log(f'📄 Attaching full-page context for page **{page_num}**')
-    state.add_analysis(
-        f"📄 Full-page context attached for page `{page_num}`"
-    )
-
-    full_page_img = Image.open(
-        io.BytesIO(image_bytes_list[page_num])
-    )
-    state.add_image(full_page_img)
-    
-    client = openai.OpenAI()
-
-    tools = [
-        {
-            "type": "function",
-            "function": {
-                "name": "merge_tiles",
-                "description": "Stitches high-resolution image tiles together into a single zoomed-in view. Use this to read small text, dimensions, or symbols.",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "tile_indexes": {
-                            "type": "array",
-                            "items": {"type": "integer"},
-                            "description": "A list of integer tile IDs from the Grid Map to stitch together."
-                        }
-                    },
-                    "required": ["tile_indexes"]
-                }
-            }
-        }
-    ]
-
-    page_text = text_list[page_num]
-    relevant_tile_meta = tile_metadata[page_num]
-    b64_full_page = base64.b64encode(image_bytes_list[page_num]).decode()
-
-    system_prompt = """
-    You are a Lead AEC Visual Investigator supporting a Compliance Planner.
-    
-    Your mission is to extract **verifiable, high-fidelity evidence** from this drawing page.
-    You must ground every claim in either:
-    - a **Zoomed Tile Image** (via `merge_tiles`) or
-    - a **Direct Text Quote** from the OCR page text.
-    
-    Guesses, assumptions, and general descriptions are not allowed.
-    
-    ========================
-    MANDATORY WORKFLOW
-    ========================
-    1) ORIENT
-       - Review the full-page image and the Grid Map to identify candidate regions.
-       - Decide which tiles likely contain the required evidence. Utilize the tile metadata to assist with this tasl.
-    
-    2) ZOOM (REQUIRED)
-       - You MUST call `merge_tiles(tile_indexes=[...])` before making ANY factual claim about symbols, dimensions, labels, or locations.
-       - Always request ALL tiles needed in a SINGLE call.
-       - If the first zoom is insufficient, call again with additional tiles.
-       - Call the zoom until you have found all relevant tiles, refer to the tile metadata to assist in your search.
-    
-    3) VERIFY
-       - Read the zoomed image carefully.
-       - Extract exact values, tags, room names, and directional cues.
-    
-    4) REPORT
-       - Return the Findings Packet in strict JSON format.
-    
-    ========================
-    WHAT COUNTS AS PROOF
-    ========================
-    - Dimension values (e.g., “36\"”, “1 HR RATED”)
-    - Explicit labels (e.g., “EXIT”, “STAIR A”, “R-2”, “COLUMN C3”)
-    - Symbol legends that define a mark
-    - Path continuity that can be visually traced across tiles
-    - OCR text snippets
-    
-    ========================
-    FINDINGS RULES
-    ========================
-    - Every bullet in `findings` MUST cite either:
-      - `[Tile <ID>]` or
-      - `"Quoted text"`
-    - If a claim cannot be verified from the zoomed tiles or text, mark it as **Unverified**.
-    - Be comprehensive in this report, your supervisor only has access to the report you give in findings, not the full page text or other image data you have.
-    - Do NOT repeat planner instructions — only report what you observe.
-    
-    ========================
-    VISUAL POINTERS RULES
-    ========================
-    - Exclude orientation-only or whitespace tiles.
-    - Include ALL tiles needed to re-trace a path or confirm a relationship.
-    - **Your superviser will ONLY see the tiles that you reference here, be comprehensive when returning these tiles.**
-    
-    ========================
-    FULL PAGE USEFULNESS
-    ========================
-    Set `true` ONLY if the finding requires spatial context across the entire page, or if your zoom is missing information.
-    (e.g., tracing egress path, riser continuity, system routing).
-    Otherwise set `false`.
-    
-    ========================
-    JSON FORMAT (STRICT)
-    ========================
-    {
-      "findings": "<markdown string with bullet points and citations>",
-      "visual_pointers": [list of <int>],
-      "textual_evidence": ["<exact quotes from PAGE TEXT>"],
-      "full_page_usefulness": <true|false>,
-      "limitations": "<what could not be verified and why>"
-    }
-    
-    ========================
-    FAILURE CONDITIONS
-    ========================
-    - If no relevant evidence exists on this page, return:
-      {
-        "findings": "No relevant technical evidence found for the planner's instruction.",
-        "visual_pointers": [],
-        "textual_evidence": [],
-        "full_page_usefulness": false,
-        "limitations": "This page does not contain the requested information or it is not legible at available resolution."
-      }
-    
-    Return ONLY valid JSON.
-    """
-
-    messages = [
-        {"role": "system", "content": system_prompt},
-        {
-            "role": "user",
-            "content": [
-                {"type": "text", "text": f"Planner Instruction:\n{expert_instructions}"},
-                {"type": "text", "text": f"Page Context:\n{page_text}"},
-                {"type": "text", "text": f"Available Grid Map:\n{json.dumps(relevant_tile_meta)}"},
-                {
-                    "type": "image_url",
-                    "image_url": {
-                        "url": f"data:image/png;base64,{b64_full_page}"
-                    }
-                }
-            ]
-        }
-    ]
-
-    MAX_TURNS = 3
-
-    for turn in range(MAX_TURNS):
-        response = client.chat.completions.create(
-            model="gpt-4o",
-            messages=messages,
-            tools=tools,
-            tool_choice="auto"
-        )
-
-        msg = response.choices[0].message
-        messages.append(msg)
-
-        if msg.content:
-            try:
-                res = extract_json(msg.content)
-                
-                
-                state.add_analysis(
-                    f"🟨 Page Analyst\n{res.get('findings','')}"
-                )
-                raw_pointers = res.get("visual_pointers", [])
-                tile_idxs = sanitize_tile_indices(raw_pointers)
-                
-                
-                if tile_idxs and tile_idxs != '[]':
-                    stitched_bytes = merge_tiles(
-                        tile_indexes=tile_idxs,
-                        page_num=page_num
-                    )
-    
-                    state.add_log(f'📸 Staging {len(tile_idxs)} tiles for final audit...')
-                    
-                    # Store these to use AFTER the chat finishes
-                    state.add_staged_image_part(
-                        types.Part.from_bytes(
-                            data=stitched_bytes,  # <-- 'data=' is required here
-                            mime_type="image/png"
-                        )
-                    )
-                    
-
-                    stitched_img = Image.open(
-                        io.BytesIO(stitched_bytes)
-                    )
-                    state.add_image(stitched_img)
-                    
-                    
-                state.add_staged_image_part(
-                    types.Part.from_bytes(
-                        data=image_bytes_list[page_num],  # <-- 'data=' is required here
-                        mime_type="image/png"
-                    )
-                )
-
-                return res
-            except:
-                pass
-            
-        if msg.tool_calls:
-            tool_results = []
-            image_blocks = []
-        
-            for call in msg.tool_calls:
-                if call.function.name == "merge_tiles":
-                    args = json.loads(call.function.arguments)
-                    idxs = args["tile_indexes"]
-        
-                    stitched_bytes = merge_tiles(
-                        tile_indexes=idxs,
-                        page_num=page_num
-                    )
-        
-                    b64_tile = base64.b64encode(stitched_bytes).decode()
-        
-                    tool_results.append({
-                        "role": "tool",
-                        "tool_call_id": call.id,
-                        "content": json.dumps({
-                            "status": "success",
-                            "tiles": idxs
-                        })
-                    })
-        
-                    image_blocks.append(
-                        {
-                            "type": "image_url",
-                            "image_url": {
-                                "url": f"data:image/png;base64,{b64_tile}"
-                            }
-                        }
-                    )
-        
-            for tool_msg in tool_results:
-                messages.append(tool_msg)
-        
-            messages.append({
-                "role": "user",
-                "content": [
-                    {
-                        "type": "text",
-                        "text": "Here are the high-resolution zooms you requested. Analyze exits, locations, and any capacity labels."
-                    },
-                    *image_blocks
-                ]
-            })
-        
-            continue
-
-        messages.append({
-            "role": "user",
-            "content": "Return the FINAL JSON now."
-        })
-
-    raise RuntimeError("No FINAL JSON output from Page Expert")
-
-# Set up Gemini planner
-tools_list = [search_page_text, nyc_legal_sub_agent, execute_page_expert]
-import time
-planner = genai.Client()
-planner_model = "gemini-3-flash-preview"
-planner_prompt = f"""
-    You are the Lead Architectural Compliance Planner for NYC Building Code and Zoning review.
-    
-    Your role is to coordinate specialist sub-agents and deliver a **proof-carrying compliance verdict**
-    based ONLY on:
-    - OCR-extracted drawing text
-    - High-resolution visual evidence (tile zooms)
-    - Official NYC Code citations
-    
-    You must NOT speculate or rely on architectural norms.
-    
-    ========================
-    DRAWING INDEX (Page Metadata)
-    ========================
-    Use this index to select pages for visual inspection.
-    Avoid irrelevant sheets (e.g., Site, Civil, Utility, Stormwater) unless zoning or site compliance is explicitly required.
-    {json.dumps(page_metadata)}
-    
-    ========================
-    SPECIALIST SUB-AGENTS
-    ========================
-    None of these agents have access to your chat history or internal thought process.
-    They know only how to access information (text, images or code) and what information you give them in the research goal.
-    If they need more context or specific instructions YOU MUST PROVIDE IT WHEN CALLING THEM in the research goal.
-    
-    1) `search_page_text`
-       Purpose: FAST signal extractor.
-       Use to identify code-triggering facts:
-       - Occupancy classification
-       - Building height / stories / high-rise
-       - Construction type
-       - Scope of work (new, alteration, addition, change of occupancy)
-       - Fire protection systems
-       Output is used ONLY to constrain legal research.
-    
-    2) `nyc_legal_sub_agent`
-       Purpose: Definitive legal authority.
-       Use to retrieve governing NYC Code sections, exceptions, and cross-references.
-       Always pass a focused topic derived from Phase 1 signals.
-    
-    3) `execute_page_expert`
-       Purpose: High-resolution visual verification.
-       Use to confirm compliance or non-compliance by zooming tiles.
-       This agent provides the ONLY acceptable visual proof.
-    
-    ========================
-    MANDATORY PHASED WORKFLOW
-    ========================
-    PHASE 1 — SIGNAL EXTRACTION
-    - Use `search_page_text` on candidate pages to determine:
-      occupancy, height, construction type, system presence, and scope.
-    - If signals are missing or ambiguous, expand to additional pages.
-    - Do NOT proceed until you have enough facts to define legal applicability.
-    
-    PHASE 2 — LEGAL SCOPING
-    - Convert Phase 1 signals into a focused legal topic.
-    - Call `nyc_legal_sub_agent`.
-    - Extract governing sections, exceptions, and edge cases.
-    
-    PHASE 3 — VISUAL VERIFICATION
-    - Identify the SINGLE most relevant page for proof.
-    - Call `execute_page_expert` with precise instructions tied to legal requirements
-      (e.g., “Verify exit door clear width at Stair A serving R-2 occupancy”).
-    - Ensure returned findings include tile IDs and/or text quotes.
-    
-    PHASE 4 — SYNTHESIS & VERDICT
-    - Compare visual findings directly against legal requirements.
-    - Resolve conflicts:
-      - If legal text and visual evidence disagree → flag as **Non-Compliant or Ambiguous**
-      - If evidence is missing → flag as **Unverified**
-    - Cite both:
-      - NYC Code Section(s)
-      - Tile ID(s) or OCR quotes
-    
-    ========================
-    FINAL OUTPUT FORMAT (STRICT MARKDOWN)
-    ========================
-    ### Compliance Verdict
-    **Status:** Compliant | Non-Compliant | Unverified | Ambiguous
-    
-    ### Legal Basis
-    - **[Code Type] §[Section] — [Title]**
-      - Requirement:
-      - Exceptions Considered:
-    
-    ### Visual Evidence
-    - Finding: <short statement>
-      - Proof: [Tile ID(s)] or "Quoted OCR Text"
-    
-    ### Reasoning
-    - Step-by-step comparison between legal requirement and observed condition
-    
-    ### Limitations
-    - What could not be verified and why
-    
-    ========================
-    CONTROL RULES
-    ========================
-    - NEVER call `nyc_legal_sub_agent` before `search_page_text`
-    - NEVER issue a final verdict without calling `execute_page_expert`
-    - If no page contains sufficient proof, return **Unverified**
-    - Prefer false negatives over false positives
-    *** CRITICAL VISUAL PROTOCOL ***
-    - When `execute_page_expert` returns, it will explicitly state "VISUAL_PROOF_PENDING".
-    - When you see this, your ONLY response must be: "Awaiting visual proof."
-    - DO NOT attempt to guess the verdict.
-    - DO NOT complain about missing images.
-    - Simply wait. The user will immediately send the images in the next turn.
-    
-    
-    ========================
-    QUALITY STANDARD
-    ========================
-    This output should be defensible to a DOB plan examiner or legal reviewer.
-    Every claim must be traceable to law and evidence.
-    """
-
-config = types.GenerateContentConfig(
-    system_instruction=planner_prompt,
-    tools=tools_list
-)
-
-chat = planner.chats.create(model=planner_model, config=config)
-
-
-def agent_worker(user_question):
-    state.clear()
-    state.add_log(f'🚀 Starting analysis for: **{user_question}**')
-    state.add_analysis("🧠 Planner initialized. Awaiting tool calls...")
-
-    # 1. Initialize the Stateful Chat
-    chat = planner.chats.create(model=planner_model, config=config)
-    response = chat.send_message(user_question)
-    
-    # 2. Track images throughout the conversation
-
-    # 3. Standard Tool Loop (Phases 1-3)
-    while response.candidates[0].content.parts[0].function_call:
-        tool_responses = []
-
-        for part in response.candidates[0].content.parts:
-            if part.function_call:
-                name = part.function_call.name
-                args = part.function_call.args
-                state.add_log(f'🛠️ Tool Call: **{name}**')
-
-                func = globals()[name]
-                result = func(**args)
-
-                tool_responses.append(
-                    types.Part.from_function_response(name=name, response={"result": result})
-                )
-
-        # Send tool results back to the stateful chat
-        response = chat.send_message(tool_responses)
-
-    # -----------------------------------------------------------------
-    # PHASE 4: THE POST-CHAT HANDOFF (The "Visual Audit")
-    # -----------------------------------------------------------------
-    
-    # At this point, the while loop has ended. 
-    # 'response.text' contains the model's preliminary answer.
-    
-    audit_images = state.get_staged_images() 
-
-    if audit_images:
-        state.add_log(f"👁️ Preliminary answer received. Performing audit with {len(audit_images)} images...")
-
-        # 1. Construct the audit parts
-        # Ensure 'text=' is used for the Part constructor
-        audit_parts = [
-            types.Part.from_text(
-                text="You have provided a preliminary verdict. Now, look at these images "
-                     "to verify your findings. If the visual evidence contradicts your "
-                     "text-based search, update your verdict now. "
-                ),
-            *audit_images
-        ]
-    
-        try:
-            # 2. Send directly through the 'chat' session
-            # This automatically appends to history and maintains the session state
-            final_response = chat.send_message(audit_parts)
-            
-            state.final_answer = final_response.text
-    
-        except Exception as e:
-            # If the above fails, try the explicit message keyword
-            state.add_log("🔄 Retrying audit with explicit message keyword...")
-            final_response = chat.send_message(message=audit_parts)
-            state.final_answer = final_response.text
-            
-    else:
-        state.add_log("⚠️ No images found in state. Skipping visual audit.")
-        state.final_answer = response.text
-    
-    state.add_log('🏁 **ANALYSIS COMPLETE**')
-    state.done = True
-
-    
-def run_agentic_workflow(user_question):
-    state.done = False
-    state.final_answer = ""
-
-    thread = threading.Thread(
-        target=agent_worker,
-        args=(user_question,),
-        daemon=True
-    )
-    thread.start()
-
-    while not state.done:
-        with state.lock:
-            logs = "\n\n".join(state.log_messages)
-            analysis = "\n\n".join(state.analysis_messages)
-            chapter = state.current_chapter
-            images = list(state.current_images)
-
-        yield (
-            logs,
-            analysis,
-            chapter,
-            images,
-            "*Analysis in progress...*"
-        )
-        time.sleep(0.25)
-
-    with state.lock:
-        logs = "\n\n".join(state.log_messages)
-        analysis = "\n\n".join(state.analysis_messages)
-        chapter = state.current_chapter
-        images = list(state.current_images)
-        final = state.final_answer
-
-    yield (
-        logs,
-        analysis,
-        chapter,
-        images,
-        final
-    )
-
-
-# Build Gradio Interface
-with gr.Blocks(title="AEC Compliance Agent") as demo:
-    gr.Markdown("# 🏗️ AEC Compliance Analysis Agent")
-    gr.Markdown("Ask questions about NYC Building Code compliance for your construction drawings.")
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            question_input = gr.Textbox(
-                label="Your Question",
-                placeholder="e.g., Does this building comply with egress requirements for 738 occupants?",
-                lines=3
-            )
-            submit_btn = gr.Button("🔍 Analyze", variant="primary", size="lg")
-            
-            gr.Markdown("### 📋 Analysis Log")
-            log_output = gr.Markdown(value="", height=400)
-            
-        with gr.Column(scale=1):
-            gr.Markdown("### 🧠 Sub-Agent Analysis")
-            analysis_output = gr.Markdown(value="", height=600)
-            
-        with gr.Column(scale=1):
-            gr.Markdown("### 📖 Code Chapter")
-            chapter_output = gr.Markdown(value="*No chapter loaded yet*", height=600)
-    
-    with gr.Row():
-        gr.Markdown("### 🖼️ Retrieved Images")
-    
-    with gr.Row():
-        image_gallery = gr.Gallery(
-            label="Visual Evidence",
-            show_label=True,
-            columns=2,
-            height=400,
-            object_fit="contain"
-        )
-    
-    with gr.Row():
-        gr.Markdown("### ✅ Final Compliance Verdict")
-    
-    with gr.Row():
-        final_output = gr.Markdown(value="*Analysis pending...*")
-    
-    submit_btn.click(
-            fn=run_agentic_workflow,
-            inputs=[question_input],
-            outputs=[
-                log_output,
-                analysis_output,   # NEW SLOT
-                chapter_output,
-                image_gallery,
-                final_output
-            ]
-        )
-
-if __name__ == "__main__":
-    demo.queue().launch(
-        inbrowser=True,
-        auth=("username", os.getenv("PASSWORD"))  # only share the password
+import gradio as gr
+import json
+import os
+import re
+from google import genai
+from google.genai import types
+import chromadb
+from chromadb.utils import embedding_functions
+from collections import Counter
+import base64
+import io
+from PIL import Image
+import matplotlib.pyplot as plt
+import openai
+from datetime import datetime
+import threading
+
+os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
+os.environ["GEMINI_API_KEY"] = os.getenv("GEMINI_API_KEY")
+
+# Global state for the interface
+class InterfaceState:
+    def __init__(self):
+        self.log_messages = []
+        self.analysis_messages = []
+        self.current_chapter = ""
+        self.current_images = []
+        self.staged_audit_images = []
+        self.final_answer = ""
+        self.done = False
+        self.lock = threading.Lock()
+    
+    def add_log(self, message):
+        timestamp = datetime.now().strftime("%H:%M:%S")
+        with self.lock:
+            self.log_messages.append(f"**[{timestamp}]** {message}")
+            return "\n\n".join(self.log_messages)
+
+    def add_analysis(self, message):
+        timestamp = datetime.now().strftime("%H:%M:%S")
+        with self.lock:
+            self.analysis_messages.append(f"**[{timestamp}]** {message}")
+            return "\n\n".join(self.analysis_messages)
+
+    def set_chapter(self, chapter_text):
+        with self.lock:
+            self.current_chapter = chapter_text
+            return chapter_text
+
+    def add_image(self, img_pil):
+        with self.lock:
+            self.current_images.append(img_pil)
+            return self.current_images.copy()
+        
+    def add_staged_image_part(self, image_part):
+        """Thread-safe method to stage images for the Gemini Audit."""
+        with self.lock:
+            self.staged_audit_images.append(image_part)
+            # Log it so we can verify it happened in the console
+            print(f"DEBUG: Staged image part. Total staged: {len(self.staged_audit_images)}")
+
+    def get_staged_images(self):
+        """Safely retrieve the staged images for the audit turn."""
+        with self.lock:
+            return list(self.staged_audit_images) # Return a copy to prevent mutation
+    
+    def clear(self):
+        with self.lock:
+            self.log_messages.clear()
+            self.analysis_messages.clear()
+            self.current_chapter = ""
+            self.current_images.clear()
+            self.final_answer = ""
+            self.done = False
+            
+    
+
+state = InterfaceState()
+
+# Load your data (same as original)
+with open('Preprocessed Files/page_metadata.json', 'r') as json_file:
+    page_metadata = json.load(json_file)
+page_metadata = {int(k): v for k, v in page_metadata.items()}
+
+with open('Preprocessed Files/text_list.json', 'r') as json_file:
+    text_list = json.load(json_file)
+    
+with open('Preprocessed Files/tile_metadata.json', 'r') as json_file:
+    tile_metadata = json.load(json_file)
+tile_metadata = {
+    int(outer_k): {
+        int(inner_k): inner_v
+        for inner_k, inner_v in outer_v.items()
+    }
+    for outer_k, outer_v in tile_metadata.items()
+}
+
+def load_fullpage_images(folder="Images"):
+    files = os.listdir(folder)
+    page_files = []
+    for f in files:
+        match = re.search(r"page_(\d+)_fullpage\.png", f)
+        if match:
+            page_num = int(match.group(1))
+            page_files.append((page_num, f))
+    page_files.sort(key=lambda x: x[0])
+    image_bytes_list = []
+    for page_num, filename in page_files:
+        path = os.path.join(folder, filename)
+        with open(path, "rb") as f:
+            img_bytes = f.read()
+        image_bytes_list.append(img_bytes)
+    return image_bytes_list
+
+def load_tile_images(page):
+    files = os.listdir('Tiles')
+    page_files = []
+    for f in files:
+        match = re.search(f"page_{page}_tile_(\d+)\.png", f)
+        if match:
+            page_num = int(match.group(1))
+            page_files.append((page_num, f))
+    page_files.sort(key=lambda x: x[0])
+    image_bytes_list = []
+    for page_num, filename in page_files:
+        path = os.path.join('Tiles', filename)
+        with open(path, "rb") as f:
+            img_bytes = f.read()
+        image_bytes_list.append(img_bytes)
+    return image_bytes_list
+
+image_bytes_list = load_fullpage_images()
+
+tile_bytes = {}
+for page in range(44):
+    tile_list = load_tile_images(page)
+    if tile_list:
+        tile_bytes[page] = load_tile_images(page)
+
+# Vector Code Base
+chroma_client = chromadb.PersistentClient(path="nyc_code_db")
+embedding_model = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
+collection = chroma_client.get_collection(name="nyc_building_codes", embedding_function=embedding_model)
+
+all_pending_images = []
+
+# Modified tool functions with Gradio updates
+def search_page_text(page_number: int, research_goal: str):
+    state.add_log(f'🔍 Searching page **{page_metadata[page_number]["sheet_title"]}** for details')
+    
+    state.add_analysis(
+        f'🔍 Searching page {page_metadata[page_number]["sheet_title"]} with prompt\n{research_goal}'
+    )
+    
+    raw_text = text_list[page_number]
+    
+    client = openai.OpenAI()
+    response = client.chat.completions.create(
+        model="gpt-5-mini", 
+        messages=[
+            {"role": "system", "content": """
+            You are a Fast NYC Plans Examiner Signal Agent.
+            
+            Your ONLY job is to extract **code-relevant signals** from the OCR text of a SINGLE drawing page.
+            You do NOT interpret the law and you do NOT summarize design intent.
+            
+            Your output will be used to CONSTRAIN a downstream legal research agent.
+            
+            ========================
+            WHAT TO EXTRACT
+            ========================
+            Look only for information that determines which parts of the NYC Code apply such as:
+            
+            - Occupancy classification (e.g., R-2, A-3, M, S, F, mixed-use)
+            - Building height (stories, feet, high-rise indicators)
+            - Construction type (I, II, III, IV, V)
+            - Fire protection systems (sprinklers, standpipes, fire alarm, smoke control)
+            - Means of egress references (stairs, exits, exit access, doors, corridors)
+            - Structural system hints (steel, concrete, load-bearing walls, columns, transfer girders)
+            - Mechanical / fuel / plumbing system mentions (boilers, gas piping, HVAC type, shafts)
+            - Zoning or special district references (if present)
+            - Scope flags (new building, alteration, addition, change of occupancy, retrofit)
+            
+            However only return relevant signals to the provided research goal.
+            
+            ========================
+            OUTPUT FORMAT (STRICT MARKDOWN)
+            ========================
+            Return ONLY the following sections:
+            
+            ### Code-Relevant Signals
+            - Bullet list of extracted facts
+            
+            ### Likely Governing Code Domains
+            - One-line list chosen from: Administrative, Building, Mechanical, FuelGas, Plumbing, Fire
+            
+            ### Text Evidence
+            - Short quoted snippets from the page that support each signal
+            
+            ========================
+            RULES
+            ========================
+            - Do NOT speculate
+            - If a signal is not present, omit it
+            - Prefer exact phrases over paraphrase
+            - Keep total length under 500 words
+            - No legal conclusions, no compliance advice
+            """},
+            {"role": "user", "content": f"PAGE TEXT:\n{raw_text}\n\nRESEARCH GOAL: {research_goal}\n\nReturn a breif but comprehensive Markdown summary of your findings and justification with text snippets."}
+        ]
+    )
+    
+    analysis_text = response.choices[0].message.content
+
+    state.add_analysis(
+        f"🟦 Text Analyst (Page {page_number})\n{analysis_text}"
+    )
+    
+    return {
+        "page": page_number,
+        "summary": analysis_text
+    }
+
+def discover_code_locations(query: str):
+    state.add_log(f'📚 Searching NYC Code for: **{query}**')
+    
+    results = collection.query(
+        query_texts=[query],
+        n_results=25,
+        include=["metadatas", "documents"]
+    )
+    
+    if not results['metadatas'][0]:
+        return "No results found. Try a different technical keyword."
+
+    metas = results['metadatas'][0]
+    docs = results['documents'][0]
+    
+    category_chapter_pairs = [f"{m['code_type']} | Ch. {m['parent_major']}" for m in metas]
+    counts = Counter(category_chapter_pairs)
+    chapter_summary = "\n".join([f"- {pair} ({count} hits)" for pair, count in counts.most_common(5)])
+
+    section_reports = []
+    for m, doc in zip(metas, docs):
+        report = (
+            f"ID: {m['section_full']} | Code: {m['code_type']} | Chapter: {m['parent_major']}\n"
+            f"Snippet: {doc}"
+        )
+        section_reports.append(report)
+
+    output = (
+        "### CODE DISCOVERY REPORT ###\n"
+        f"MOST RELEVANT CHAPTERS:\n{chapter_summary}\n\n"
+        "TOP RELEVANT SECTIONS:\n" + 
+        "\n---\n".join(section_reports)
+    )
+    
+    return output
+
+def fetch_full_chapter(code_type: str, chapter_id: str):
+    state.add_log(f'📖 Fetching Chapter **{chapter_id}** from **{code_type}** code')
+    
+    try:
+        chapter_data = collection.get(
+            where={
+                "$and": [
+                    {"code_type": {"$eq": code_type}},
+                    {"parent_major": {"$eq": chapter_id}}
+                ]
+            },
+            include=["documents", "metadatas"]
+        )
+
+        if not chapter_data['documents']:
+            return f"No documentation found for {code_type} Chapter {chapter_id}."
+
+        sections = sorted(zip(chapter_data['metadatas'], chapter_data['documents']), 
+                          key=lambda x: x[0]['section_full'])
+
+        full_text = f"## FULL LEGAL TEXT: {code_type.upper()} CODE - CHAPTER {chapter_id}\n\n"
+
+        for meta, doc in sections:
+            blocks = doc.split("[CONT.]:")
+            unique_blocks = []
+            for b in blocks:
+                clean_b = b.strip()
+                if clean_b and clean_b not in unique_blocks:
+                    unique_blocks.append(clean_b)
+            
+            clean_doc = " ".join(unique_blocks)
+            full_text += f"### SECTION {meta['section_full']}\n{clean_doc}\n\n---\n\n"
+
+        # Update the chapter display
+        state.set_chapter(full_text)
+        
+        return full_text
+
+    except Exception as e:
+        return f"Error retrieving chapter content: {str(e)}"
+
+def nyc_legal_sub_agent(research_goal: str):
+    state.add_log(f'⚖️ Investigating NYC Code for: **{research_goal}**')
+    
+    state.add_analysis(
+        f"⚖️ Legal Analyst is searching\n{research_goal}"
+    )
+    
+    client = openai.OpenAI()
+    
+    internal_tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "discover_code_locations",
+                "description": "Scans NYC code in a semantic vector database. Use this FIRST to find which chapters/sections are relevant.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "query": {"type": "string", "description": "semantic search string for a vector database (Not a keyword search use a full sentence)"}
+                    },
+                    "required": ["query"]
+                }
+            }
+        },
+        {
+            "type": "function",
+            "function": {
+                "name": "fetch_full_chapter",
+                "description": "Retrieves the full legal text of a specific chapter for deep analysis.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "code_type": {
+                            "type": "string", 
+                            "enum": ["Administrative", "Building", "FuelGas", "Mechanical", "Plumbing"],
+                            "description": "The specific NYC code volume to search."
+                        },
+                        "chapter_id": {"type": "string", "description": "The chapter number string"}
+                    },
+                    "required": ["code_type", "chapter_id"]
+                }
+            }
+        }
+    ]
+
+    messages = [
+        {"role": "system", "content": """
+        You are a Senior NYC Building Code Consultant and Legal Research Agent.
+        
+        Your task is to produce a **definitive, citation-backed legal report** that can be used directly by a downstream orchestration agent. 
+        Accuracy, traceability, and completeness matter more than brevity.
+        
+        ========================
+        PRIMARY OBJECTIVE
+        ========================
+        Given a research goal, identify and analyze relevant NYC Code provisions, including:
+        - Governing sections
+        - Exceptions
+        - Cross-references
+        - Related chapters that modify, limit, or expand the rule
+        
+        Every legal claim MUST be supported by a specific code citation.
+        
+        You are operating in FAST LEGAL MODE.
+
+        SEARCH BUDGET:
+        - Maximum of 2 calls to `discover_code_locations`
+        - Maximum of 2 calls to `fetch_full_chapter`
+        
+        STOP CONDITIONS:
+        - If the first chapter fetch contains governing text AND exceptions, STOP and synthesize.
+        - Only fetch a second chapter if the first chapter explicitly cross-references another chapter.
+        
+        PRIORITY ORDER:
+        1) Governing rule section
+        2) Exceptions
+        3) Cross-references that MODIFY the rule
+        Ignore definitions and administrative content unless directly referenced.
+        
+        GOOD ENOUGH STANDARD:
+        If you can identify:
+        - The governing section
+        - At least one exception or limitation
+        You must STOP and report.
+        
+        ========================
+        TOOL STRATEGY (MANDATORY)
+        ========================
+        This is a semantic vector database, NOT a keyword index. Always search in full English questions.
+        
+        1) FIRST — Call `discover_code_locations`
+           - Use a natural-language query describing the legal requirement you are trying to find
+           - Example: "What NYC Building Code sections regulate emergency egress width in residential buildings"
+           NEVER use a keyword search thi will not work you are searching a vector database. 
+           If you know what chaoter you need call the fetch_full_chapter tool instead.
+           If you perform TWO consecutive `discover_code_locations` calls
+            and both return no new relevant chapters or sections:
+            
+            You MUST stop searching and do one of the following:
+            - Conclude that the table/section does NOT exist as a standalone provision in the NYC Code corpus, OR
+            - Conclude that the requirement is embedded within the previously retrieved sections
+            
+            Then proceed to report findings using the closest governing section.
+            
+            DO NOT continue reformulating the same query.
+            You MUST NOT call `discover_code_locations` more than once for the same legal concept.
+            If a new query is semantically similar to a prior query, STOP and move forward with analysis.
+            
+        2) SECOND — Call `fetch_full_chapter`
+           - If multiple relevant sections appear in the same chapter
+           - OR if a section contains exceptions, references, or conditional language
+           - OR if you know what section of the code is relevant and want to see a full chapter
+        
+        3) THIRD — Follow Cross-References
+           - If a section says "See Section X", "As required by Chapter Y", or "Except as permitted in..."
+           - You MUST search and retrieve those sections as well
+        
+        4) STOP ONLY WHEN
+           - All exceptions are reviewed
+           - All cross-references are resolved
+           - No additional modifying sections remain
+        
+        ========================
+        OUTPUT FORMAT (STRICT)
+        ========================
+        Return a structured legal report in the following format:
+        
+        ### Legal Summary
+        Brief, plain-language explanation of what the code requires.
+        
+        ### Governing Code Sections
+        - **[Code Type] §[Section Number] — [Title]**
+          - Summary:
+          - Key Requirements:
+          - Applicability Conditions:
+          - Exceptions:
+        
+        ### Cross-References Analyzed
+        - **§[Section Number] — [Title]**
+          - Why It Matters:
+          - Impact on Main Rule:
+        
+        ### Edge Cases & Enforcement Notes
+        - Special conditions (building type, occupancy class, height, system type, jurisdictional notes)
+        - Common misinterpretations
+        - DOB or FDNY enforcement implications (if relevant)
+        
+        ### Compliance Checklist
+        - Bullet list of actionable compliance steps derived from the code
+        
+        ========================
+        QUALITY RULES
+        ========================
+        - NEVER summarize without citing
+        - NEVER assume jurisdiction, building type, or occupancy unless the code explicitly states it
+        - If legal text is ambiguous, flag it as **Interpretive**
+        - Prefer quoting short legal phrases when clarity matters
+        
+        ========================
+        TONE
+        ========================
+        Professional. Precise. Legal-research quality. No speculation.
+        """},
+        {"role": "user", "content": f"Analyze the NYC building code with this goal: {research_goal}"}
+    ]
+
+    for _ in range(20):
+        response = client.chat.completions.create(
+            model="gpt-5-mini",
+            messages=messages,
+            tools=internal_tools,
+            tool_choice="auto"
+        )
+        
+        msg = response.choices[0].message
+        messages.append(msg)
+
+        if not msg.tool_calls:
+            
+            break
+
+        for tool_call in msg.tool_calls:
+            func_name = tool_call.function.name
+            args = json.loads(tool_call.function.arguments)
+            
+            if func_name == "discover_code_locations":
+                result = discover_code_locations(args['query'])
+            elif func_name == "fetch_full_chapter":
+                result = fetch_full_chapter(args['code_type'], args['chapter_id'])
+            
+            messages.append({
+                "role": "tool",
+                "tool_call_id": tool_call.id,
+                "content": result
+            })
+    
+    state.add_analysis(
+        f"🟨 Legal Analyst\n{msg.content}"
+    )
+    
+    return msg.content
+
+def merge_tiles(tile_indexes: list[int], page_num: int):
+    state.add_log(f'🔬 Stitching tiles **{tile_indexes}** from page **{page_num}**')
+    
+    images = []
+    positions = []
+    
+    tiles = tile_bytes[page_num]
+    tiles_coords_dict = tile_metadata[page_num]
+
+    for index in tile_indexes:
+        if index < 0 or index >= len(tiles):
+            raise ValueError(f"Tile index {index} out of range")
+
+        img_bytes = tiles[index]
+        if img_bytes is None:
+            raise ValueError(f"No image bytes found for tile {index}")
+
+        img = Image.open(io.BytesIO(img_bytes)).convert('RGBA')
+        images.append(img)
+
+        x = tiles_coords_dict[index]['coords'][0]
+        y = tiles_coords_dict[index]['coords'][1]
+        positions.append((x, y))
+
+    if not images:
+        return None
+
+    min_x = min(x for x, y in positions)
+    min_y = min(y for x, y in positions)
+    normalized_positions = [(x - min_x, y - min_y) for x, y in positions]
+
+    total_width = max(pos[0] + img.width for pos, img in zip(normalized_positions, images))
+    total_height = max(pos[1] + img.height for pos, img in zip(normalized_positions, images))
+
+    stitched_image = Image.new('RGB', (total_width, total_height), (255, 255, 255))
+
+    for img, pos in zip(images, normalized_positions):
+        stitched_image.paste(img, pos)
+    
+    # Add to image gallery
+    
+    output_buffer = io.BytesIO()
+    stitched_image.save(output_buffer, format='PNG')
+    stitched_bytes = output_buffer.getvalue()
+
+    return stitched_bytes
+
+def extract_json(s: str):
+    s = s.strip()
+    start = s.find("{")
+    end = s.rfind("}")
+    if start == -1 or end == -1 or end < start:
+        raise ValueError("No JSON object found in model output:\n" + repr(s))
+    json_str = s[start:end+1]
+    return json.loads(json_str)
+
+def sanitize_tile_indices(data):
+    """
+    Forcefully converts various LLM outputs into a clean list of integers.
+    Handles: [1, 2], ["1", "2"], "1, 2, 3", "[1, 2, 3]", and None.
+    """
+    if not data:
+        return []
+    
+    # If it's already a list, ensure all elements are integers
+    if isinstance(data, list):
+        clean_list = []
+        for item in data:
+            try:
+                # This handles strings inside the list like ["1", "2"]
+                clean_list.append(int(str(item).strip()))
+            except (ValueError, TypeError):
+                continue
+        return clean_list
+
+    # If it's a string, use Regex to find all sequences of digits
+    if isinstance(data, str):
+        # findall returns all non-overlapping matches of the pattern
+        numbers = re.findall(r'\d+', data)
+        return [int(n) for n in numbers]
+
+    return []
+
+def execute_page_expert(expert_instructions: str, page_num: int):
+    state.add_log(f'👁️ Spawning Page Expert for page **{page_num}**')
+    state.add_analysis(f"👁️ Page Expert searching for {expert_instructions}")
+    state.add_log(f'📄 Attaching full-page context for page **{page_num}**')
+    state.add_analysis(
+        f"📄 Full-page context attached for page `{page_num}`"
+    )
+
+    full_page_img = Image.open(
+        io.BytesIO(image_bytes_list[page_num])
+    )
+    state.add_image(full_page_img)
+    
+    client = openai.OpenAI()
+
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "merge_tiles",
+                "description": "Stitches high-resolution image tiles together into a single zoomed-in view. Use this to read small text, dimensions, or symbols.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "tile_indexes": {
+                            "type": "array",
+                            "items": {"type": "integer"},
+                            "description": "A list of integer tile IDs from the Grid Map to stitch together."
+                        }
+                    },
+                    "required": ["tile_indexes"]
+                }
+            }
+        }
+    ]
+
+    page_text = text_list[page_num]
+    relevant_tile_meta = tile_metadata[page_num]
+    b64_full_page = base64.b64encode(image_bytes_list[page_num]).decode()
+
+    system_prompt = """
+    You are a Lead AEC Visual Investigator supporting a Compliance Planner.
+    
+    Your mission is to extract **verifiable, high-fidelity evidence** from this drawing page.
+    You must ground every claim in either:
+    - a **Zoomed Tile Image** (via `merge_tiles`) or
+    - a **Direct Text Quote** from the OCR page text.
+    
+    Guesses, assumptions, and general descriptions are not allowed.
+    
+    ========================
+    MANDATORY WORKFLOW
+    ========================
+    1) ORIENT
+       - Review the full-page image and the Grid Map to identify candidate regions.
+       - Decide which tiles likely contain the required evidence. Utilize the tile metadata to assist with this tasl.
+    
+    2) ZOOM (REQUIRED)
+       - You MUST call `merge_tiles(tile_indexes=[...])` before making ANY factual claim about symbols, dimensions, labels, or locations.
+       - Always request ALL tiles needed in a SINGLE call.
+       - If the first zoom is insufficient, call again with additional tiles.
+       - Call the zoom until you have found all relevant tiles, refer to the tile metadata to assist in your search.
+    
+    3) VERIFY
+       - Read the zoomed image carefully.
+       - Extract exact values, tags, room names, and directional cues.
+    
+    4) REPORT
+       - Return the Findings Packet in strict JSON format.
+    
+    ========================
+    WHAT COUNTS AS PROOF
+    ========================
+    - Dimension values (e.g., “36\"”, “1 HR RATED”)
+    - Explicit labels (e.g., “EXIT”, “STAIR A”, “R-2”, “COLUMN C3”)
+    - Symbol legends that define a mark
+    - Path continuity that can be visually traced across tiles
+    - OCR text snippets
+    
+    ========================
+    FINDINGS RULES
+    ========================
+    - Every bullet in `findings` MUST cite either:
+      - `[Tile <ID>]` or
+      - `"Quoted text"`
+    - If a claim cannot be verified from the zoomed tiles or text, mark it as **Unverified**.
+    - Be comprehensive in this report, your supervisor only has access to the report you give in findings, not the full page text or other image data you have.
+    - Do NOT repeat planner instructions — only report what you observe.
+    
+    ========================
+    VISUAL POINTERS RULES
+    ========================
+    - Exclude orientation-only or whitespace tiles.
+    - Include ALL tiles needed to re-trace a path or confirm a relationship.
+    - **Your superviser will ONLY see the tiles that you reference here, be comprehensive when returning these tiles.**
+    
+    ========================
+    FULL PAGE USEFULNESS
+    ========================
+    Set `true` ONLY if the finding requires spatial context across the entire page, or if your zoom is missing information.
+    (e.g., tracing egress path, riser continuity, system routing).
+    Otherwise set `false`.
+    
+    ========================
+    JSON FORMAT (STRICT)
+    ========================
+    {
+      "findings": "<markdown string with bullet points and citations>",
+      "visual_pointers": [list of <int>],
+      "textual_evidence": ["<exact quotes from PAGE TEXT>"],
+      "full_page_usefulness": <true|false>,
+      "limitations": "<what could not be verified and why>"
+    }
+    
+    ========================
+    FAILURE CONDITIONS
+    ========================
+    - If no relevant evidence exists on this page, return:
+      {
+        "findings": "No relevant technical evidence found for the planner's instruction.",
+        "visual_pointers": [],
+        "textual_evidence": [],
+        "full_page_usefulness": false,
+        "limitations": "This page does not contain the requested information or it is not legible at available resolution."
+      }
+    
+    Return ONLY valid JSON.
+    """
+
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": f"Planner Instruction:\n{expert_instructions}"},
+                {"type": "text", "text": f"Page Context:\n{page_text}"},
+                {"type": "text", "text": f"Available Grid Map:\n{json.dumps(relevant_tile_meta)}"},
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/png;base64,{b64_full_page}"
+                    }
+                }
+            ]
+        }
+    ]
+
+    MAX_TURNS = 3
+
+    for turn in range(MAX_TURNS):
+        response = client.chat.completions.create(
+            model="gpt-4o",
+            messages=messages,
+            tools=tools,
+            tool_choice="auto"
+        )
+
+        msg = response.choices[0].message
+        messages.append(msg)
+
+        if msg.content:
+            try:
+                res = extract_json(msg.content)
+                
+                
+                state.add_analysis(
+                    f"🟨 Page Analyst\n{res.get('findings','')}"
+                )
+                raw_pointers = res.get("visual_pointers", [])
+                tile_idxs = sanitize_tile_indices(raw_pointers)
+                
+                
+                if tile_idxs and tile_idxs != '[]':
+                    stitched_bytes = merge_tiles(
+                        tile_indexes=tile_idxs,
+                        page_num=page_num
+                    )
+    
+                    state.add_log(f'📸 Staging {len(tile_idxs)} tiles for final audit...')
+                    
+                    # Store these to use AFTER the chat finishes
+                    state.add_staged_image_part(
+                        types.Part.from_bytes(
+                            data=stitched_bytes,  # <-- 'data=' is required here
+                            mime_type="image/png"
+                        )
+                    )
+                    
+
+                    stitched_img = Image.open(
+                        io.BytesIO(stitched_bytes)
+                    )
+                    state.add_image(stitched_img)
+                    
+                    
+                state.add_staged_image_part(
+                    types.Part.from_bytes(
+                        data=image_bytes_list[page_num],  # <-- 'data=' is required here
+                        mime_type="image/png"
+                    )
+                )
+
+                return res
+            except:
+                pass
+            
+        if msg.tool_calls:
+            tool_results = []
+            image_blocks = []
+        
+            for call in msg.tool_calls:
+                if call.function.name == "merge_tiles":
+                    args = json.loads(call.function.arguments)
+                    idxs = args["tile_indexes"]
+        
+                    stitched_bytes = merge_tiles(
+                        tile_indexes=idxs,
+                        page_num=page_num
+                    )
+        
+                    b64_tile = base64.b64encode(stitched_bytes).decode()
+        
+                    tool_results.append({
+                        "role": "tool",
+                        "tool_call_id": call.id,
+                        "content": json.dumps({
+                            "status": "success",
+                            "tiles": idxs
+                        })
+                    })
+        
+                    image_blocks.append(
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/png;base64,{b64_tile}"
+                            }
+                        }
+                    )
+        
+            for tool_msg in tool_results:
+                messages.append(tool_msg)
+        
+            messages.append({
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "Here are the high-resolution zooms you requested. Analyze exits, locations, and any capacity labels."
+                    },
+                    *image_blocks
+                ]
+            })
+        
+            continue
+
+        messages.append({
+            "role": "user",
+            "content": "Return the FINAL JSON now."
+        })
+
+    raise RuntimeError("No FINAL JSON output from Page Expert")
+
+# Set up Gemini planner
+tools_list = [search_page_text, nyc_legal_sub_agent, execute_page_expert]
+import time
+planner = genai.Client()
+planner_model = "gemini-3-flash-preview"
+planner_prompt = f"""
+    You are the Lead Architectural Compliance Planner for NYC Building Code and Zoning review.
+    
+    Your role is to coordinate specialist sub-agents and deliver a **proof-carrying compliance verdict**
+    based ONLY on:
+    - OCR-extracted drawing text
+    - High-resolution visual evidence (tile zooms)
+    - Official NYC Code citations
+    
+    You must NOT speculate or rely on architectural norms.
+    
+    ========================
+    DRAWING INDEX (Page Metadata)
+    ========================
+    Use this index to select pages for visual inspection.
+    Avoid irrelevant sheets (e.g., Site, Civil, Utility, Stormwater) unless zoning or site compliance is explicitly required.
+    {json.dumps(page_metadata)}
+    
+    ========================
+    SPECIALIST SUB-AGENTS
+    ========================
+    None of these agents have access to your chat history or internal thought process.
+    They know only how to access information (text, images or code) and what information you give them in the research goal.
+    If they need more context or specific instructions YOU MUST PROVIDE IT WHEN CALLING THEM in the research goal.
+    
+    1) `search_page_text`
+       Purpose: FAST signal extractor.
+       Use to identify code-triggering facts:
+       - Occupancy classification
+       - Building height / stories / high-rise
+       - Construction type
+       - Scope of work (new, alteration, addition, change of occupancy)
+       - Fire protection systems
+       Output is used ONLY to constrain legal research.
+    
+    2) `nyc_legal_sub_agent`
+       Purpose: Definitive legal authority.
+       Use to retrieve governing NYC Code sections, exceptions, and cross-references.
+       Always pass a focused topic derived from Phase 1 signals.
+    
+    3) `execute_page_expert`
+       Purpose: High-resolution visual verification.
+       Use to confirm compliance or non-compliance by zooming tiles.
+       This agent provides the ONLY acceptable visual proof.
+    
+    ========================
+    MANDATORY PHASED WORKFLOW
+    ========================
+    PHASE 1 — SIGNAL EXTRACTION
+    - Use `search_page_text` on candidate pages to determine:
+      occupancy, height, construction type, system presence, and scope.
+    - If signals are missing or ambiguous, expand to additional pages.
+    - Do NOT proceed until you have enough facts to define legal applicability.
+    
+    PHASE 2 — LEGAL SCOPING
+    - Convert Phase 1 signals into a focused legal topic.
+    - Call `nyc_legal_sub_agent`.
+    - Extract governing sections, exceptions, and edge cases.
+    
+    PHASE 3 — VISUAL VERIFICATION
+    - Identify the SINGLE most relevant page for proof.
+    - Call `execute_page_expert` with precise instructions tied to legal requirements
+      (e.g., “Verify exit door clear width at Stair A serving R-2 occupancy”).
+    - Ensure returned findings include tile IDs and/or text quotes.
+    
+    PHASE 4 — SYNTHESIS & VERDICT
+    - Compare visual findings directly against legal requirements.
+    - Resolve conflicts:
+      - If legal text and visual evidence disagree → flag as **Non-Compliant or Ambiguous**
+      - If evidence is missing → flag as **Unverified**
+    - Cite both:
+      - NYC Code Section(s)
+      - Tile ID(s) or OCR quotes
+
+    **NEVER CALL THE SAME AGENT FOR THE SAME TASK TWICE REFER TO PREVIOUS ANSWERS WHEN ABLE**
+    
+    ========================
+    FINAL OUTPUT FORMAT (STRICT MARKDOWN)
+    ========================
+    ### Compliance Verdict
+    **Status:** Compliant | Non-Compliant | Unverified | Ambiguous
+    
+    ### Legal Basis
+    - **[Code Type] §[Section] — [Title]**
+      - Requirement:
+      - Exceptions Considered:
+    
+    ### Visual Evidence
+    - Finding: <short statement>
+      - Proof: [Tile ID(s)] or "Quoted OCR Text"
+    
+    ### Reasoning
+    - Step-by-step comparison between legal requirement and observed condition
+    
+    ### Limitations
+    - What could not be verified and why
+    
+    ========================
+    CONTROL RULES
+    ========================
+    - NEVER call `nyc_legal_sub_agent` before `search_page_text`
+    - NEVER issue a final verdict without calling `execute_page_expert`
+    - If no page contains sufficient proof, return **Unverified**
+    - Prefer false negatives over false positives
+    *** CRITICAL VISUAL PROTOCOL ***
+    - When `execute_page_expert` returns, it will explicitly state "VISUAL_PROOF_PENDING".
+    - When you see this, your ONLY response must be: "Awaiting visual proof."
+    - DO NOT attempt to guess the verdict.
+    - DO NOT complain about missing images.
+    - Simply wait. The user will immediately send the images in the next turn.
+    
+    
+    ========================
+    QUALITY STANDARD
+    ========================
+    This output should be defensible to a DOB plan examiner or legal reviewer.
+    Every claim must be traceable to law and evidence.
+    """
+
+config = types.GenerateContentConfig(
+    system_instruction=planner_prompt,
+    tools=tools_list
+)
+
+chat = planner.chats.create(model=planner_model, config=config)
+
+
+def agent_worker(user_question):
+    state.clear()
+    state.add_log(f'🚀 Starting analysis for: **{user_question}**')
+    state.add_analysis("🧠 Planner initialized. Awaiting tool calls...")
+
+    # 1. Initialize the Stateful Chat
+    chat = planner.chats.create(model=planner_model, config=config)
+    response = chat.send_message(user_question)
+    
+    # 2. Track images throughout the conversation
+
+    # 3. Standard Tool Loop (Phases 1-3)
+    while response.candidates[0].content.parts[0].function_call:
+        tool_responses = []
+
+        for part in response.candidates[0].content.parts:
+            if part.function_call:
+                name = part.function_call.name
+                args = part.function_call.args
+                state.add_log(f'🛠️ Tool Call: **{name}**')
+
+                func = globals()[name]
+                result = func(**args)
+
+                tool_responses.append(
+                    types.Part.from_function_response(name=name, response={"result": result})
+                )
+
+        # Send tool results back to the stateful chat
+        response = chat.send_message(tool_responses)
+
+    # -----------------------------------------------------------------
+    # PHASE 4: THE POST-CHAT HANDOFF (The "Visual Audit")
+    # -----------------------------------------------------------------
+    
+    # At this point, the while loop has ended. 
+    # 'response.text' contains the model's preliminary answer.
+    
+    audit_images = state.get_staged_images() 
+
+    if audit_images:
+        state.add_log(f"👁️ Preliminary answer received. Performing audit with {len(audit_images)} images...")
+
+        # 1. Construct the audit parts
+        # Ensure 'text=' is used for the Part constructor
+        audit_parts = [
+            types.Part.from_text(
+                text="You have provided a preliminary verdict. Now, look at these images "
+                     "to verify your findings. If the visual evidence contradicts your "
+                     "text-based search, update your verdict now. "
+                ),
+            *audit_images
+        ]
+    
+        try:
+            # 2. Send directly through the 'chat' session
+            # This automatically appends to history and maintains the session state
+            final_response = chat.send_message(audit_parts)
+            
+            state.final_answer = final_response.text
+    
+        except Exception as e:
+            # If the above fails, try the explicit message keyword
+            state.add_log("🔄 Retrying audit with explicit message keyword...")
+            final_response = chat.send_message(message=audit_parts)
+            state.final_answer = final_response.text
+            
+    else:
+        state.add_log("⚠️ No images found in state. Skipping visual audit.")
+        state.final_answer = response.text
+    
+    state.add_log('🏁 **ANALYSIS COMPLETE**')
+    state.done = True
+
+    
+def run_agentic_workflow(user_question):
+    state.done = False
+    state.final_answer = ""
+
+    thread = threading.Thread(
+        target=agent_worker,
+        args=(user_question,),
+        daemon=True
+    )
+    thread.start()
+
+    while not state.done:
+        with state.lock:
+            logs = "\n\n".join(state.log_messages)
+            analysis = "\n\n".join(state.analysis_messages)
+            chapter = state.current_chapter
+            images = list(state.current_images)
+
+        yield (
+            logs,
+            analysis,
+            chapter,
+            images,
+            "*Analysis in progress...*"
+        )
+        time.sleep(0.25)
+
+    with state.lock:
+        logs = "\n\n".join(state.log_messages)
+        analysis = "\n\n".join(state.analysis_messages)
+        chapter = state.current_chapter
+        images = list(state.current_images)
+        final = state.final_answer
+
+    yield (
+        logs,
+        analysis,
+        chapter,
+        images,
+        final
+    )
+
+
+# Build Gradio Interface
+with gr.Blocks(title="AEC Compliance Agent") as demo:
+    gr.Markdown("# 🏗️ AEC Compliance Analysis Agent")
+    gr.Markdown("Ask questions about NYC Building Code compliance for your construction drawings.")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            question_input = gr.Textbox(
+                label="Your Question",
+                placeholder="e.g., Does this building comply with egress requirements for 738 occupants?",
+                lines=3
+            )
+            submit_btn = gr.Button("🔍 Analyze", variant="primary", size="lg")
+            
+            gr.Markdown("### 📋 Analysis Log")
+            log_output = gr.Markdown(value="", height=400)
+            
+        with gr.Column(scale=1):
+            gr.Markdown("### 🧠 Sub-Agent Analysis")
+            analysis_output = gr.Markdown(value="", height=600)
+            
+        with gr.Column(scale=1):
+            gr.Markdown("### 📖 Code Chapter")
+            chapter_output = gr.Markdown(value="*No chapter loaded yet*", height=600)
+    
+    with gr.Row():
+        gr.Markdown("### 🖼️ Retrieved Images")
+    
+    with gr.Row():
+        image_gallery = gr.Gallery(
+            label="Visual Evidence",
+            show_label=True,
+            columns=2,
+            height=400,
+            object_fit="contain"
+        )
+    
+    with gr.Row():
+        gr.Markdown("### ✅ Final Compliance Verdict")
+    
+    with gr.Row():
+        final_output = gr.Markdown(value="*Analysis pending...*")
+    
+    submit_btn.click(
+            fn=run_agentic_workflow,
+            inputs=[question_input],
+            outputs=[
+                log_output,
+                analysis_output,   # NEW SLOT
+                chapter_output,
+                image_gallery,
+                final_output
+            ]
+        )
+
+if __name__ == "__main__":
+    demo.queue().launch(
+        inbrowser=True,
+        auth=("username", os.getenv("PASSWORD"))  # only share the password
     )