Update blaxel_main.py

blaxel_main.py

@@ -1,4 +1,4 @@
-# blaxel_main.py (New Secure Version)
+# blaxel_main.py (Updated: New Provider Keys + Core Logic Only)
 from fastapi import FastAPI
 from pydantic import BaseModel
 from sse_starlette.sse import EventSourceResponse
@@ -6,44 +6,44 @@ import os
 import uvicorn
 from typing import Dict, Optional
 
-# 1. Import our "Brain"
 from agent_logic import StrategicSelectorAgent
 
 app = FastAPI()
 
-# 2. Update Pydantic model to accept keys (some are optional)
+# 1. Update Pydantic model to accept all provider keys
 class ApiKeys(BaseModel):
     google: str
     anthropic: Optional[str] = None
     sambanova: Optional[str] = None
+    # New providers added:
+    openai: Optional[str] = None
+    nebius: Optional[str] = None
 
 class ProblemRequest(BaseModel):
     problem: str
+    # Expansion flag removed - we are sticking to the robust core loop
     keys: ApiKeys
 
 @app.post("/solve_problem")
 async def solve_problem(request: ProblemRequest):
     """
-    This endpoint now creates a *new* agent for every request,
-    and passes the user's API keys to it.
-    This is stateless and secure.
+    Standard endpoint. Initializes the agent with all provided keys
+    and streams the solution process.
     """
-    
     async def stream_solution():
         try:
-            # 3. Pass keys to the agent's constructor
-            mudabbir_ai = StrategicSelectorAgent(api_keys=request.keys.model_dump())
+            # Pass the dictionary of keys (including new ones) to the agent
+            # request.keys.dict() will include openai and nebius if sent
+            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:
-            # This will catch errors (like invalid keys)
             yield f"AGENT ERROR: {e}"
-            print(f"AGENT ERROR: {e}") # Also log to backend
+            print(f"AGENT ERROR: {e}") 
 
     return EventSourceResponse(stream_solution())
 
-
 if __name__ == "__main__":
     HOST = os.getenv("BL_SERVER_HOST", "0.0.0.0")
     PORT = int(os.getenv("BL_SERVER_PORT", "8080"))