Fix chroma retrieval

src/backend/rag.py

@@ -4,53 +4,117 @@ from src.backend.llm import llm_generate
 
 EMBED_MODEL = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
+
 def retrieve(username: str, notebook_id: str, query: str, k=6):
     col = get_collection(username, notebook_id)
-    qemb = EMBED_MODEL.encode([query], normalize_embeddings=True).tolist()
-    res = col.query(query_embeddings=qemb, n_results=k, include=["documents","metadatas","ids"])
+
+    qemb = EMBED_MODEL.encode(
+        [query],
+        normalize_embeddings=True
+    ).tolist()
+
+    # Compatible with HF Chroma version
+    res = col.query(
+        query_embeddings=qemb,
+        n_results=k,
+        include=["documents", "metadatas", "distances"]
+    )
+
+    # Safe extraction (prevents crashes)
+    ids = res.get("ids", [[]])[0]
+    docs = res.get("documents", [[]])[0]
+    mets = res.get("metadatas", [[]])[0]
+    dists = res.get("distances", [[]])[0]
+
     hits = []
-    for i in range(len(res["ids"][0])):
-        hits.append({"id": res["ids"][0][i], "doc": res["documents"][0][i], "meta": res["metadatas"][0][i]})
+
+    for i in range(len(docs)):
+        hits.append({
+            "id": ids[i] if i < len(ids) else f"chunk_{i}",
+            "doc": docs[i],
+            "meta": mets[i] if i < len(mets) else {},
+            "distance": dists[i] if i < len(dists) else None
+        })
+
     return hits
 
+
 def format_sources(hits):
     lines = []
+
     for i, h in enumerate(hits, start=1):
-        m = h["meta"]
+        m = h["meta"] or {}
+
         loc = ""
-        if m.get("page"): loc = f"p.{m['page']}"
-        if m.get("slide"): loc = f"slide {m['slide']}"
-        lines.append(f"[S{i}] {m.get('source_title','source')} {loc}".strip())
+        if m.get("page"):
+            loc = f"p.{m['page']}"
+
+        if m.get("slide"):
+            loc = f"slide {m['slide']}"
+
+        title = m.get("source_title", "source")
+
+        lines.append(f"[S{i}] {title} {loc}".strip())
+
     return "\n".join(lines)
 
+
 def context_block(hits):
     blocks = []
+
     for i, h in enumerate(hits, start=1):
-        m = h["meta"]
+        m = h["meta"] or {}
+
         loc = ""
-        if m.get("page"): loc = f"(page {m['page']})"
-        if m.get("slide"): loc = f"(slide {m['slide']})"
-        blocks.append(f"[S{i}] {m.get('source_title','source')} {loc}\n{h['doc']}")
+        if m.get("page"):
+            loc = f"(page {m['page']})"
+
+        if m.get("slide"):
+            loc = f"(slide {m['slide']})"
+
+        title = m.get("source_title", "source")
+
+        blocks.append(
+            f"[S{i}] {title} {loc}\n{h['doc']}"
+        )
+
     return "\n\n---\n\n".join(blocks)
 
+
 def rag_answer(query: str, hits):
+
     if not hits:
         return "Not found in the provided sources. (No indexed chunks yet.)"
+
     prompt = f"""
-You are a research assistant. Answer ONLY using the sources below.
+You are a research assistant.
+
+Answer ONLY using the sources below.
+
 Every non-trivial claim must end with citations like [S1] or [S2].
-If not present in sources, say: Not found in the provided sources.
+
+If not present in sources say:
+Not found in the provided sources.
 
 Question:
 {query}
 
+
 Sources list:
 {format_sources(hits)}
 
+
 Source excerpts:
 {context_block(hits)}
 
+
 Answer with citations:
 """
-    ans = llm_generate(prompt, max_new_tokens=450, temperature=0.2)
+
+    ans = llm_generate(
+        prompt,
+        max_new_tokens=450,
+        temperature=0.2
+    )
+
     return f"{ans}\n\nSources:\n{format_sources(hits)}"