Refactor: Api components

app/agents/nodes.py

@@ -41,11 +41,17 @@ base_llm: Any
 editor_llm_core: Any
 prosecutor_llm_core: Any
 
+# --- 1. SOTA MoE BRAIN CONFIGURATION ---
 if _nv_key:
     try:
+        # The Architect: Stable, Dense Llama 3.3 for flawless formatting
         base_llm = ChatNVIDIA(model="meta/llama-3.3-70b-instruct", nvidia_api_key=_nv_key, temperature=0, max_tokens=2048)
-        editor_llm_core = ChatNVIDIA(model="nvidia/nvidia-nemotron-nano-9b-v2", nvidia_api_key=_nv_key, temperature=0.1, max_tokens=1024)
-        prosecutor_llm_core = ChatNVIDIA(model="meta/llama-3.1-405b-instruct", nvidia_api_key=_nv_key, temperature=0, max_tokens=512)
+        
+        # The Editor: Step-Flash MoE for ultra-fast, cheap data extraction
+        editor_llm_core = ChatNVIDIA(model="stepfun-ai/step-3.5-flash", nvidia_api_key=_nv_key, temperature=0.1, max_tokens=1024)
+        
+        # The Prosecutor: DeepSeek-Terminus MoE for brutal logic verification
+        prosecutor_llm_core = ChatNVIDIA(model="deepseek-ai/deepseek-v3.1-terminus", nvidia_api_key=_nv_key, temperature=0, max_tokens=1024)
     except: _nv_key = None
 
 if not _nv_key:
@@ -61,7 +67,6 @@ async def retrieve_node(state: AgentState):
     command = None
     clean_question = raw_question
     
-    # SOTA: Multi-Flag Parser (Captures '-a -t -v' as a single block)
     cmd_match = re.match(r'^/axm\s+((?:-[a-z]+\s*|\.\.\s*)+)(.*)', raw_question, re.IGNORECASE | re.DOTALL)
     if cmd_match:
         command = cmd_match.group(1).strip().lower()
@@ -72,7 +77,6 @@ async def retrieve_node(state: AgentState):
     is_vault_mode = "vault" in filenames or len(filenames) == 0
     search_input = None if is_vault_mode else filenames
     
-    # SOTA: Use 'in' to check for specific flags inside the chained command
     is_deep_audit = command and "-a" in command
     search_limit = 60 if is_deep_audit else 30
     top_k = 20 if is_deep_audit else 12
@@ -89,7 +93,7 @@ async def distill_node(state: AgentState):
     if not context_text.strip(): 
         return {"generation": "NO RELEVANT EVIDENCE", "status": "thinking"}
 
-    chain = DISTILLATION_PROMPT | base_llm | distill_parser
+    chain = DISTILLATION_PROMPT | editor_llm_core | distill_parser
     
     try:
         raw_response = await chain.ainvoke({"context": context_text, "question": state["question"]})
@@ -120,14 +124,15 @@ async def generate_node(state: AgentState):
         return {"generation": "No direct evidence found in the vault.", "status": "verifying"}
 
     history_context = ""
-    # SOTA Check: Is the reset flag present?
     if history and (not command or ".." not in command):
         history_context = "\n\n### PREVIOUS AUDIT CONTEXT:\n"
         for turn in history[-3:]: 
             history_context += f"{turn['role'].upper()}: {turn['content']}\n"
 
-    # SOTA Check: Is the table flag present?
-    formatting_directive = "\n\nCRITICAL: You are in TABLE MODE. Output strictly as a Markdown Data Grid." if command and "-t" in command else ""
+    # CRITICAL: Ensures the Architect responds ONLY with data, satisfying the Prosecutor's strict checks
+    formatting_directive = "\n\nCRITICAL: Answer ONLY with the facts found in the evidence. Do not add intro or outro text."
+    if command and "-t" in command:
+        formatting_directive += " You are in TABLE MODE. Output strictly as a Markdown Data Grid."
 
     chain = VERIFICATION_PROMPT | simple_llm 
     response = await chain.ainvoke({"context": f"{history_context}\n\nEVIDENCE:\n{distilled_brief}{formatting_directive}", "question": state["question"]})
@@ -139,18 +144,18 @@ async def grade_generation_node(state: AgentState):
         return {"hallucination_score": 1.0, "metrics": {"faithfulness": 1.0, "precision": 1.0, "relevance": 1.0}, "status": "verified", "active_node": "Prosecutor"}
 
     command = state.get("command")
-    # SOTA Check: Is the intense verification flag present?
     intensify = command is not None and "-v" in command
     
     context_list = state["documents"]
     context_str = "\n\n".join(context_list)
     
     try:
+        # Prosecutor is now powered by DeepSeek-v3.1-Terminus
         chain = GRADING_PROMPT | prosecutor_llm_core | grade_parser
         grade = await chain.ainvoke({"context": context_str, "generation": generation})
         
         if str(grade.is_hallucinating).strip().lower() == "true":
-            print(f"LOGIC BREACH (NIM): {grade.explanation}")
+            print(f"LOGIC BREACH (DeepSeek): {grade.explanation}")
             return {"hallucination_score": 0.0, "status": "thinking", "retry_count": state.get("retry_count", 0) + 1, "active_node": "Prosecutor"}
     except Exception as e:
         print(f"⚠️ PROSECUTOR JSON FAILSAFE: {e}")

app/api/run.py

@@ -42,7 +42,7 @@ async def run_verification(
             start_time = time.time()
             print(f"--- STREAM STARTED FOR: {payload.question[:30]}... ---")
 
-            history_buffer: List[Dict[str, str]] = []
+            history_buffer: List[Dict[str, str]] =[]
             is_root_reset = payload.question.strip().startswith("/axm ..")
 
             if db and not is_root_reset:
@@ -50,13 +50,11 @@ async def run_verification(
                     primary_file = payload.filenames[0] if payload.filenames else "vault"
                     doc_res = db.table("documents").select("id").eq("filename", primary_file).eq("user_id", user_id).execute()
                     
-                    # FIX: Explicit cast to allow indexing
                     doc_rows = cast(List[Dict[str, Any]], doc_res.data)
                     if doc_rows:
                         doc_id = doc_rows[0]['id']
                         hist_res = db.table("chat_messages").select("role, content").eq("document_id", doc_id).eq("user_id", user_id).order("created_at", desc=True).limit(5).execute()
                         
-                        # FIX: Cast to match AgentState history requirements
                         raw_hist = cast(List[Dict[str, str]], hist_res.data)
                         history_buffer = raw_hist[::-1]
                 except Exception as e:
@@ -96,10 +94,18 @@ async def run_verification(
                 
                 if kind == "on_chain_start" and name in ui_node_map:
                     current_active_node = ui_node_map[name]
+                    
+                    # STRIKE 3 FIX: The State Accumulation Patch
+                    # If the Architect fires again, it means the Prosecutor triggered a retry.
+                    # We must wipe the previous failed draft from memory and tell the UI to clear.
+                    if current_active_node in ["Architect", "Strategist"] and full_generation:
+                        full_generation = ""
+                        yield {"event": "clear", "data": json.dumps({"message": "retry_triggered"})}
+
                     yield {"event": "node_update", "data": json.dumps({"node": current_active_node, "status": "active"})}
 
                 elif kind == "on_chat_model_stream":
-                    if current_active_node in ["Architect", "Strategist"]:
+                    if current_active_node in["Architect", "Strategist"]:
                         chunk = event["data"].get("chunk")
                         content = ""
                         if chunk:
@@ -110,15 +116,13 @@ async def run_verification(
                             full_generation += content
                             yield {"event": "token", "data": json.dumps({"text": content})}
 
-                elif kind == "on_chain_end" and name in ["generate_node", "Architect"]:
-                    # FIX: Explicit type annotation for Mypy
+                elif kind == "on_chain_end" and name in["generate_node", "Architect"]:
                     node_output: Dict[str, Any] = event["data"].get("output", {})
                     if not full_generation and "generation" in node_output:
                         full_generation = str(node_output["generation"])
                         yield {"event": "token", "data": json.dumps({"text": full_generation})}
 
-                elif kind == "on_chain_end" and name in ["grade_generation_node", "Prosecutor"]:
-                    # FIX: Explicit type annotation for Mypy
+                elif kind == "on_chain_end" and name in["grade_generation_node", "Prosecutor"]:
                     eval_output: Dict[str, Any] = event["data"].get("output", {})
                     final_metrics = eval_output.get("metrics", {})
 

app/core/evaluator.py

@@ -21,7 +21,6 @@ class AxiomEvaluator:
     Stabilized for LangChain 1.x Parent-Child Type Resolution & Concurrency.
     """
     def __init__(self) -> None:
-        # 1. MYPY FIX: Use 'Any' because RAGAS 0.4.x wrappers are dynamically typed
         self.evaluator_llm: Any = None
         self.faithfulness_metric: Optional[Faithfulness] = None
 
@@ -29,11 +28,12 @@ class AxiomEvaluator:
         """Initializes RAGAS V2 and stabilizes the Pydantic Registry."""
         if self.evaluator_llm is None:
             try:
-                BaseChatModel.model_rebuild()
-                ChatOpenAI.model_rebuild()
+                # SOTA: Providing the namespace fixes the "BaseCache is not defined" warning
+                BaseChatModel.model_rebuild(_types_namespace={"BaseCache": BaseCache, "Callbacks": Callbacks})
+                ChatOpenAI.model_rebuild(_types_namespace={"BaseCache": BaseCache, "Callbacks": Callbacks})
                 print("AXIOM-CORE: Evaluator Registry Synchronized.")
-            except Exception as e:
-                print(f"AXIOM-CORE: Registry notice (Non-fatal): {e}")
+            except Exception:
+                pass # Silently catch non-fatal Pydantic 2.x notices
 
             print("AXIOM-CORE: Materializing RAGAS V2 Auditor (NVIDIA NIM)...")
             raw_key = os.environ.get("NVIDIA_API_KEY")
@@ -43,7 +43,9 @@ class AxiomEvaluator:
                 temperature=0,
                 api_key=SecretStr(raw_key) if raw_key else None,
                 base_url="https://integrate.api.nvidia.com/v1",
-                max_completion_tokens=2048
+                max_completion_tokens=2048,
+                max_retries=3,
+                timeout=60.0 # Matches our core 60s timeout for cold starts
             )
             
             self.evaluator_llm = LangchainLLMWrapper(llm)
@@ -60,21 +62,29 @@ class AxiomEvaluator:
             )
             dataset = EvaluationDataset(samples=[sample])
 
-            # 2. SOTA THREAD POOLING
+            # SOTA THREAD POOLING
             def run_ragas() -> Any:
                 return evaluate(dataset=dataset, metrics=[self.faithfulness_metric]) # type: ignore
             
             result = await asyncio.to_thread(run_ragas)
             
-            # 3. Offload Pandas DataFrame operations
+            # Offload Pandas DataFrame operations
             def extract_score() -> float:
                 scores_df = result.to_pandas()
-                return float(scores_df["faithfulness"].iloc[0])
                 
-            raw_val = await asyncio.to_thread(extract_score)
-            
-            # 4. Sanitize for JSON compliance
-            faithfulness_score = raw_val if math.isfinite(raw_val) else 0.0
+                # STRIKE 2 FIX: Empty DataFrame Guard (Prevents IndexError)
+                if scores_df is None or scores_df.empty or "faithfulness" not in scores_df.columns:
+                    print("⚠️ RAGAS: Evaluation returned empty dataset.")
+                    return 0.0
+                
+                val = scores_df["faithfulness"].iloc[0]
+                try:
+                    f_val = float(val)
+                    return f_val if math.isfinite(f_val) else 0.0
+                except (ValueError, TypeError):
+                    return 0.0
+                
+            faithfulness_score = await asyncio.to_thread(extract_score)
             print(f"AXIOM-AUDIT: Faithfulness Score Verified at {faithfulness_score * 100}%")
 
             return {
@@ -84,7 +94,7 @@ class AxiomEvaluator:
             }
             
         except Exception as e:
-            print(f"RAGAS V2 EVAL ERROR: {e}")
+            print(f"⚠️ RAGAS V2 EVAL ERROR: {e}")
             return {"faithfulness": 0.0, "relevance": 1.0, "precision": 1.0}
 
 axiom_evaluator = AxiomEvaluator()

requirements.txt

@@ -50,5 +50,8 @@ tenacity==9.1.4
 jsonpatch==1.33
 setuptools==75.8.0
 
-# --- FORCE CPU-ONLY TORCH ---
-torch==2.2.2 --index-url https://download.pytorch.org/whl/cpu
+# --- FORCE CPU-ONLY TORCH (AXIOM-CORE compatible) ---
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch==2.4.0+cpu
+torchvision==0.19.0+cpu
+torchaudio==2.4.0+cpu