Add token count and explanation

app.py

@@ -18,6 +18,7 @@ Architecture:
 import os
 import json
 import logging
+import time
 from pathlib import Path
 from typing import Optional
 
@@ -256,12 +257,14 @@ def retrieve_relevant_context(query: str, top_k: int = TOP_K_RESULTS) -> list[di
     return contexts
 
 
-def call_llm(prompt: str) -> str:
+def call_llm(prompt: str) -> dict:
     """
     Make a call to the LLM via Azure Foundry (GPT-5).
 
     Calls the Azure OpenAI-compatible chat/completions endpoint.
     Endpoint URL is loaded from config.json; API key from AZURE_API_KEY env var.
+
+    Returns a dict with 'content' (str) and 'total_tokens' (int).
     """
     headers = {
         "api-key": AZURE_API_KEY,
@@ -278,7 +281,9 @@ def call_llm(prompt: str) -> str:
         resp = http_requests.post(LLM_ENDPOINT_URL, headers=headers, json=payload, timeout=60)
         resp.raise_for_status()
         data = resp.json()
-        return data["choices"][0]["message"]["content"].strip()
+        content = data["choices"][0]["message"]["content"].strip()
+        total_tokens = data.get("usage", {}).get("total_tokens", 0)
+        return {"content": content, "total_tokens": total_tokens}
     except http_requests.exceptions.HTTPError as e:
         logger.error(f"LLM API call failed: {e} — {resp.text}")
         raise HTTPException(status_code=503, detail=f"LLM service unavailable: {str(e)}")
@@ -318,29 +323,54 @@ def rag_query(query: str, top_k: int = TOP_K_RESULTS) -> dict:
     2. Retrieve relevant context from vector store
     3. Build augmented prompt with context
     4. Call LLM to generate answer
-    5. Return answer with source metadata
+    5. Return answer with source metadata, explanation, token count, and timing
     """
+    start_time = time.perf_counter()
+
     # Step 1: Retrieve relevant document chunks
     contexts = retrieve_relevant_context(query, top_k=top_k)
 
     if not contexts:
+        elapsed_ms = round((time.perf_counter() - start_time) * 1000, 2)
         return {
             "answer": "No documents have been ingested yet. Please upload documents first.",
             "sources": [],
-            "query": query,
+            "explanation": "No documents found in the vector store to retrieve context from.",
+            "total_token": 0,
+            "run_time_in_ms": elapsed_ms,
         }
 
     # Step 2: Build the augmented prompt
     prompt = build_rag_prompt(query, contexts)
 
     # Step 3: Generate answer from LLM
-    answer = call_llm(prompt)
+    llm_result = call_llm(prompt)
+    raw_content = llm_result["content"]
+    total_token = llm_result["total_tokens"]
+
+    # Parse structured JSON response from LLM (handle markdown code fences)
+    json_str = raw_content.strip()
+    if json_str.startswith("```"):
+        json_str = json_str.split("\n", 1)[-1]
+        json_str = json_str.rsplit("```", 1)[0].strip()
+
+    try:
+        parsed = json.loads(json_str)
+        answer = parsed["answer"]
+        explanation = parsed["explanation"]
+    except (json.JSONDecodeError, KeyError):
+        answer = raw_content
+        explanation = "LLM did not return a structured explanation."
+
+    elapsed_ms = round((time.perf_counter() - start_time) * 1000, 2)
 
     # Step 4: Return structured response
     return {
         "answer": answer,
         "sources": [{"source": ctx["source"], "score": ctx["similarity_score"]} for ctx in contexts],
-        "query": query,
+        "explanation": explanation,
+        "total_token": total_token,
+        "run_time_in_ms": elapsed_ms,
     }
 
 
@@ -421,7 +451,9 @@ async def query_endpoint(request: QueryRequest):
     Returns:
         - answer (str): The generated answer
         - sources (list): Source documents used for context
-        - query (str): The original query
+        - explanation (str): Explanation of the retrieval and answer logic
+        - total_token (int): Total token count from the LLM call
+        - run_time_in_ms (float): Pipeline execution time in milliseconds
     """
     logger.info(f"Query received: {request.query}")
     result = rag_query(request.query, top_k=request.top_k)
@@ -463,15 +495,19 @@ async def health_check():
 # Step 6: Gradio UI for Interactive Demo
 # ---------------------------------------------------------------------------
 
-def gradio_query(question: str) -> str:
+def gradio_query(question: str) -> tuple[str, str, str, str]:
     """Handle queries from the Gradio chat interface."""
     if not question.strip():
-        return "Please enter a question."
+        return "Please enter a question.", "", "", ""
     result = rag_query(question)
     sources_text = "\n".join(
         f"  - {s['source']} (relevance: {s['score']:.2f})" for s in result["sources"]
     )
-    return f"{result['answer']}\n\n📚 Sources:\n{sources_text}" if result["sources"] else result["answer"]
+    answer = f"{result['answer']}\n\n📚 Sources:\n{sources_text}" if result["sources"] else result["answer"]
+    explanation = result.get("explanation", "")
+    token_info = str(result.get("total_token", 0))
+    run_time = f"{result.get('run_time_in_ms', 0)} ms"
+    return answer, explanation, token_info, run_time
 
 
 def gradio_ingest(text: str, source_name: str) -> str:
@@ -504,7 +540,15 @@ with gr.Blocks(title="RAG Chat API - Gustave Eiffel Hackathon") as demo:
             )
         query_button = gr.Button("Ask", variant="primary")
         query_output = gr.Textbox(label="Answer", lines=8, interactive=False)
-        query_button.click(fn=gradio_query, inputs=query_input, outputs=query_output)
+        query_explanation = gr.Textbox(label="Explanation", lines=3, interactive=False)
+        with gr.Row():
+            query_tokens = gr.Textbox(label="Total Tokens", interactive=False)
+            query_runtime = gr.Textbox(label="Run Time", interactive=False)
+        query_button.click(
+            fn=gradio_query,
+            inputs=query_input,
+            outputs=[query_output, query_explanation, query_tokens, query_runtime],
+        )
 
     with gr.Tab("📄 Ingest Documents"):
         gr.Markdown("Add new documents to the knowledge base.")

prompts/rag_prompt.txt

@@ -2,6 +2,12 @@ You are a helpful assistant. Answer the user's question based ONLY on the provid
 If the context does not contain enough information to answer, say "I don't have enough information to answer this question."
 Always be concise and factual.
 
+You MUST respond in valid JSON with exactly two fields:
+- "answer": your concise factual answer to the question
+- "explanation": a brief explanation of your reasoning, describing which parts of the context you used and why they led to your answer
+
+Do NOT include any text outside the JSON object.
+
 Context:
 {context}