Spaces:
Sleeping
Sleeping
File size: 1,788 Bytes
1fee1c2 | 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 | import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
from dotenv import load_dotenv
load_dotenv()
from chroma import get_collection # noqa: E402
def retrieve_chunks(query: str, top_k: int = 5) -> list[str]:
"""Search the vector store for document chunks most relevant to the query.
Embeds the query using the same OpenAI embedding model used during ingestion,
performs a cosine-similarity search against the ChromaDB collection, and
returns the top-k matching text chunks formatted with source metadata so
the caller can cite the origin of each passage.
Args:
query: The natural-language question or search string to look up.
top_k: Maximum number of chunks to return (default 5).
Returns:
A list of strings, each formatted as:
"[source: <filename>, page: <n>, chunk: <i>]\\n<chunk text>"
Page number is omitted for DOCX files where it is unavailable.
"""
collection = get_collection()
count = collection.count()
if count == 0:
return ["No documents have been indexed yet. Please ingest documents first."]
results = collection.query(
query_texts=[query],
n_results=min(top_k, count),
include=["documents", "metadatas"],
)
chunks: list[str] = []
for doc, meta in zip(results["documents"][0], results["metadatas"][0]):
source = meta.get("source_file", "unknown")
chunk_idx = meta.get("chunk_index", "?")
page = meta.get("page_number")
if page is not None:
header = f"[source: {source}, page: {page}, chunk: {chunk_idx}]"
else:
header = f"[source: {source}, chunk: {chunk_idx}]"
chunks.append(f"{header}\n{doc}")
return chunks
|