Update app.py

app.py

@@ -55,8 +55,8 @@ DATA_DIR = Path("/data_V5") if Path("/data_V5").is_dir() else Path("./data_V5")
 CHROMA_PERSIST_DIR = str(DATA_DIR / "chroma_db")
 TRAIN_DOCS_DIR = Path("./train_data")
 COLLECTION_NAME = "rag_documents"
-CHUNK_SIZE = 1200
-CHUNK_OVERLAP = 200
+CHUNK_SIZE = 512
+CHUNK_OVERLAP = 50
 TOP_K_RESULTS = 3
 
 # Settings loaded from data/config.json (runtime config).
@@ -76,8 +76,11 @@ EMBEDDING_ENDPOINT_URL = _config["embedding"]["endpoint_url"]
 EMBEDDING_MODEL_NAME = _config["embedding"]["model"]
 
 # LLM (Azure OpenAI)
-SMALL_LLM = _config["llm_small"]
-LARGE_LLM = _config["llm_large"]
+LLM_ENDPOINT_URL = _config["llm"]["endpoint_url"]
+LLM_MODEL_NAME = _config["llm"]["model"]
+LLM_MAX_TOKENS = _config["llm"].get("max_completion_tokens", 512)
+LLM_TEMPERATURE = _config["llm"].get("temperature", 0.7)
+LLM_TOP_P = _config["llm"].get("top_p", 0.95)
 
 # Azure API key from environment variable (shared by both LLM and embedding endpoints)
 AZURE_API_KEY = os.environ.get("AZURE_API_KEY")
@@ -119,7 +122,7 @@ logger.info(f"ChromaDB collection '{COLLECTION_NAME}' ready. Documents: {collect
 # We call the Azure OpenAI-compatible endpoint directly via requests.
 # Endpoint URL and model are configured in config.json.
 
-logger.info(f"LLMs configurés : small={SMALL_LLM['model']} / large={LARGE_LLM['model']}")
+logger.info(f"LLM configured: {LLM_MODEL_NAME} via {LLM_ENDPOINT_URL}")
 
 
 # ---------------------------------------------------------------------------
@@ -272,34 +275,23 @@ def build_rag_prompt(query: str, contexts: list[dict]) -> str:
     prompt = RAG_PROMPT_TEMPLATE.format(context=context_text, question=query)
     return prompt
 
-def _call_model(cfg: dict, prompt: str) -> dict:
-    return call_llm_with_metrics(
-        prompt,
-        endpoint_url=cfg["endpoint_url"],
-        api_key=AZURE_API_KEY,
-        model=cfg["model"],
-        max_completion_tokens=cfg["max_completion_tokens"],
-        temperature=cfg["temperature"],
-        top_p=cfg["top_p"],
-    )
 
+def rag_query(query: str, top_k: int = TOP_K_RESULTS) -> dict:
+    """
+    End-to-end RAG pipeline: Query → Retrieve → Generate.
 
-def _parse_llm_json(raw_content: str):
-    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.get("explanation", "")
-        confident = bool(parsed.get("confident", True))  # absent => on suppose confiant (pas d'escalade)
-        return answer, explanation, confident
-    except (json.JSONDecodeError, KeyError):
-        return raw_content, "LLM did not return a structured explanation.", True
+    This demonstrates Step 2d: How to make a RAG system end-to-end.
 
-def rag_query(query: str, top_k: int = TOP_K_RESULTS) -> dict:
+    Pipeline steps:
+    1. Receive user query
+    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, 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:
@@ -308,50 +300,59 @@ def rag_query(query: str, top_k: int = TOP_K_RESULTS) -> dict:
             "answer": "No documents have been ingested yet. Please upload documents first.",
             "sources": [],
             "explanation": "No documents found in the vector store to retrieve context from.",
-            "total_token": 0, "prompt_tokens": 0, "completion_tokens": 0, "cached_tokens": 0,
-            "co2_grams": None, "energy_kwh": None,
+            "total_token": 0,
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "cached_tokens": 0,
+            "co2_grams": None,
+            "energy_kwh": None,
             "run_time_in_ms": elapsed_ms,
         }
 
+    # Step 2: Build the augmented prompt
     prompt = build_rag_prompt(query, contexts)
 
-    # accumulateurs : on additionne le coût de TOUS les appels
-    tot = {"total": 0, "prompt": 0, "completion": 0, "cached": 0}
-    co2, energy = 0.0, 0.0
-
-    def _accumulate(res):
-        nonlocal co2, energy
-        for k in tot:
-            tot[k] += res["tokens"][k]
-        co2 += res["co2_grams"] or 0
-        energy += res["energy_kwh"] or 0
-
-    # 1) PETIT modèle d'abord
-    res_small = _call_model(SMALL_LLM, prompt)
-    _accumulate(res_small)
-    answer, explanation, confident = _parse_llm_json(res_small["content"])
-    model_used = SMALL_LLM["model"]
-
-    # 2) Escalade vers le GROS modèle uniquement si pas confiant
-    if not confident:
-        res_large = _call_model(LARGE_LLM, prompt)
-        _accumulate(res_large)
-        answer, explanation, _ = _parse_llm_json(res_large["content"])
-        model_used = LARGE_LLM["model"]
-
-    logger.info(f"Modèle utilisé : {model_used} (confiance petit mod��le : {confident})")
+    # Step 3: Generate answer from LLM (with token + CO2 metrics)
+    llm_result = call_llm_with_metrics(
+        prompt,
+        endpoint_url=LLM_ENDPOINT_URL,
+        api_key=AZURE_API_KEY,
+        model=LLM_MODEL_NAME,
+        max_completion_tokens=LLM_MAX_TOKENS,
+        temperature=LLM_TEMPERATURE,
+        top_p=LLM_TOP_P,
+    )
+    raw_content = llm_result["content"]
+    tokens = llm_result["tokens"]
+    total_token = tokens["total"]
+
+    # 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": c["source"], "score": c["similarity_score"], "ref_text": c["text"]} for c in contexts],
+        "sources": [{"source": ctx["source"], "score": ctx["similarity_score"], "ref_text": ctx["text"]} for ctx in contexts],
         "explanation": explanation,
-        "total_token": tot["total"],
-        "prompt_tokens": tot["prompt"],
-        "completion_tokens": tot["completion"],
-        "cached_tokens": tot["cached"],
-        "co2_grams": co2 if co2 else None,
-        "energy_kwh": energy if energy else None,
+        "total_token": total_token,
+        "prompt_tokens": tokens["prompt"],
+        "completion_tokens": tokens["completion"],
+        "cached_tokens": tokens["cached"],
+        "co2_grams": llm_result["co2_grams"],
+        "energy_kwh": llm_result["energy_kwh"],
         "run_time_in_ms": elapsed_ms,
     }
 
@@ -470,7 +471,7 @@ async def health_check():
         "status": "healthy",
         "documents_in_store": collection.count(),
         "embedding_model": EMBEDDING_MODEL_NAME,
-        "llm_models": [SMALL_LLM["model"], LARGE_LLM["model"]],
+        "llm_model": LLM_MODEL_NAME,
     }