Update core_mcp.py

core_mcp.py

@@ -185,10 +185,6 @@ Generate the complete code based on these clarified requirements."""
 def format_strategist_audit_prompt(blueprint: str, intent: str, generated_code: str) -> str:
     return f"""{HEADER_PROMPT}
 
-**AUDIT PHASE INITIATED**
-
-The Coder AI has generated the following code based on the blueprint. Your task is to AUDIT this code against the original intent and blueprint.
-
 ORIGINAL BLUEPRINT:
 {blueprint}
 
@@ -198,22 +194,14 @@ ORIGINAL INTENT:
 GENERATED CODE:
 {generated_code}
 
-**INSTRUCTIONS:**
-1. Review the code for logical errors, missing features from the blueprint, or deviations from the intent.
-2. If the code is acceptable and meets all requirements, respond with the exact string: "AUDIT_PASSED".
-3. If changes are required, provide a numbered list of specific, actionable corrections. DO NOT rewrite the code yourself. Just list the required changes.
-"""
+**[AUDIT DIRECTIVE]**
+Review the generated code for compliance with the blueprint and intent. 
 
-def format_audit_resolution_prompt(blueprint: str, audit_feedback: str) -> str:
-    return f"""{SYSTEM_PROMPT}
-
-BLUEPRINT:
-{blueprint}
+Respond with EXACTLY ONE of the following:
+1. If code is acceptable: "AUDIT_PASSED"
+2. If changes required: Provide a numbered list of required changes.
 
-AUDIT FEEDBACK (REQUIRED CHANGES):
-{audit_feedback}
-
-The previous code generation failed the audit. Regenerate the code, strictly adhering to the audit feedback above."""
+Do not generate code. Only provide audit results."""
 
 # --- UTILITIES ---
 
@@ -242,6 +230,8 @@ def sanitize_error_message(error_msg: str) -> str:
     Scrub API keys and sensitive URLs from error messages.
     Removes everything after '?' in URLs to prevent API key exposure.
     """
+    import re
+    
     # Pattern to match URLs with query parameters (containing API keys)
     # This will find "for url: https://...?key=..." and remove everything after the ?
     pattern = r'(for url: https?://[^\s?]+)\?[^\s]*'
@@ -433,9 +423,7 @@ def execute_http_tool(tool_name: str, tool_config: Dict, params: Dict[str, str],
 
 def call_ai_api(full_url: str, api_key: str, messages: List[Dict[str, str]]) -> str:
     headers = {'Content-Type': 'application/json'}
-    # We construct the URL with the key here for the request...
     url_with_key = f"{full_url}?key={api_key}"
-    
     payload = {
         'contents': messages,
         'generationConfig': {
@@ -445,17 +433,14 @@ def call_ai_api(full_url: str, api_key: str, messages: List[Dict[str, str]]) ->
     }
     try:
         response = requests.post(url_with_key, headers=headers, json=payload, timeout=120)
-        response.raise_for_status() # This is what usually triggers the error containing the URL
+        response.raise_for_status()
         data = response.json()
         if 'candidates' in data and len(data['candidates']) > 0:
             return data['candidates'][0]['content']['parts'][0]['text']
         else:
             return f"ERROR: API malformed response: {data}"
     except Exception as e:
-        # SECURITY FIX: Sanitize the exception string before returning it
-        # This strips the '?key=xyz' from the error message
-        sanitized_error = sanitize_error_message(str(e))
-        return f"ERROR: API call failed: {sanitized_error}"
+        return f"ERROR: API call failed: {str(e)}"
 
 def resolve_model(selected_label: str, model_registry: Dict[str, Dict], fallback_label: str, agent_name: str) -> Tuple[str, str, str]:
     if selected_label and selected_label in model_registry:
@@ -482,6 +467,19 @@ def resolve_model(selected_label: str, model_registry: Dict[str, Dict], fallback
         )
     return ('', '', f"FATAL: No {agent_name} models configured.")
 
+def format_audit_entry(diagnostic: Dict[str, any]) -> str:
+    entry = {
+        "transaction_id": diagnostic['transaction_id'],
+        "timestamp": diagnostic['timestamp'],
+        "tool": diagnostic['tool_name'],
+        "status": diagnostic['status'],
+        "latency": diagnostic['latency'],
+        "request": diagnostic['request_payload'],
+        "response": diagnostic['response_payload'],
+        "error": diagnostic['error_message']
+    }
+    return json.dumps(entry, indent=2)
+
 # --- ENGINE ---
 
 class ConductorEngine:
@@ -491,10 +489,8 @@ class ConductorEngine:
         self.system_log: List[str] = []
         self.audit_log: List[Dict[str, Any]] = []
         self.loop_count: int = 0
-        self.max_loops: int = 6 if enable_audit else 3 # Increased loop limit for audit
-        self.enable_audit = enable_audit
-        self.final_output_content = ""
-        
+        self.max_loops: int = 6
+        self.enable_audit: bool = enable_audit
         self.coder_model_id, self.coder_url, coder_log = resolve_model(
             coder_model_label, CONFIG['coder_models'], CONFIG['coder_model_default_label'], 'Coder'
         )
@@ -522,18 +518,17 @@ class ConductorEngine:
             "coder_log": self.coder_history,
             "strategist_log": self.strategist_history,
             "audit_log": self.audit_log,
-            "final_output": self.final_output_content
+            "final_output": ""
         }
 
     def execute_workflow(self, blueprint: str, intent: str) -> Dict[str, Any]:
         """ORIGINAL: Non-streaming version (kept for backwards compatibility)"""
-        # Note: Audit logic not implemented in legacy non-streaming version
         self.coder_history = []
         self.strategist_history = []
         self.system_log = []
         self.audit_log = []
         self.loop_count = 0
-        self.final_output_content = ""
+        final_code = ""
 
         self.log_system("INITIATING CONDUCTOR ENGINE...")
         for log_msg in self.model_resolution_logs:
@@ -557,8 +552,8 @@ class ConductorEngine:
             
             if "DETECTED" not in coder_response:
                 self.log_system("MISSION COMPLETE: Ambiguity resolved or Code generated.")
-                self.final_output_content = coder_response
-                return self._compile_results(coder_response)
+                final_code = coder_response
+                return self._compile_results(final_code)
             
             self.log_system("Ambiguity DETECTED. Escalating to Strategist.")
             strategist_prompt = format_strategist_prompt(blueprint, intent, coder_response)
@@ -588,13 +583,12 @@ class ConductorEngine:
         return self._compile_results(coder_response)
 
     def execute_workflow_streaming(self, blueprint: str, intent: str):
-        """NEW: Streaming version that yields updates as they happen"""
+        """STREAMING version with TASK 3 & 4 modifications"""
         self.coder_history = []
         self.strategist_history = []
         self.system_log = []
         self.audit_log = []
         self.loop_count = 0
-        self.final_output_content = ""
 
         self.log_system("INITIATING CONDUCTOR ENGINE...")
         yield self._current_state()
@@ -610,7 +604,7 @@ class ConductorEngine:
 
         self.log_system(f"Engaged: Coder ({self.coder_model_id}), Strategist ({self.strategist_model_id})")
         if self.enable_audit:
-            self.log_system("AUDIT PROTOCOL: ENABLED")
+            self.log_system("AUDIT MODE: ENABLED")
         yield self._current_state()
         
         initial_prompt = format_initial_prompt(blueprint)
@@ -626,44 +620,44 @@ class ConductorEngine:
             self.log_coder("assistant", coder_response)
             yield self._current_state()
             
-            # CHECK FOR COMPLETION OR AMBIGUITY
             if "DETECTED" not in coder_response:
-                # Code generated. Check if Audit is required.
-                if self.enable_audit:
-                    self.log_system("Code Generated. INITIATING AUDIT PHASE...")
-                    yield self._current_state()
-                    
-                    audit_prompt = format_strategist_audit_prompt(blueprint, intent, coder_response)
-                    self.log_strategist("user", audit_prompt)
-                    yield self._current_state()
-                    
-                    strategist_response = call_ai_api(self.strategist_url, CONFIG['STRATEGIST_API_KEY'], self.strategist_history)
-                    self.log_strategist("assistant", strategist_response)
-                    yield self._current_state()
-                    
-                    if "AUDIT_PASSED" in strategist_response:
-                        self.log_system("AUDIT STATUS: PASSED. Mission Complete.")
-                        self.final_output_content = coder_response
-                        yield self._current_state()
-                        return
-                    else:
-                        self.log_system("AUDIT STATUS: FAILED. Revisions Requested.")
-                        yield self._current_state()
-                        
-                        # Feed audit feedback back to Coder
-                        audit_resolution_prompt = format_audit_resolution_prompt(blueprint, strategist_response)
-                        self.log_coder("user", audit_resolution_prompt)
-                        yield self._current_state()
-                        # Loop continues to next iteration
-                        continue
-                else:
-                    # No audit, just finish
+                # TASK 4: AUDIT LOOP LOGIC
+                if not self.enable_audit:
+                    # Original termination path
                     self.log_system("MISSION COMPLETE: Ambiguity resolved or Code generated.")
-                    self.final_output_content = coder_response
-                    yield self._current_state()
+                    final_state = self._current_state()
+                    final_state["final_code"] = coder_response  # TASK 3: Add final_code key
+                    yield final_state
+                    return
+                # AUDIT MODE ACTIVATED
+                self.log_system("AUDIT PHASE INITIATED.")
+                yield self._current_state()
+                
+                audit_prompt = format_strategist_audit_prompt(blueprint, intent, coder_response)
+                self.log_strategist("user", audit_prompt)
+                yield self._current_state()
+                
+                audit_response = call_ai_api(self.strategist_url, CONFIG['STRATEGIST_API_KEY'], self.strategist_history)
+                self.log_strategist("assistant", audit_response)
+                yield self._current_state()
+                
+                if "AUDIT_PASSED" in audit_response:
+                    self.log_system("AUDIT PASSED. MISSION COMPLETE.")
+                    final_state = self._current_state()
+                    final_state["final_code"] = coder_response  # TASK 3: Add final_code key
+                    yield final_state
                     return
+                
+                # Audit failed - loop back to coder
+                self.log_system("AUDIT FAILED. Generating revisions.")
+                yield self._current_state()
+                
+                revision_prompt = format_resolution_prompt(blueprint, f"AUDIT CHANGES REQUIRED:\n{audit_response}")
+                self.log_coder("user", revision_prompt)
+                yield self._current_state()
+                # Continue loop for coder to process revisions
+                continue
             
-            # AMBIGUITY DETECTED PATH
             self.log_system("Ambiguity DETECTED. Escalating to Strategist.")
             yield self._current_state()
             
@@ -699,7 +693,6 @@ class ConductorEngine:
             yield self._current_state()
 
         self.log_system("FATAL: Loop limit reached.")
-        self.final_output_content = coder_response # Return whatever we have
         yield self._current_state()
 
     def _compile_results(self, final_output: str) -> Dict[str, Any]:
@@ -713,14 +706,14 @@ class ConductorEngine:
 
 # --- API ---
 
-app = FastAPI(title="Conductor MCP Server", version="4.1")
+app = FastAPI(title="Conductor MCP Server", version="4.0")
 
 class ExecutionRequest(BaseModel):
     blueprint: str
     intent: str
     coder_model_label: Optional[str] = None
     strategist_model_label: Optional[str] = None
-    enable_audit: Optional[bool] = False
+    enable_audit: Optional[bool] = False  # TASK 4: New field
 
 class ExecutionResponse(BaseModel):
     system_log: List[str]
@@ -738,7 +731,7 @@ async def execute_conductor(req: ExecutionRequest):
 
 @app.post("/execute_stream")
 async def execute_conductor_stream(req: ExecutionRequest):
-    """NEW: Streaming endpoint that yields real-time updates"""
+    """STREAMING endpoint with TASK 4 audit support"""
     
     async def generate_updates():
         engine = ConductorEngine(req.coder_model_label, req.strategist_model_label, req.enable_audit)
@@ -752,4 +745,4 @@ async def execute_conductor_stream(req: ExecutionRequest):
     return StreamingResponse(generate_updates(), media_type="text/event-stream")
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)