vector_store.py

"""Vector store operations for document retrieval."""

from __future__ import annotations

from typing import List, Sequence, Tuple

from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_openai import OpenAIEmbeddings

from config import COLLECTION_NAME, EMBEDDING_MODEL, VECTOR_DIR
from helpers import _chunk_section_or_preview, _clip


def load_vector_store() -> Tuple[Chroma | None, str | None]:
    """Load the Chroma vector store."""
    if not VECTOR_DIR.exists():
        return None, "Vector store not found. Run `python build_vectordb.py --recreate` first."

    try:
        embeddings = OpenAIEmbeddings(model=EMBEDDING_MODEL)
        store = Chroma(
            collection_name=COLLECTION_NAME,
            persist_directory=str(VECTOR_DIR),
            embedding_function=embeddings,
        )
        return store, None
    except Exception as exc:  # pragma: no cover
        return None, f"Failed to load vector store: {exc}"


def retrieve_report_chunks(store: Chroma, query: str, state_slugs: Sequence[str], k: int = 8) -> List[Document]:
    """Retrieve relevant report chunks from vector store."""
    target_states = set(state_slugs)
    candidates = store.similarity_search(query, k=max(24, k * 4))
    filtered = [doc for doc in candidates if doc.metadata.get("state") in target_states]

    if not filtered and len(target_states) == 1:
        only_state = next(iter(target_states))
        try:
            filtered = store.similarity_search(query, k=k, filter={"state": only_state})
        except Exception:
            filtered = []

    deduped: List[Document] = []
    seen = set()
    for doc in filtered:
        key = (doc.metadata.get("source"), doc.metadata.get("chunk_id"), doc.page_content[:100])
        if key in seen:
            continue
        seen.add(key)
        deduped.append(doc)
        if len(deduped) >= k:
            break

    return deduped


def format_report_context(docs: Sequence[Document]) -> Tuple[str, List[str]]:
    """Format retrieved documents into context string and source references."""
    snippets: List[str] = []
    source_refs: List[str] = []

    for idx, doc in enumerate(docs, start=1):
        state_label = doc.metadata.get("state_label") or doc.metadata.get("state", "Unknown")
        section_or_preview = _chunk_section_or_preview(doc.page_content)
        snippets.append(f"[S{idx}] ({state_label} | {section_or_preview})\n{_clip(doc.page_content)}")
        source_refs.append(f"- [S{idx}] {state_label} | {section_or_preview}")

    return "\n\n".join(snippets), source_refs