Update blaxel_main.py

blaxel_main.py

@@ -1,39 +1,45 @@
-# blaxel_main.py (Updated for Milestone 1)
+# blaxel_main.py (New Secure Version)
 from fastapi import FastAPI
 from pydantic import BaseModel
 from sse_starlette.sse import EventSourceResponse
 import os
 import uvicorn
+from typing import Dict, Optional
 
-# --- 1. Import our new "Brain" ---
+# 1. Import our "Brain"
 from agent_logic import StrategicSelectorAgent
 
 app = FastAPI()
 
+# 2. Update Pydantic model to accept keys (some are optional)
+class ApiKeys(BaseModel):
+    google: str
+    anthropic: Optional[str] = None
+    sambanova: Optional[str] = None
+
 class ProblemRequest(BaseModel):
     problem: str
-
-# --- 2. Initialize MudabbirAI (The Brain) ---
-# We create one instance when the server starts.
-mudabbir_ai = StrategicSelectorAgent()
-
+    keys: ApiKeys
 
 @app.post("/solve_problem")
 async def solve_problem(request: ProblemRequest):
     """
-    This endpoint now calls our real agent's .solve() method
-    and streams its thought process back to Gradio.
+    This endpoint now creates a *new* agent for every request,
+    and passes the user's API keys to it.
+    This is stateless and secure.
     """
-
-    # --- 3. Call the agent's .solve() generator ---
-    # We are returning the agent's generator directly.
-    # FastAPI and EventSourceResponse will handle streaming it.
+    
     async def stream_solution():
         try:
+            # 3. Pass keys to the agent's constructor
+            mudabbir_ai = StrategicSelectorAgent(api_keys=request.keys.dict())
+            
             async for status_update in mudabbir_ai.solve(request.problem):
                 yield status_update
         except Exception as e:
-            yield f"AGENT ERROR: {str(e)}"
+            # This will catch errors (like invalid keys)
+            yield f"AGENT ERROR: {e}"
+            print(f"AGENT ERROR: {e}") # Also log to backend
 
     return EventSourceResponse(stream_solution())