Spaces:
Sleeping
Sleeping
File size: 2,604 Bytes
e2df10e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """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
|