added docs

backend/app/rag/retriever.py

@@ -9,9 +9,14 @@ from app.rag import embeddings
 from app.rag.store import COLLECTION_DOC_CHUNKS, COLLECTION_FACTS, search
 
 
-def _user_filter(user_id: int) -> qm.Filter:
+def _document_filter(user_id: int, document_id: int | None = None) -> qm.Filter:
+    must: list[qm.FieldCondition] = [
+        qm.FieldCondition(key="user_id", match=qm.MatchValue(value=user_id))
+    ]
+    if document_id is not None:
+        must.append(qm.FieldCondition(key="document_id", match=qm.MatchValue(value=document_id)))
     return qm.Filter(
-        must=[qm.FieldCondition(key="user_id", match=qm.MatchValue(value=user_id))]
+        must=must
     )
 
 
@@ -25,7 +30,7 @@ async def retrieve_memory_facts(
     vectors = await embeddings.embed_texts([query])
     if not vectors:
         return []
-    hits = search(COLLECTION_FACTS, vectors[0], limit=k, query_filter=_user_filter(user_id))
+    hits = search(COLLECTION_FACTS, vectors[0], limit=k, query_filter=_document_filter(user_id))
     results: list[dict[str, Any]] = []
     for hit in hits:
         payload = hit.payload or {}
@@ -41,7 +46,7 @@ async def retrieve_memory_facts(
 
 
 async def retrieve_document_chunks(
-    query: str, *, user_id: int, top_k: int = 6
+    query: str, *, user_id: int, top_k: int = 6, document_id: int | None = None
 ) -> list[dict[str, Any]]:
     settings = get_settings()
     if not settings.rag_enabled or not query.strip():
@@ -50,7 +55,10 @@ async def retrieve_document_chunks(
     if not vectors:
         return []
     hits = search(
-        COLLECTION_DOC_CHUNKS, vectors[0], limit=top_k, query_filter=_user_filter(user_id)
+        COLLECTION_DOC_CHUNKS,
+        vectors[0],
+        limit=top_k,
+        query_filter=_document_filter(user_id, document_id),
     )
     out: list[dict[str, Any]] = []
     for hit in hits:
@@ -58,8 +66,16 @@ async def retrieve_document_chunks(
         out.append(
             {
                 "document_id": payload.get("document_id"),
+                "chunk_id": payload.get("chunk_id"),
                 "chunk_index": payload.get("chunk_index"),
                 "title": payload.get("title", ""),
+                "filename": payload.get("filename", ""),
+                "mime_type": payload.get("mime_type", ""),
+                "page_start": payload.get("page_start"),
+                "page_end": payload.get("page_end"),
+                "section_title": payload.get("section_title", ""),
+                "section_path": payload.get("section_path", ""),
+                "source_ref": payload.get("source_ref", ""),
                 "content": payload.get("content", ""),
                 "score": hit.score,
             }