vectorstore/bills_vectorstore.py

# vectorstore/bills_vectorstore.py
from __future__ import annotations
import os, json, hashlib, time
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeout
from pathlib import Path
from typing import Dict, List, Optional, Iterable, Any

from dotenv import load_dotenv, find_dotenv
load_dotenv(find_dotenv())

try:
    from langchain_chroma import Chroma
except Exception:
    from langchain_community.vectorstores import Chroma

from langchain_openai import OpenAIEmbeddings
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter

DEFAULT_EMBED_MODEL = os.getenv("OPENAI_EMBED_MODEL", "text-embedding-3-small")
DEFAULT_PERSIST_DIR = "data/bills_vectorstore"
DEFAULT_COLLECTION = "bills"
DEFAULT_MANIFEST = "data/bills_vectorstore_manifest.json"

def get_embeddings(model: Optional[str] = None) -> OpenAIEmbeddings:
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        raise RuntimeError("OPENAI_API_KEY is not set. Check your .env or environment.")
    return OpenAIEmbeddings(api_key=api_key, model=model or DEFAULT_EMBED_MODEL, chunk_size=32, request_timeout=120)

def _sha256(text: str) -> str:
    import hashlib
    return hashlib.sha256(text.encode("utf-8")).hexdigest()

def _bill_id(b: Dict[str, Any]) -> str:
    # Use LegiScan bill_id when available for uniqueness across sessions;
    # fall back to state_billnumber for backwards compatibility.
    lid = b.get("bill_id")
    if lid:
        return str(lid)
    return f"{b.get('state','Unknown')}_{b.get('bill_number','Unknown')}"

def _bill_text(b: Dict[str, Any]) -> str:
    title = b.get("title") or ""
    summary = b.get("description") or ""
    txt = b.get("text") or ""
    return f"Title: {title}\n\nSummary: {summary}\n\nFull Text:\n{txt}"

def _bill_hash(b: Dict[str, Any]) -> str:
    payload = json.dumps({
        "title": b.get("title"),
        "description": b.get("description"),
        "text": b.get("text"),
        "status": b.get("status"),
        "last_action_date": b.get("last_action_date"),
    }, ensure_ascii=False, sort_keys=True)
    return _sha256(payload)

def _manifest_load(path: str) -> Dict[str, Dict[str, str]]:
    p = Path(path)
    if not p.exists():
        return {}
    try:
        return json.loads(p.read_text(encoding="utf-8"))
    except Exception:
        return {}

def _manifest_save(path: str, data: Dict[str, Dict[str, str]]) -> None:
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    Path(path).write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")

def _clean_metadata(meta: Dict[str, Any]) -> Dict[str, Any]:
    """Keep only metadata values that Chroma accepts: str/int/float/bool and not None."""
    allowed_types = (str, int, float, bool)
    cleaned: Dict[str, Any] = {}
    for k, v in meta.items():
        if v is None:
            continue
        if isinstance(v, allowed_types):
            cleaned[k] = v
        else:
            # If you prefer to drop complex types instead of stringifying, replace with `continue`
            cleaned[k] = str(v)
    return cleaned


def _make_doc(b: Dict[str, Any]) -> Document:
    sponsors_list = b.get("sponsors") or []
    if isinstance(sponsors_list, list):
        sponsors_str = "; ".join(map(str, sponsors_list))
    else:
        sponsors_str = str(sponsors_list) if sponsors_list else ""

    flat_iapp = []
    iapp = b.get("iapp_categories")
    if isinstance(iapp, dict):
        for k, v in iapp.items():
            if isinstance(v, list):
                for sub in v:
                    flat_iapp.append(f"{k}:{sub}")
    iapp_str = "; ".join(flat_iapp) if flat_iapp else ""

    meta = {
        "doc_id": _bill_id(b),
        "state": b.get("state"),
        "session_year": b.get("session_year"),
        "legislative_body": b.get("chamber") or b.get("legislative_body") or None,
        "status": b.get("status"),
        "title": b.get("title"),
        "bill_number": b.get("bill_number"),
        "sponsors": sponsors_str,
        "last_action_date": b.get("last_action_date"),
        "iapp_flat": iapp_str,
    }

    meta = _clean_metadata(meta)

    return Document(page_content=_bill_text(b), metadata=meta)

def _load_bills(source_json_path: str) -> List[Dict[str, Any]]:
    data = json.loads(Path(source_json_path).read_text(encoding="utf-8"))
    if not isinstance(data, list):
        raise ValueError(f"{source_json_path} must contain a list of bills")
    return data

def load_vectorstore(
    persist_dir: str = DEFAULT_PERSIST_DIR,
    collection: str = DEFAULT_COLLECTION,
    embeddings: Optional[OpenAIEmbeddings] = None,
) -> Chroma:
    embeddings = embeddings or get_embeddings()
    Path(persist_dir).mkdir(parents=True, exist_ok=True)
    return Chroma(
        collection_name=collection,
        persist_directory=persist_dir,
        embedding_function=embeddings,
    )

def _chunk_bill(b: Dict[str, Any], *, size: int = 1500, overlap: int = 200) -> List[Document]:
    text = _bill_text(b)
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=size, chunk_overlap=overlap,
        separators=["\n\n", "\n", ". ", " ", ""]
    )
    pieces = splitter.split_text(text) or ["(no content)"]

    docs: List[Document] = []
    base_meta = {
        "doc_id": _bill_id(b),
        "state": b.get("state"),
        "session_year": b.get("session_year"),
        "legislative_body": b.get("chamber") or b.get("legislative_body") or None,
        "status": b.get("status"),
        "title": b.get("title"),
        "bill_number": b.get("bill_number"),
        "sponsors": (("; ".join(map(str, b.get("sponsors") or [])))
                     if isinstance(b.get("sponsors"), list)
                     else (b.get("sponsors") or "")),
        "last_action_date": b.get("last_action_date"),
    }

    iapp = b.get("iapp_categories") or {}
    flat = []
    if isinstance(iapp, dict):
        for k, v in iapp.items():
            if isinstance(v, list):
                for sub in v:
                    flat.append(f"{k}:{sub}")
    base_meta["iapp_flat"] = "; ".join(flat)

    # 🔑 Clean out None / bad types before using this as metadata
    base_meta = _clean_metadata(base_meta)


    total = len(pieces)
    for i, chunk in enumerate(pieces):
        m = dict(base_meta)
        m["chunk_index"] = i
        m["chunk_total"] = total
        docs.append(Document(page_content=chunk, metadata=m))
    return docs

def upsert_from_bills_json(
    source_json_path: str = "data/known_bills_visualize.json",
    persist_dir: str = DEFAULT_PERSIST_DIR,
    collection: str = DEFAULT_COLLECTION,
    manifest_path: str = DEFAULT_MANIFEST,
    embed_model: Optional[str] = None,
    batch_size: int = 128,
) -> Dict[str, int]:
    t0 = time.time()
    bills = _load_bills(source_json_path)
    embeddings = get_embeddings(embed_model)
    vs = load_vectorstore(persist_dir, collection, embeddings)
    manifest = _manifest_load(manifest_path)
    manifest_meta = manifest.get("_meta", {})
    if manifest_meta.get("embed_model") != (embed_model or DEFAULT_EMBED_MODEL):
        manifest = {}
        manifest["_meta"] = {"embed_model": embed_model or DEFAULT_EMBED_MODEL}

    to_docs, to_ids = [], []
    added, skipped = 0, 0

    def _add_batch(docs, ids):
        """Add a batch of documents with a thread-based 5-minute timeout."""
        with ThreadPoolExecutor(max_workers=1) as executor:
            future = executor.submit(vs.add_documents, documents=docs, ids=ids)
            try:
                future.result(timeout=300)
            except FuturesTimeout:
                raise TimeoutError("Embedding batch exceeded 300s hard timeout")

    for b in bills:
        if not (b.get("text") or b.get("description") or b.get("title")):
            skipped += 1
            continue
        doc_id = _bill_id(b)
        hsh = _bill_hash(b)
        if manifest.get(doc_id, {}).get("hash") == hsh:
            skipped += 1
            continue
        try:
            vs.delete(where={"doc_id": doc_id})
        except Exception:
            pass
        chunks = _chunk_bill(b)
        for d in chunks:
            to_docs.append(d)
            to_ids.append(f"{doc_id}::c{d.metadata['chunk_index']}")
            if len(to_docs) >= batch_size:
                _add_batch(to_docs, to_ids)
                to_docs, to_ids = [], []
        manifest[doc_id] = {"hash": hsh}
        added += 1

    if to_docs:
        _add_batch(to_docs, to_ids)

    if hasattr(vs, "persist"):
        vs.persist()

    manifest["_meta"] = {"embed_model": embed_model or DEFAULT_EMBED_MODEL}
    _manifest_save(manifest_path, manifest)

    return {
        "total_bills": len(bills),
        "embedded": added,
        "skipped_unchanged": skipped,
        "elapsed_sec": int(time.time() - t0),
    }

def get_retriever(persist_dir=DEFAULT_PERSIST_DIR, collection=DEFAULT_COLLECTION, k=8, filter_kwargs=None):
    vs = load_vectorstore(persist_dir=persist_dir, collection=collection)
    search_kwargs = {"k": k}
    if filter_kwargs:
        search_kwargs["filter"] = filter_kwargs
    return vs.as_retriever(search_kwargs=search_kwargs)

def similarity_search(
    query: str,
    k: int = 5,
    where: Optional[Dict[str, Any]] = None,
    persist_dir: str = DEFAULT_PERSIST_DIR,
    collection: str = DEFAULT_COLLECTION,
):
    vs = load_vectorstore(persist_dir=persist_dir, collection=collection)
    filt = where if (where and len(where) > 0) else None   # <-- key line
    return vs.similarity_search(query, k=k, filter=filt)