feat: add cited answers and configurable retrieval

app/engine/context_builder.py

@@ -0,0 +1,58 @@
+from langchain_core.documents import Document
+
+from app.core.config import settings
+
+
+def build_context(documents: list[Document]) -> str:
+    parts: list[str] = []
+    total_chars = 0
+    for document in dedupe_documents(documents):
+        source = document.metadata.get("source", "unknown")
+        page = document.metadata.get("page")
+        label = f"{source}, page {page}" if page else source
+        content = document.page_content.strip()
+        block = f"Source: {label}\n{content}"
+        if total_chars + len(block) > settings.MAX_CONTEXT_CHARS:
+            break
+        parts.append(block)
+        total_chars += len(block)
+    return "\n\n".join(parts)
+
+
+def dedupe_documents(documents: list[Document]) -> list[Document]:
+    seen: set[str] = set()
+    unique: list[Document] = []
+    for document in documents:
+        key = document.metadata.get("chunk_id") or document.page_content[:120]
+        if key in seen:
+            continue
+        seen.add(key)
+        unique.append(document)
+    return unique
+
+
+def source_citations(documents: list[Document]) -> list[dict]:
+    citations: list[dict] = []
+    for document in dedupe_documents(documents):
+        snippet = " ".join(document.page_content.split())[:240]
+        citations.append(
+            {
+                "source": document.metadata.get("source", "unknown"),
+                "page": document.metadata.get("page"),
+                "chunk_id": document.metadata.get("chunk_id"),
+                "doc_id": document.metadata.get("doc_id"),
+                "snippet": snippet,
+            }
+        )
+    return citations
+
+
+def format_sources(documents: list[Document]) -> str:
+    citations = source_citations(documents)
+    if not citations:
+        return ""
+    lines = ["", "Sources:"]
+    for citation in citations:
+        page = f", page {citation['page']}" if citation.get("page") else ""
+        lines.append(f"- {citation['source']}{page}: {citation['snippet']}")
+    return "\n".join(lines)

app/engine/query_transform.py

@@ -0,0 +1,15 @@
+import re
+
+
+def normalize_query(query: str) -> str:
+    return re.sub(r"\s+", " ", query).strip()
+
+
+def query_variants(query: str) -> list[str]:
+    normalized = normalize_query(query)
+    variants = [normalized]
+    if "error" in normalized.lower() and "troubleshoot" not in normalized.lower():
+        variants.append(f"troubleshoot {normalized}")
+    if "how" in normalized.lower() and "steps" not in normalized.lower():
+        variants.append(f"{normalized} steps")
+    return list(dict.fromkeys(variants))

app/engine/retriever.py

@@ -1,31 +1,47 @@
-from langchain_community.embeddings import FastEmbedEmbeddings
-from langchain_qdrant import FastEmbedSparse, QdrantVectorStore, RetrievalMode
-from langchain_community.cross_encoders import HuggingFaceCrossEncoder
-from langchain.retrievers.document_compressors import CrossEncoderReranker
+import logging
+
 from langchain.retrievers import ContextualCompressionRetriever
-from qdrant_client import QdrantClient
+from langchain.retrievers.document_compressors import CrossEncoderReranker
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 
 from app.core.config import settings
+from app.engine.indexer import open_vector_store
+from app.engine.query_transform import query_variants
 
-def get_reranked_retriever():
-    if settings.QDRANT_URL:
-        client = QdrantClient(url=settings.QDRANT_URL)
-    else:
-        client = QdrantClient(path=settings.QDRANT_LOCATION)
-        
-    dense_embeddings = FastEmbedEmbeddings(model_name=settings.DENSE_EMBEDDING_MODEL)
-    sparse_embeddings = FastEmbedSparse(model_name=settings.SPARSE_EMBEDDING_MODEL)
-
-    qdrant = QdrantVectorStore(
-        client=client,
-        collection_name=settings.COLLECTION_NAME,
-        embedding=dense_embeddings,
-        sparse_embedding=sparse_embeddings,
-        retrieval_mode=RetrievalMode.HYBRID,
-    )
+logger = logging.getLogger(__name__)
 
-    base_retriever = qdrant.as_retriever(search_kwargs={"k": 15})
-    model = HuggingFaceCrossEncoder(model_name=settings.RERANKER_MODEL)
-    compressor = CrossEncoderReranker(model=model, top_n=3)
 
+def get_retriever():
+    qdrant = open_vector_store()
+    base_retriever = qdrant.as_retriever(search_kwargs={"k": settings.RETRIEVAL_TOP_K})
+
+    if not settings.RERANKER_ENABLED:
+        return base_retriever
+
+    model = HuggingFaceCrossEncoder(model_name=settings.RERANKER_MODEL)
+    compressor = CrossEncoderReranker(model=model, top_n=settings.RERANKER_TOP_N)
     return ContextualCompressionRetriever(base_compressor=compressor, base_retriever=base_retriever)
+
+
+def retrieve_documents(question: str):
+    retriever = get_retriever()
+    documents = []
+    seen = set()
+    for query in query_variants(question):
+        for document in retriever.invoke(query):
+            key = document.metadata.get("chunk_id") or document.page_content[:120]
+            if key in seen:
+                continue
+            seen.add(key)
+            documents.append(document)
+    logger.info(
+        "retrieval completed query_count=%s returned_chunks=%s reranker_enabled=%s",
+        len(query_variants(question)),
+        len(documents),
+        settings.RERANKER_ENABLED,
+    )
+    return documents[: settings.RETRIEVAL_TOP_K]
+
+
+def get_reranked_retriever():
+    return get_retriever()

app/graph/workflow.py

@@ -1,17 +1,19 @@
 import json
-from typing import List, TypedDict
+from typing import List, Optional, TypedDict
 from langchain_core.prompts import PromptTemplate
 from langchain_core.documents import Document
 from langchain_ollama import ChatOllama
 from langgraph.graph import START, END, StateGraph
 
 from app.core.config import settings
-from app.engine.retriever import get_reranked_retriever
+from app.engine.context_builder import build_context, source_citations
+from app.engine.retriever import retrieve_documents
 
 class GraphState(TypedDict):
     question: str
     generation: str
     documents: List[Document]
+    sources: Optional[list[dict]]
     run_count: int
 
 llm = ChatOllama(model=settings.OLLAMA_MODEL, temperature=0, base_url=settings.OLLAMA_BASE_URL)
@@ -21,9 +23,8 @@ def retrieve(state: GraphState):
     print("--- NODE: RETRIEVE DOCS ---")
     question = state["question"]
     run_count = state.get("run_count", 0)
-    retriever = get_reranked_retriever()
-    documents = retriever.invoke(question)
-    return {"documents": documents, "question": question, "run_count": run_count}
+    documents = retrieve_documents(question)
+    return {"documents": documents, "sources": source_citations(documents), "question": question, "run_count": run_count}
 
 def grade_documents(state: GraphState):
     print("--- NODE: GRADE DOCUMENT RELEVANCE ---")
@@ -58,9 +59,9 @@ def generate(state: GraphState):
     documents = state["documents"]
     run_count = state.get("run_count", 0) + 1
     
-    context = "\n\n".join(doc.page_content for doc in documents)
+    context = build_context(documents)
     prompt = PromptTemplate(
-        template="""You are a Support Docs Copilot. Use the retrieved context to answer the question concisely. If you don't know the answer, say "I don't know".
+        template="""You are a Support Docs Copilot. Use only the retrieved context to answer the question concisely. If the context does not contain the answer, say "I don't know".
         Question: {question} 
         Context: {context} 
         Answer:""",
@@ -68,7 +69,7 @@ def generate(state: GraphState):
     )
     rag_chain = prompt | llm
     generation = rag_chain.invoke({"context": context, "question": question})
-    return {"generation": generation.content, "run_count": run_count}
+    return {"generation": generation.content, "sources": source_citations(documents), "run_count": run_count}
 
 def decide_to_generate(state: GraphState):
     if not state["documents"]:
@@ -86,7 +87,7 @@ def check_hallucinations(state: GraphState):
         print("--- ROUTE: MAX RETRIES REACHED ---")
         return "end"
         
-    context = "\n\n".join(doc.page_content for doc in documents)
+    context = build_context(documents)
     prompt = PromptTemplate(
         template="""You are evaluating whether a generated answer is fully grounded in the retrieved facts.
         Facts: \n\n {context} \n\n

app/main.py

@@ -10,6 +10,7 @@ from langchain_ollama import ChatOllama
 from app.core.config import settings
 from app.core.dependencies import check_ollama, check_qdrant
 from app.core.logging import configure_logging, logger
+from app.engine.context_builder import build_context, format_sources
 from app.graph.workflow import compile_workflow
 from app.guardrails.validators import DetectPromptInjection
 
@@ -21,9 +22,17 @@ input_guard = Guard().use(DetectPromptInjection, on_fail="exception")
 class ChatRequest(BaseModel):
     query: str
 
+class SourceCitation(BaseModel):
+    source: str
+    page: int | None = None
+    chunk_id: str | None = None
+    doc_id: str | None = None
+    snippet: str
+
 class ChatResponse(BaseModel):
     query: str
     answer: str
+    sources: list[SourceCitation] = []
 
 @app.get("/health")
 async def health_endpoint():
@@ -42,6 +51,8 @@ async def ready_endpoint():
 @app.post("/chat", response_model=ChatResponse)
 async def chat_endpoint(request: ChatRequest):
     started_at = time.perf_counter()
+    if len(request.query) > settings.MAX_QUERY_LENGTH:
+        raise HTTPException(status_code=400, detail="Query is too long.")
     if settings.ENABLE_GUARDRAILS:
         try:
             input_guard.validate(request.query)
@@ -52,14 +63,17 @@ async def chat_endpoint(request: ChatRequest):
     try:
         final_state = rag_agent.invoke(initial_state)
         answer = final_state.get("generation", "Unable to compile answer.")
+        sources = final_state.get("sources", [])
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
     logger.info("chat completed latency_ms=%s", round((time.perf_counter() - started_at) * 1000, 2))
-    return ChatResponse(query=request.query, answer=answer)
+    return ChatResponse(query=request.query, answer=answer, sources=sources)
 
 @app.post("/chat/stream")
 async def chat_stream_endpoint(request: ChatRequest):
+    if len(request.query) > settings.MAX_QUERY_LENGTH:
+        raise HTTPException(status_code=400, detail="Query is too long.")
     if settings.ENABLE_GUARDRAILS:
         try:
             input_guard.validate(request.query)
@@ -76,9 +90,9 @@ async def chat_stream_endpoint(request: ChatRequest):
             yield "I am sorry, no reliable matching documentation was found."
             return
 
-        context = "\n\n".join(doc.page_content for doc in documents)
+        context = build_context(documents)
         prompt = PromptTemplate(
-            template="""You are a Support Docs Copilot. Use the retrieved context to answer the question concisely. If you don't know the answer, say "I don't know".
+            template="""You are a Support Docs Copilot. Use only the retrieved context to answer the question concisely. If you don't know the answer, say "I don't know".
             Question: {question} 
             Context: {context} \n\nAnswer:""",
             input_variables=["question", "context"],
@@ -90,6 +104,7 @@ async def chat_stream_endpoint(request: ChatRequest):
             if chunk.content:
                 yield chunk.content
                 await asyncio.sleep(0.01)
+        yield format_sources(documents)
         logger.info("stream completed latency_ms=%s", round((time.perf_counter() - started_at) * 1000, 2))
 
     return StreamingResponse(token_generator(), media_type="text/event-stream")