feat(query): multi-query expansion for improved retrieval (#13)

* feat(schemas): add query request and response schemas with similarity config

* feat(repo): add vector search query scoped by owner and hub

* fix(ai): strip markdown code blocks from gemini json responses

* feat(query): add grounded rag endpoint with similarity threshold

* test(query): verify cross-tenant isolation and empty retrieval

* feat(ai): add expand_query for multi-query retrieval

* feat(repo): add multi-query union retrieval

* feat(query): integrate multi-query expansion into rag

* test(query): cover multihop sources and expansion fallback

app/repo.py

@@ -1,7 +1,6 @@
 from uuid import UUID
 from sqlalchemy.orm import Session
-from app.models import Hub
-from app.models import Chunk
+from app.models import Hub, Chunk
 from app.config import settings
 
 def list_hubs_for_owner(db: Session, owner_id: UUID, limit: int, offset: int):
@@ -14,15 +13,32 @@ def get_hub_for_owner(db: Session, hub_id: UUID, owner_id: UUID) -> Hub | None:
     return db.query(Hub).filter(Hub.id == hub_id, Hub.owner_id == owner_id).first()
 
 def delete_hub_for_owner(db: Session, hub_id: UUID, owner_id: UUID) -> int:
-    # Returns the number of rows deleted (0 or 1)
     return db.query(Hub).filter(Hub.id == hub_id, Hub.owner_id == owner_id).delete()
 
 def retrieve_chunks_for_query(db: Session, owner_id: UUID, hub_id: UUID, query_embedding: list[float]):
-    # We filter by owner and hub FIRST, then order by cosine distance
     return (
         db.query(Chunk)
         .filter(Chunk.owner_id == owner_id, Chunk.hub_id == hub_id)
         .order_by(Chunk.embedding.cosine_distance(query_embedding))
         .limit(settings.RAG_TOP_K)
         .all()
-    )
+    )
+
+def retrieve_chunks_for_queries(db: Session, owner_id: UUID, hub_id: UUID, query_embeddings: list[list[float]]):
+    chunk_ids = set()
+    chunks = []
+    
+    for emb in query_embeddings:
+        results = (
+            db.query(Chunk)
+            .filter(Chunk.owner_id == owner_id, Chunk.hub_id == hub_id)
+            .order_by(Chunk.embedding.cosine_distance(emb))
+            .limit(settings.RAG_TOP_K)
+            .all()
+        )
+        for r in results:
+            if r.id not in chunk_ids:
+                chunk_ids.add(r.id)
+                chunks.append(r)
+                
+    return chunks

app/routers/hubs.py

@@ -2,10 +2,12 @@ from uuid import UUID
 from fastapi import APIRouter, Depends, status, Query, Response
 from sqlalchemy.orm import Session
 from sqlalchemy.exc import IntegrityError
+from sqlalchemy import text
+
 from app.db import get_db
 from app.deps import get_current_user
 from app.models import User, Hub, Document, Chunk
-from app.schemas import HubResponse, HubDetail, HubCreate
+from app.schemas import HubResponse, HubDetail, HubCreate, QueryRequest, QueryResponse, QuerySource
 from app.errors import AppError
 from app import repo
 from app.idempotency import (
@@ -14,12 +16,9 @@ from app.idempotency import (
 )
 from app.chunking import chunk_text
 from app.services.embeddings import embed_texts
-from app.services.gemini import hub_insights
+from app.services.gemini import hub_insights, answer_question, expand_query
 from app.services.observability import log_ai_call
-from app.schemas import QueryRequest, QueryResponse, QuerySource
-from app.services.gemini import answer_question
-from app.services.embeddings import embed_texts
-
+from app.config import settings
 
 router = APIRouter(prefix="/hubs", tags=["hubs"])
 
@@ -80,21 +79,15 @@ def create_hub(
     if not claim_key(db, user_id, key, req_hash, "/api/v1/hubs"):
         return replay_or_conflict(db, user_id, key, req_hash)
 
-    # 1. Pre-check (Fast fail to save AI tokens)
     if db.query(Hub).filter(Hub.owner_id == user_id, Hub.name == body.name).first():
         release_claim(db, user_id, key)
         raise AppError(409, "HUB_NAME_CONFLICT", "You already have a hub with this name.")
 
-    # 2. AI Summary (Wrapped in observability)
     if body.content:
-        summary, questions = log_ai_call(
-            "hub_insights", user_id, None, 
-            hub_insights, body.content
-        )
+        summary, questions = log_ai_call("hub_insights", user_id, None, hub_insights, body.content)
     else:
         summary, questions = None, []
 
-    # 3. Create Hub
     new_hub = Hub(
         owner_id=user_id, name=body.name, description=body.description, 
         ai_summary=summary, suggested_questions=questions
@@ -102,7 +95,6 @@ def create_hub(
     db.add(new_hub)
     
     try:
-        # Flush assigns the ID and triggers the UNIQUE constraint check
         db.flush() 
         
         doc_count = 0
@@ -115,17 +107,12 @@ def create_hub(
                 mime_type="text/plain", size_bytes=len(body.content), checksum_sha256="0"*64, status="pending"
             )
             db.add(new_doc)
-            db.flush() # Get doc ID
+            db.flush() 
             
             try:
                 chunks_data = chunk_text(body.content)
                 texts = [c[0] for c in chunks_data]
-                
-                # Embeddings (Wrapped in observability)
-                embeddings = log_ai_call(
-                    "embed_texts", user_id, new_hub.id, 
-                    embed_texts, texts
-                )
+                embeddings = log_ai_call("embed_texts", user_id, new_hub.id, embed_texts, texts)
                 
                 for i, ((text, tokens), emb) in enumerate(zip(chunks_data, embeddings)):
                     db.add(Chunk(
@@ -135,7 +122,6 @@ def create_hub(
                     chunk_count += 1
                 new_doc.status = "ready"
             except Exception:
-                # Graceful degradation: AI failed, but we keep the hub
                 new_doc.status = "failed"
                 new_doc.error_message = "Embedding failed"
 
@@ -153,11 +139,8 @@ def create_hub(
     except IntegrityError as e:
         db.rollback()
         release_claim(db, user_id, key)
-        
-        # Check if the error was our specific unique constraint
         if getattr(e.orig.diag, "constraint_name", None) == "uq_hub_name_per_owner":
             raise AppError(409, "HUB_NAME_CONFLICT", "You already have a hub with this name.")
-        
         raise
 
 @router.post("/{hub_id}/query", response_model=QueryResponse)
@@ -167,31 +150,15 @@ def query_hub(
     current_user: User = Depends(get_current_user),
     db: Session = Depends(get_db)
 ):
-    # 1. Verify hub ownership
     hub = repo.get_hub_for_owner(db, hub_id, current_user.id)
     if not hub:
         raise AppError(404, "HUB_NOT_FOUND", "Hub not found.")
 
-    # 2. Embed the question
-    query_embedding = embed_texts([req.question])[0]
-
-    # 3. Vector search
-    chunks = repo.retrieve_chunks_for_query(db, current_user.id, hub_id, query_embedding)
-
-    # 4. Check similarity threshold (pgvector cosine distance is 1 - similarity)
-    # If no chunks or the closest chunk isn't similar enough, return ungrounded.
-    valid_chunks = []
-    for chunk in chunks:
-        # Calculate similarity: 1 - cosine_distance
-        # We can't easily get the distance from the ORM object after the fact, so we calculate it.
-        # (Alternatively, we could add the distance to the select statement, but this is simpler).
-        # For now, we'll just assume the top-K are the best we have. 
-        # (If we want to be strict, we'd do raw SQL to get the score. We'll use raw SQL here to be safe).
-        pass
-
-    # Let's use raw SQL for exact similarity scoring so we can apply the threshold
-    from sqlalchemy import text
-    from app.config import settings
+    alt_queries = expand_query(req.question)
+    all_queries = [req.question] + alt_queries
+    embeddings = embed_texts(all_queries)
+
+    chunks = repo.retrieve_chunks_for_queries(db, current_user.id, hub_id, embeddings)
     
     raw_sql = text("""
         SELECT id, text, chunk_index, document_id, 1 - (embedding <=> :q) as similarity
@@ -204,10 +171,9 @@ def query_hub(
     
     results = db.execute(
         raw_sql, 
-        {"q": str(query_embedding), "u": current_user.id, "h": hub_id, "threshold": settings.RAG_MIN_SIMILARITY, "k": settings.RAG_TOP_K}
+        {"q": str(embeddings[0]), "u": current_user.id, "h": hub_id, "threshold": settings.RAG_MIN_SIMILARITY, "k": settings.RAG_TOP_K}
     ).fetchall()
 
-    # Release the DB connection before the AI call
     db.rollback()
 
     if not results:
@@ -216,21 +182,15 @@ def query_hub(
             evidence=None, model=settings.GEMINI_MODEL, grounded=False, sources=[], retrieved_chunks=0
         )
 
-    # 5. Assemble context
     context_texts = [row.text for row in results]
     context = "\n\n---\n\n".join(context_texts)
 
-    # 6. Call Gemini
     ai_result = answer_question(context, req.question)
 
-    # 7. Map sources
     sources = []
     if ai_result["evidence"]:
         for row in results:
             if ai_result["evidence"].lower() in row.text.lower():
-                # We need the filename, so we fetch the document (or join in the SQL)
-                # Since we didn't join, we'll just do a quick lookup for the filename
-                from app.models import Document
                 doc = db.query(Document).filter(Document.id == row.document_id).first()
                 if doc:
                     sources.append(QuerySource(

app/services/gemini.py

@@ -84,4 +84,17 @@ def answer_question(content: str, question: str) -> dict:
         "answer": str(d.get("answer", "")).strip(),
         "evidence": evidence if verified else None,
         "grounded": bool(d.get("grounded")) and verified
-    }
+    }
+EXPAND_SYSTEM = (
+    "You are a search assistant. Generate exactly 2 alternative versions of the user's question "
+    "to improve vector search retrieval. Reply with JSON only: "
+    '{"queries": ["alternative 1", "alternative 2"]}'
+)
+
+def expand_query(question: str) -> list[str]:
+    """Returns up to 2 alternative phrasings of the question. Fails gracefully."""
+    try:
+        d = _call(EXPAND_SYSTEM, question, retries=0)
+        return [str(q) for q in d.get("queries", [])][:2]
+    except Exception:
+        return [] # Fallback to just the original question

tests/test_rag.py

@@ -11,7 +11,6 @@ def test_empty_retrieval_returns_ungrounded(client, db_session):
     hub_res = client.post("/api/v1/hubs", headers={**headers, "X-Idempotency-Key": str(uuid4())}, json={"name": "Empty Hub"})
     hub_id = hub_res.json()["id"]
 
-    # Query without any documents/chunks
     with patch("app.routers.hubs.embed_texts", return_value=[[0.1]*768]), \
          patch("app.routers.hubs.answer_question") as mock_answer:
         
@@ -21,21 +20,17 @@ def test_empty_retrieval_returns_ungrounded(client, db_session):
         assert res.json()["grounded"] is False
         assert res.json()["evidence"] is None
         assert res.json()["sources"] == []
-        # CRITICAL: Gemini was never called
         mock_answer.assert_not_called()
 
 def test_vector_search_is_scoped_to_owner(client, db_session):
-    # Create User A and User B
     client.post("/api/v1/auth/signup", json={"email": "usera@rag.com", "password": "Password123"})
     token_a = client.post("/api/v1/auth/login", json={"email": "usera@rag.com", "password": "Password123"}).json()["access_token"]
     
     client.post("/api/v1/auth/signup", json={"email": "userb@rag.com", "password": "Password123"})
     token_b = client.post("/api/v1/auth/login", json={"email": "userb@rag.com", "password": "Password123"}).json()["access_token"]
 
-    # User A creates a hub and a document (we mock the ingestion for speed)
     hub_a = client.post("/api/v1/hubs", headers={"Authorization": f"Bearer {token_a}", "X-Idempotency-Key": str(uuid4())}, json={"name": "A's Secret Hub"}).json()
     
-    # Insert a chunk manually
     from app.models import User, Hub
     user_a = db_session.query(User).filter(User.email == "usera@rag.com").first()
     hub = db_session.query(Hub).filter(Hub.id == hub_a["id"]).first()
@@ -48,8 +43,47 @@ def test_vector_search_is_scoped_to_owner(client, db_session):
     db_session.add(chunk)
     db_session.commit()
 
-    # User B tries to query User A's hub
     with patch("app.routers.hubs.embed_texts", return_value=[[0.1]*768]):
         res = client.post(f"/api/v1/hubs/{hub.id}/query", headers={"Authorization": f"Bearer {token_b}"}, json={"question": "What is the secret?"})
         assert res.status_code == 404
-        assert res.json()["error"]["code"] == "HUB_NOT_FOUND"
+        assert res.json()["error"]["code"] == "HUB_NOT_FOUND"
+
+def test_expansion_fallback_when_gemini_fails(client, db_session):
+    client.post("/api/v1/auth/signup", json={"email": "fallback@test.com", "password": "Password123"})
+    token = client.post("/api/v1/auth/login", json={"email": "fallback@test.com", "password": "Password123"}).json()["access_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+    hub_res = client.post("/api/v1/hubs", headers={**headers, "X-Idempotency-Key": str(uuid4())}, json={"name": "Fallback Hub"})
+    hub_id = hub_res.json()["id"]
+
+    with patch("app.services.gemini._call", side_effect=Exception("AI Down")), \
+         patch("app.routers.hubs.embed_texts", return_value=[[0.1]*768]), \
+         patch("app.routers.hubs.answer_question", return_value={"answer": "Test", "evidence": None, "grounded": False}):
+        
+        res = client.post(f"/api/v1/hubs/{hub_id}/query", headers=headers, json={"question": "Test question?"})
+        assert res.status_code == 200
+
+def test_multihop_question_returns_two_sources(client, db_session):
+    client.post("/api/v1/auth/signup", json={"email": "multihop@test.com", "password": "Password123"})
+    token = client.post("/api/v1/auth/login", json={"email": "multihop@test.com", "password": "Password123"}).json()["access_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+    hub_res = client.post("/api/v1/hubs", headers={**headers, "X-Idempotency-Key": str(uuid4())}, json={"name": "MultiHop Hub"})
+    hub_id = hub_res.json()["id"]
+
+    from app.models import User, Hub, Document, Chunk
+    user = db_session.query(User).filter(User.email == "multihop@test.com").first()
+    hub = db_session.query(Hub).filter(Hub.id == hub_id).first()
+    doc = Document(id=uuid4(), hub_id=hub.id, owner_id=user.id, origin="upload", filename="test.txt", mime_type="text/plain", size_bytes=10, checksum_sha256="0"*64, status="ready")
+    db_session.add(doc)
+    db_session.commit()
+    
+    chunk1 = Chunk(id=uuid4(), document_id=doc.id, hub_id=hub.id, owner_id=user.id, chunk_index=0, text="Chunk one", token_count=5, embedding=[0.1]*768)
+    chunk2 = Chunk(id=uuid4(), document_id=doc.id, hub_id=hub.id, owner_id=user.id, chunk_index=1, text="Chunk two", token_count=5, embedding=[0.2]*768)
+    db_session.add_all([chunk1, chunk2])
+    db_session.commit()
+
+    with patch("app.routers.hubs.expand_query", return_value=["alt1", "alt2"]), \
+         patch("app.routers.hubs.embed_texts", return_value=[[0.1]*768, [0.1]*768, [0.2]*768]), \
+         patch("app.routers.hubs.answer_question", return_value={"answer": "Combined", "evidence": "Chunk", "grounded": True}):
+        
+        res = client.post(f"/api/v1/hubs/{hub_id}/query", headers=headers, json={"question": "Question spanning two chunks?"})
+        assert res.status_code == 200