agent2_vectorizer.py

"""
Agent 2: Vectorization Engine
Embeds the structured plan extraction (from Agent 1) and predefined compliance rules
into a searchable vector store using OpenAI embeddings + FAISS.

Chunking strategy:
  - Plan side:  one chunk per logical element (room, door, window, stair, setback group)
  - Rules side: one chunk per rule (split by rule number or double-newline)
"""

import os
import json
import re
import pickle
from pathlib import Path
from typing import Optional
import numpy as np
import faiss
from openai import OpenAI

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

EMBEDDING_MODEL = "text-embedding-3-large"
EMBEDDING_DIM = 3072  # Dimension for text-embedding-3-large


# ─────────────────────────────────────────────
# Embedding utilities
# ─────────────────────────────────────────────
def get_embeddings(texts: list[str], batch_size: int = 100) -> np.ndarray:
    """Embed a list of texts in batches, return numpy array."""
    all_embeddings = []
    
    for i in range(0, len(texts), batch_size):
        batch = texts[i : i + batch_size]
        response = client.embeddings.create(
            model=EMBEDDING_MODEL,
            input=batch
        )
        batch_embeddings = [item.embedding for item in response.data]
        all_embeddings.extend(batch_embeddings)
    
    return np.array(all_embeddings, dtype=np.float32)


# ─────────────────────────────────────────────
# Plan chunking: JSON → human-readable text chunks
# ─────────────────────────────────────────────
def chunk_plan_extraction(plan_data: dict) -> list[dict]:
    """
    Convert structured plan JSON into embeddable text chunks.
    Each chunk retains metadata for traceability in Agent 3.
    
    Returns list of:
      { "text": str, "type": str, "id": str, "raw": dict, "confidence": str }
    """
    chunks = []
    
    # ── Rooms ──
    for room in plan_data.get("rooms", []):
        name = room.get("name", "Unnamed Room")
        parts = [f"Room: {name}"]
        if room.get("width_ft") and room.get("length_ft"):
            parts.append(f"Dimensions: {room['width_ft']} ft x {room['length_ft']} ft")
        if room.get("area_sqft"):
            parts.append(f"Area: {room['area_sqft']} sq ft")
        if room.get("ceiling_height_ft"):
            parts.append(f"Ceiling height: {room['ceiling_height_ft']} ft")
        if room.get("floor_level"):
            parts.append(f"Floor level: {room['floor_level']}")
        
        chunks.append({
            "text": ". ".join(parts) + ".",
            "type": "room",
            "id": f"room_{name.replace(' ', '_').lower()}",
            "raw": room,
            "confidence": room.get("confidence", "medium"),
            "sheet": room.get("_sheet", "unknown")
        })
    
    # ── Doors ──
    for door in plan_data.get("doors", []):
        door_id = door.get("id", "unknown")
        parts = [f"Door ID: {door_id}"]
        if door.get("type"):
            parts.append(f"Type: {door['type']}")
        if door.get("width_in"):
            parts.append(f"Width: {door['width_in']} inches ({door['width_in']/12:.2f} ft)")
        if door.get("height_in"):
            parts.append(f"Height: {door['height_in']} inches")
        if door.get("location"):
            parts.append(f"Location: {door['location']}")
        if door.get("swing_direction"):
            parts.append(f"Swing direction: {door['swing_direction']}")
        
        chunks.append({
            "text": ". ".join(parts) + ".",
            "type": "door",
            "id": f"door_{door_id}",
            "raw": door,
            "confidence": door.get("confidence", "medium"),
            "sheet": door.get("_sheet", "unknown")
        })
    
    # ── Windows ──
    for window in plan_data.get("windows", []):
        win_id = window.get("id", "unknown")
        parts = [f"Window ID: {win_id}"]
        if window.get("type"):
            parts.append(f"Type: {window['type']}")
        if window.get("width_in"):
            parts.append(f"Width: {window['width_in']} inches")
        if window.get("height_in"):
            parts.append(f"Height: {window['height_in']} inches")
        if window.get("sill_height_in"):
            parts.append(f"Sill height: {window['sill_height_in']} inches from floor")
        if window.get("net_clear_area_sqft"):
            parts.append(f"Net clear opening area: {window['net_clear_area_sqft']} sq ft")
        if window.get("egress_capable") is not None:
            parts.append(f"Egress capable: {window['egress_capable']}")
        if window.get("location"):
            parts.append(f"Location: {window['location']}")
        
        chunks.append({
            "text": ". ".join(parts) + ".",
            "type": "window",
            "id": f"window_{win_id}",
            "raw": window,
            "confidence": window.get("confidence", "medium"),
            "sheet": window.get("_sheet", "unknown")
        })
    
    # ── Stairs ──
    for stair in plan_data.get("stairs", []):
        stair_id = stair.get("id", "stair_1")
        parts = [f"Staircase ID: {stair_id}"]
        if stair.get("width_in"):
            parts.append(f"Width: {stair['width_in']} inches")
        if stair.get("rise_in"):
            parts.append(f"Riser height: {stair['rise_in']} inches")
        if stair.get("run_in"):
            parts.append(f"Tread depth: {stair['run_in']} inches")
        if stair.get("num_risers"):
            parts.append(f"Number of risers: {stair['num_risers']}")
        if stair.get("handrail_present") is not None:
            parts.append(f"Handrail present: {stair['handrail_present']}")
        
        chunks.append({
            "text": ". ".join(parts) + ".",
            "type": "stair",
            "id": f"stair_{stair_id}",
            "raw": stair,
            "confidence": stair.get("confidence", "medium"),
            "sheet": stair.get("_sheet", "unknown")
        })
    
    # ── Setbacks ──
    setbacks = plan_data.get("setbacks", {})
    if any(v is not None for v in setbacks.values()):
        parts = ["Site setbacks"]
        for side in ["front", "rear", "left_side", "right_side"]:
            val = setbacks.get(f"{side}_ft")
            if val is not None:
                parts.append(f"{side.replace('_', ' ').title()}: {val} ft")
        
        chunks.append({
            "text": ". ".join(parts) + ".",
            "type": "setback",
            "id": "setbacks",
            "raw": setbacks,
            "confidence": "medium",
            "sheet": "site_plan"
        })
    
    # ── Annotations / General notes ──
    for i, annotation in enumerate(plan_data.get("annotations", [])):
        if isinstance(annotation, str) and len(annotation.strip()) > 10:
            chunks.append({
                "text": f"Plan annotation: {annotation}",
                "type": "annotation",
                "id": f"annotation_{i}",
                "raw": {"text": annotation},
                "confidence": "high",
                "sheet": "general"
            })
    
    return chunks


# ─────────────────────────────────────────────
# Rules chunking: text → individual rule chunks
# ─────────────────────────────────────────────
def chunk_rules(rules_text: str) -> list[dict]:
    """
    Parse rules text into individual rule chunks.
    Supports formats:
      - Numbered rules: "1. Rule text" or "Rule 1:" or "R-001:"
      - Section-based: "Section 3.2 - ..."
      - Double-newline separated blocks
    """
    chunks = []
    
    # Try numbered rule pattern first
    numbered_pattern = re.compile(
        r'(?:^|\n)(?:Rule\s*)?(\d+[\.\-\)]\d*\.?\d*\.?)\s*(.+?)(?=\n(?:Rule\s*)?\d+[\.\-\)]\d*|$)',
        re.DOTALL
    )
    matches = list(numbered_pattern.finditer(rules_text))
    
    if len(matches) >= 2:
        for match in matches:
            rule_id = match.group(1).strip().rstrip(".")
            rule_text = match.group(2).strip()
            if len(rule_text) > 10:
                chunks.append({
                    "text": f"Rule {rule_id}: {rule_text}",
                    "rule_id": rule_id,
                    "raw_text": rule_text,
                    "type": "rule"
                })
    else:
        # Fall back to double-newline splitting
        blocks = re.split(r'\n\s*\n', rules_text.strip())
        for i, block in enumerate(blocks):
            block = block.strip()
            if len(block) > 10:
                chunks.append({
                    "text": block,
                    "rule_id": f"R-{i+1:03d}",
                    "raw_text": block,
                    "type": "rule"
                })
    
    return chunks


# ─────────────────────────────────────────────
# FAISS index builder
# ─────────────────────────────────────────────
def build_faiss_index(embeddings: np.ndarray) -> faiss.IndexFlatIP:
    """
    Build an inner-product FAISS index (equivalent to cosine similarity
    when vectors are L2-normalized).
    """
    # Normalize for cosine similarity
    faiss.normalize_L2(embeddings)
    index = faiss.IndexFlatIP(EMBEDDING_DIM)
    index.add(embeddings)
    return index


# ─────────────────────────────────────────────
# Main Agent 2 pipeline
# ─────────────────────────────────────────────
def vectorize(
    plan_json_path: str,
    rules_text_path: str,
    output_dir: str = "vector_store"
) -> dict:
    """
    Full vectorization pipeline.
    
    Args:
        plan_json_path:  Path to extracted_plan.json from Agent 1
        rules_text_path: Path to plain text file containing compliance rules
        output_dir:      Directory to save indexes and chunk metadata
    
    Returns:
        Dictionary with loaded indexes and chunk metadata (passed to Agent 3)
    """
    
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    # ── Load inputs ──
    print("[Agent 2] Loading plan extraction...")
    with open(plan_json_path) as f:
        plan_data = json.load(f)
    
    print("[Agent 2] Loading rules text...")
    with open(rules_text_path) as f:
        rules_text = f.read()
    
    # ── Chunk ──
    print("[Agent 2] Chunking plan elements...")
    plan_chunks = chunk_plan_extraction(plan_data)
    print(f"[Agent 2] Plan chunks: {len(plan_chunks)}")
    
    print("[Agent 2] Chunking rules...")
    rule_chunks = chunk_rules(rules_text)
    print(f"[Agent 2] Rule chunks: {len(rule_chunks)}")
    
    # ── Embed plan chunks ──
    print("[Agent 2] Embedding plan chunks...")
    plan_texts = [c["text"] for c in plan_chunks]
    plan_embeddings = get_embeddings(plan_texts)
    
    # ── Embed rule chunks ──
    print("[Agent 2] Embedding rule chunks...")
    rule_texts = [c["text"] for c in rule_chunks]
    rule_embeddings = get_embeddings(rule_texts)
    
    # ── Build FAISS indexes ──
    print("[Agent 2] Building FAISS indexes...")
    plan_index = build_faiss_index(plan_embeddings.copy())
    rule_index = build_faiss_index(rule_embeddings.copy())
    
    # ── Save to disk ──
    plan_index_path = f"{output_dir}/plan.index"
    rule_index_path = f"{output_dir}/rules.index"
    plan_chunks_path = f"{output_dir}/plan_chunks.pkl"
    rule_chunks_path = f"{output_dir}/rule_chunks.pkl"
    embeddings_path  = f"{output_dir}/embeddings.npz"
    
    faiss.write_index(plan_index, plan_index_path)
    faiss.write_index(rule_index, rule_index_path)
    
    with open(plan_chunks_path, "wb") as f:
        pickle.dump(plan_chunks, f)
    with open(rule_chunks_path, "wb") as f:
        pickle.dump(rule_chunks, f)
    
    np.savez(
        embeddings_path,
        plan_embeddings=plan_embeddings,
        rule_embeddings=rule_embeddings
    )
    
    print(f"[Agent 2] Vector store saved to: {output_dir}/")
    
    return {
        "plan_index": plan_index,
        "rule_index": rule_index,
        "plan_chunks": plan_chunks,
        "rule_chunks": rule_chunks,
        "plan_embeddings": plan_embeddings,
        "rule_embeddings": rule_embeddings,
        "output_dir": output_dir
    }


def load_vector_store(output_dir: str = "vector_store") -> dict:
    """Load a previously saved vector store from disk."""
    
    plan_index = faiss.read_index(f"{output_dir}/plan.index")
    rule_index = faiss.read_index(f"{output_dir}/rules.index")
    
    with open(f"{output_dir}/plan_chunks.pkl", "rb") as f:
        plan_chunks = pickle.load(f)
    with open(f"{output_dir}/rule_chunks.pkl", "rb") as f:
        rule_chunks = pickle.load(f)
    
    emb = np.load(f"{output_dir}/embeddings.npz")
    
    return {
        "plan_index": plan_index,
        "rule_index": rule_index,
        "plan_chunks": plan_chunks,
        "rule_chunks": rule_chunks,
        "plan_embeddings": emb["plan_embeddings"],
        "rule_embeddings": emb["rule_embeddings"],
        "output_dir": output_dir
    }


# ─────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────
if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 3:
        print("Usage: python agent2_vectorizer.py <extracted_plan.json> <rules.txt>")
        sys.exit(1)
    
    result = vectorize(
        plan_json_path=sys.argv[1],
        rules_text_path=sys.argv[2],
        output_dir="vector_store"
    )
    
    print(f"\nVector store ready:")
    print(f"  Plan chunks indexed:  {len(result['plan_chunks'])}")
    print(f"  Rule chunks indexed:  {len(result['rule_chunks'])}")