Update app.py

app.py

@@ -2,12 +2,41 @@ from fastapi import FastAPI, WebSocket
 from fastapi.responses import HTMLResponse
 import uvicorn
 import json
+from fastapi import Request, HTTPException
+from pydantic import BaseModel
+from typing import List, Optional
+import uuid
+import time
+import asyncio
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatCompletionRequest(BaseModel):
+    model: str = "gemini-2.5-pro-exp-03-25"
+    messages: List[ChatMessage]
+    temperature: Optional[float] = 0.7
+    stream: Optional[bool] = False
+
+class ChatCompletionResponseChoice(BaseModel):
+    index: int = 0
+    message: ChatMessage
+    finish_reason: str = "stop"
+
+class ChatCompletionResponse(BaseModel):
+    id: str
+    object: str = "chat.completion"
+    created: int
+    model: str
+    choices: List[ChatCompletionResponseChoice]
 
 app = FastAPI()
 
 class ConnectionManager:
     def __init__(self):
         self.active_connections = {}  # WebSocket: source
+        self.response_queues = {}     # request_id: asyncio.Queue
 
     async def connect(self, websocket: WebSocket):
         await websocket.accept()
@@ -26,6 +55,15 @@ class ConnectionManager:
         if websocket in self.active_connections:
             del self.active_connections[websocket]
 
+    async def wait_for_response(self, request_id: str, timeout: int = 30):
+        queue = asyncio.Queue(maxsize=1)
+        self.response_queues[request_id] = queue
+        try:
+            return await asyncio.wait_for(queue.get(), timeout=timeout)
+        finally:
+            self.response_queues.pop(request_id, None)
+
+
 manager = ConnectionManager()
 
 @app.get("/")
@@ -244,6 +282,43 @@ async def get_proxy():
     """)
 
 
+@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
+async def chat_completions(request: ChatCompletionRequest):
+    request_id = str(uuid.uuid4())
+    # Find proxy connection
+    proxy_ws = next((ws for ws, src in manager.active_connections.items() if src == 'proxy'), None)
+    if not proxy_ws:
+        raise HTTPException(503, "Proxy client not connected")
+    # Get user message
+    user_message = next((m for m in request.messages if m.role == "user"), None)
+    if not user_message:
+        raise HTTPException(400, "No user message found")
+    # Send to proxy
+    proxy_msg = {
+        "request_id": request_id,
+        "content": user_message.content,
+        "source": "api",
+        "destination": "proxy",
+        "model": request.model,
+        "temperature": request.temperature
+    }
+    await proxy_ws.send_text(json.dumps(proxy_msg))
+    # Wait for response from proxy
+    try:
+        response_content = await manager.wait_for_response(request_id)
+    except asyncio.TimeoutError:
+        raise HTTPException(504, "Proxy response timeout")
+    # Return OpenAI-compatible response
+    return ChatCompletionResponse(
+        id=request_id,
+        created=int(time.time()),
+        model=request.model,
+        choices=[ChatCompletionResponseChoice(
+            message=ChatMessage(role="assistant", content=response_content)
+        )]
+    )
+
+
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
     await manager.connect(websocket)
@@ -253,13 +328,17 @@ async def websocket_endpoint(websocket: WebSocket):
         init_msg = json.loads(data)
         if 'source' in init_msg:
             manager.set_source(websocket, init_msg['source'])
-        
         # Handle messages
         while True:
             message = await websocket.receive_text()
             msg_data = json.loads(message)
-            await manager.send_to_destination(msg_data['destination'], message)
-            
+            # If this is a response to an API request
+            if 'request_id' in msg_data and msg_data.get('destination') == 'api':
+                queue = manager.response_queues.get(msg_data['request_id'])
+                if queue:
+                    await queue.put(msg_data['content'])
+            else:
+                await manager.send_to_destination(msg_data['destination'], message)
     except Exception as e:
         manager.remove(websocket)
         await websocket.close()