Update app/llm_chain.py

app/llm_chain.py

@@ -3,17 +3,17 @@ llm_chain.py — LangChain LCEL pipeline for multi-LLM code debugging.
 
 Flow:
   1. Build one LCEL chain per panel model  (PANEL_PROMPT | llm | StrOutputParser)
-  2. Run all 3 chains concurrently with asyncio.gather  (ainvoke)
-  3. Feed all results into the judge chain               (JUDGE_PROMPT | llm | StrOutputParser)
-  4. Parse judge output into {reasoning, final_answer}
-  5. Return structured dict consumed by the FastAPI endpoint
+  2. Run both chains concurrently          (asyncio.gather + ainvoke)
+  3. Feed results into judge chain         (JUDGE_PROMPT | llm | StrOutputParser)
+  4. Parse judge output into {error_analysis, corrected_code}
+  5. Return structured dict consumed by FastAPI
 """
 import asyncio
+import re
 import time
 from typing import Any
 
 from langchain_core.output_parsers import StrOutputParser
-from langchain_core.runnables import RunnablePassthrough
 
 from app.config import settings
 from app.llm_factory import make_panel_llm, make_judge_llm
@@ -22,22 +22,30 @@ from app.prompts import PANEL_PROMPT, JUDGE_PROMPT
 
 # ── Helpers ───────────────────────────────────────────────────────────────────
 
-def _split_judge_output(raw: str) -> tuple[str, str]:
+def _parse_judge_output(raw: str) -> tuple[str, str]:
     """
-    Split judge output on the markdown headers we asked for.
-    Returns (reasoning, final_answer).
+    Parse judge response into (error_analysis, corrected_code).
+    Splits on the exact section headers defined in JUDGE_PROMPT.
     """
-    reasoning    = ""
-    final_answer = raw.strip()
+    error_analysis = ""
+    corrected_code = ""
 
-    if "## ✅ Final Answer" in raw:
-        parts        = raw.split("## ✅ Final Answer", 1)
-        final_answer = parts[1].strip()
-        reasoning    = parts[0].replace("## 🔍 Reasoning", "").strip()
-    elif "## 🔍 Reasoning" in raw:
-        reasoning    = raw.split("## 🔍 Reasoning", 1)[1].strip()
+    if "## ✅ Corrected Code" in raw:
+        parts          = raw.split("## ✅ Corrected Code", 1)
+        corrected_code = parts[1].strip()
+        if "## 🐛 Error Analysis" in parts[0]:
+            error_analysis = parts[0].split("## 🐛 Error Analysis", 1)[1].strip()
+        else:
+            error_analysis = parts[0].strip()
+    elif "## 🐛 Error Analysis" in raw:
+        error_analysis = raw.split("## 🐛 Error Analysis", 1)[1].strip()
+        corrected_code = raw  # fallback: treat whole thing as corrected code
 
-    return reasoning, final_answer
+    # Strip leading/trailing whitespace from each
+    error_analysis = error_analysis.strip()
+    corrected_code = corrected_code.strip()
+
+    return error_analysis, corrected_code
 
 
 def _error_panel(model: str, label: str, error: str, latency_ms: float) -> dict:
@@ -54,25 +62,20 @@ def _error_panel(model: str, label: str, error: str, latency_ms: float) -> dict:
 
 def build_panel_chain(model: str, temperature: float):
     """
-    LCEL chain for a single panel model:
-        PANEL_PROMPT | ChatOpenAI | StrOutputParser
+    LCEL chain:  PANEL_PROMPT | ChatOpenAI | StrOutputParser
     Input:  {"question": str}
-    Output: str  (model's raw markdown answer)
+    Output: str
     """
     llm = make_panel_llm(model, temperature=temperature)
     return PANEL_PROMPT | llm | StrOutputParser()
 
 
 async def _run_panel(model: str, label: str, question: str, temperature: float) -> dict:
-    """
-    Invoke one panel chain, return structured result dict.
-    Errors are caught so one failing model doesn't abort everything.
-    """
     chain = build_panel_chain(model, temperature)
     start = time.perf_counter()
     try:
-        response  = await chain.ainvoke({"question": question})
-        latency   = round((time.perf_counter() - start) * 1000, 1)
+        response = await chain.ainvoke({"question": question})
+        latency  = round((time.perf_counter() - start) * 1000, 1)
         return {
             "model":      model,
             "label":      label,
@@ -89,35 +92,25 @@ async def _run_panel(model: str, label: str, question: str, temperature: float)
 
 def build_judge_chain():
     """
-    LCEL chain for the Qwen3 judge:
-        JUDGE_PROMPT | ChatOpenAI | StrOutputParser
-    Input:  {"question", "label_1", "response_1", "label_2", "response_2",
-              "label_3", "response_3"}
-    Output: str  (full judge markdown text)
+    LCEL chain:  JUDGE_PROMPT | ChatOpenAI | StrOutputParser
+    Input:  {"question", "label_1", "response_1", "label_2", "response_2"}
+    Output: str
     """
     llm = make_judge_llm()
     return JUDGE_PROMPT | llm | StrOutputParser()
 
 
 async def _run_judge(question: str, panel_results: list[dict]) -> dict:
-    """
-    Build judge inputs from panel results, invoke the judge chain,
-    and parse the structured output.
-    """
     chain = build_judge_chain()
     start = time.perf_counter()
 
-    # Pad missing responses so the prompt always has 3 slots
-    padded = []
-    for r in panel_results:
-        padded.append(r)
-    while len(padded) < 3:
+    # Always provide exactly 2 slots
+    padded = list(panel_results)
+    while len(padded) < 2:
         padded.append({"label": "N/A", "response": "[No response — model failed]"})
 
     def _content(r: dict) -> str:
-        if r.get("response"):
-            return r["response"]
-        return f"[Error: {r.get('error', 'unknown')}]"
+        return r["response"] if r.get("response") else f"[Error: {r.get('error', 'unknown')}]"
 
     judge_input = {
         "question":   question,
@@ -125,55 +118,43 @@ async def _run_judge(question: str, panel_results: list[dict]) -> dict:
         "response_1": _content(padded[0]),
         "label_2":    padded[1]["label"],
         "response_2": _content(padded[1]),
-        "label_3":    padded[2]["label"],
-        "response_3": _content(padded[2]),
     }
 
     try:
         raw     = await chain.ainvoke(judge_input)
         latency = round((time.perf_counter() - start) * 1000, 1)
-        reasoning, final_answer = _split_judge_output(raw)
+        error_analysis, corrected_code = _parse_judge_output(raw)
         return {
-            "model":         settings.JUDGE_MODEL,
-            "label":         settings.JUDGE_LABEL,
-            "reasoning":     reasoning,
-            "final_answer":  final_answer,
-            "full_response": raw.strip(),
-            "latency_ms":    latency,
-            "error":         None,
+            "model":          settings.JUDGE_MODEL,
+            "label":          settings.JUDGE_LABEL,
+            "error_analysis": error_analysis,
+            "corrected_code": corrected_code,
+            "full_response":  raw.strip(),
+            "latency_ms":     latency,
+            "error":          None,
         }
     except Exception as exc:
         latency = round((time.perf_counter() - start) * 1000, 1)
         return {
-            "model":         settings.JUDGE_MODEL,
-            "label":         settings.JUDGE_LABEL,
-            "reasoning":     "",
-            "final_answer":  None,
-            "full_response": None,
-            "latency_ms":    latency,
-            "error":         str(exc),
+            "model":          settings.JUDGE_MODEL,
+            "label":          settings.JUDGE_LABEL,
+            "error_analysis": None,
+            "corrected_code": None,
+            "full_response":  None,
+            "latency_ms":     latency,
+            "error":          str(exc),
         }
 
 
-# ── Main entry point ──────────────────────────────────────────────────────────
+# ── Main entry ────────────────────────────────────────────────────────────────
 
 async def debug_query(question: str, temperature: float = 0.3) -> dict[str, Any]:
     """
-    Full multi-LLM pipeline:
-      Phase 1 — query 3 panel models concurrently (asyncio.gather)
-      Phase 2 — send all results to the Qwen3 judge
-
-    Returns:
-    {
-        "question":  str,
-        "panel":     list[dict],   # one dict per model
-        "judge":     dict,         # reasoning + final_answer
-        "total_ms":  float,
-    }
+    Phase 1 — 2 panel models in parallel (asyncio.gather)
+    Phase 2 — StepFun Flash judge → {error_analysis, corrected_code}
     """
     total_start = time.perf_counter()
 
-    # ── Phase 1: parallel panel calls ────────────────────────────────────────
     panel_tasks = [
         _run_panel(
             model       = model,
@@ -185,14 +166,11 @@ async def debug_query(question: str, temperature: float = 0.3) -> dict[str, Any]
     ]
     panel_results: list[dict] = list(await asyncio.gather(*panel_tasks))
 
-    # ── Phase 2: judge ────────────────────────────────────────────────────────
     judge_result = await _run_judge(question, panel_results)
 
-    total_ms = round((time.perf_counter() - total_start) * 1000, 1)
-
     return {
         "question": question,
         "panel":    panel_results,
         "judge":    judge_result,
-        "total_ms": total_ms,
-    }
+        "total_ms": round((time.perf_counter() - total_start) * 1000, 1),
+    }