Spaces:
Sleeping
Sleeping
File size: 1,984 Bytes
e7f02ca | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | from rag.hybrid_retriever import HybridSearch
class Retriever:
def __init__(self):
self.hybrid_search = HybridSearch()
def retrieve_documents(
self,
query: str,
user_id: str,
idea_id: str,
top_k: int = 5
):
# -----------------------------------
# Hybrid Retrieval
# -----------------------------------
docs = self.hybrid_search.search(
query=query,
user_id=user_id,
idea_id=idea_id,
top_k=top_k
)
# -----------------------------------
# Build Citation-Aware Context
# -----------------------------------
formatted_chunks = []
for doc in docs:
if not isinstance(doc, dict):
continue
text = doc.get("text", "")
if not text:
continue
meta = doc.get("metadata", {})
filename = meta.get(
"filename",
"Unknown"
)
page_number = meta.get(
"page_number"
)
# -----------------------------------
# Citation
# -----------------------------------
if page_number:
citation = (
f"[Source: {filename}, "
f"Page {page_number}]"
)
else:
citation = (
f"[Source: {filename}]"
)
formatted_chunks.append(
f"{text}\n{citation}"
)
# -----------------------------------
# Final Context
# -----------------------------------
context = "\n\n".join(formatted_chunks)
return {
"query": query,
"context": context,
"documents": docs
} |