Update app.py

app.py

@@ -109,7 +109,12 @@ async def get_proxy():
             <!-- Connection Panel -->
             <div style="border-right: 1px solid #ccc; padding-right: 20px;">
                 <div style="margin-bottom: 20px;">
-                    <input type="password" id="apiKey" placeholder="Enter API Key" style="width: 100%;">
+                    <div style="margin-bottom: 20px;">
+                        <input type="password" id="apiKey" placeholder="LLM API Key" style="width: 100%; margin-bottom: 8px;">
+                        <input type="password" id="incomingKey" placeholder="Proxy Incoming Key" style="width: 100%;">
+                        <button onclick="initializeClient()" style="margin-top: 10px;">Connect</button>
+                    </div>
+
                     <button onclick="initializeClient()" style="margin-top: 10px;">Fetch Models</button>
                 </div>
                 <select id="modelSelect" style="width: 100%; margin-bottom: 20px;"></select>
@@ -301,6 +306,16 @@ async def get_proxy():
                 if (msg.destination === 'proxy') {
                     addMessageEntry('incoming', msg.source, 'proxy', msg.content);
                     document.getElementById('detailedStatus').textContent = `Processing ${msg.source} request...`;
+                    // check if incoming call has the correct key
+                    const expected_incoming_key = document.getElementById('incomingKey');
+                    if (msg.incomingKey != expected_incoming_key){
+                        return {
+                            request_id: msg.request_id,  // Critical addition
+                            content: "Error: Incoming Authentication Key incorrect!",
+                            source: 'proxy',
+                            destination: msg.source
+                        };
+                    }
                     
                     try {
                         const llmResponse = await agentClient.call(currentModel, msg.content, systemPrompt, conversationHistory);
@@ -344,32 +359,45 @@ async def get_proxy():
 
 
 @app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
-async def chat_completions(request: ChatCompletionRequest):
+async def chat_completions(
+    request: ChatCompletionRequest,
+    authorization: str = Header(None)  # Add Authorization header
+):
+    # Extract and validate API key
+    if not authorization or not authorization.startswith("Bearer "):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Missing or invalid Authorization header"
+        )
+    api_key = authorization[7:]  # Remove "Bearer " prefix
+
     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
+
+    # Add API key to proxy message
     proxy_msg = {
         "request_id": request_id,
         "content": user_message.content,
         "source": "api",
         "destination": "proxy",
         "model": request.model,
-        "temperature": request.temperature
+        "temperature": request.temperature,
+        "incoming_key": api_key  # Critical addition
     }
+
     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()),
@@ -380,6 +408,7 @@ async def chat_completions(request: ChatCompletionRequest):
     )
 
 
+
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
     await manager.connect(websocket)