Update proxy_server.py

proxy_server.py

@@ -17,10 +17,11 @@ from typing import AsyncGenerator, Set, Optional, Dict, Any, List
 # --- Logging Configuration ---
 logger.remove()
 log_level = os.getenv("LOG_LEVEL", "INFO").upper()
-logger.add(sys.stderr, level=log_level)
+# Ensure DEBUG level enables detailed logging
+logger.add(sys.stderr, level=log_level, format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>")
 
 # --- Environment Variable Configuration ---
-# Target OpenAI API endpoint (can be OpenAI's official API or another compatible proxy)
+# Target OpenAI API endpoint (can be OpenAI's official API or another compatible proxy like chat-api)
 OPENAI_API_ENDPOINT = os.getenv("OPENAI_API_ENDPOINT", "https://api.openai.com/v1/chat/completions")
 # API Key for authenticating with the target OpenAI endpoint (if required by the endpoint)
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
@@ -50,9 +51,8 @@ async def lifespan(app: FastAPI):
     timeout_config = httpx.Timeout(connect=CONNECT_TIMEOUT, read=READ_TIMEOUT, write=WRITE_TIMEOUT, pool=POOL_TIMEOUT)
     proxy_config = {"http://": HTTP_PROXY, "https://": HTTP_PROXY} if HTTP_PROXY else None
 
-    # --- FIX: Completely hide target endpoint info from logs ---
+    # Completely hide target endpoint info from logs
     logger.info("Initializing httpx client for upstream requests.") # Generic message
-    # --- End Fix ---
 
     if proxy_config:
         logger.info(f"Using outbound proxy: {HTTP_PROXY}") # Proxy URL might still be sensitive
@@ -97,27 +97,24 @@ async def get_api_key(key: Optional[str] = Security(api_key_header)) -> str:
     """Validate the API key provided in the request header."""
     if not VALID_API_KEYS:
         logger.warning("No PROXY_API_KEYS configured. Allowing request.")
-        # If no keys are configured, allow the request (consider raising an error in production)
-        # raise HTTPException(status_code=500, detail="Server configuration error: PROXY_API_KEYS not set.")
-        return "unsecured_dummy_key" # Return a dummy value
+        return "unsecured_dummy_key"
 
     if key is None:
         logger.warning("API key missing from request header.")
         raise HTTPException(status_code=401, detail=f"API Key required in header '{API_KEY_NAME}'")
     if key not in VALID_API_KEYS:
-        # --- FIX: Avoid logging the invalid key directly ---
         logger.warning(f"Invalid API key received (length: {len(key)}).")
-        # --- End Fix ---
         raise HTTPException(status_code=401, detail="Invalid or expired API Key")
-    # --- FIX: Avoid logging the valid key directly ---
     logger.debug(f"Valid API key received (length: {len(key)}).")
-    # --- End Fix ---
     return key
 
 # --- Format Conversion Logic ---
 
-def claude_request_to_openai_payload(claude_request: Dict[str, Any]) -> Dict[str, Any]:
-    """Converts a Claude API request body to OpenAI API format."""
+def claude_request_to_openai_payload(claude_request: Dict[str, Any], client_api_key: str) -> Dict[str, Any]:
+    """
+    Converts a Claude API request body to OpenAI API format.
+    Optionally adds a 'user' field based on the client's API key.
+    """
     messages = []
     system_prompt = claude_request.get("system")
     if system_prompt:
@@ -153,6 +150,10 @@ def claude_request_to_openai_payload(claude_request: Dict[str, Any]) -> Dict[str
         "model": claude_request.get("model", "gpt-3.5-turbo"), # Use model from request or default
         "messages": messages,
         "stream": claude_request.get("stream", False),
+        # --- Potential Fix: Add user field for chat-api ---
+        # Use the client's API key as the user identifier, or derive one if needed
+        "user": f"proxy_user_{client_api_key}", # Example: prefixing the key
+        # --- End Potential Fix ---
         # Optional parameters - only include if present in Claude request
         **({ "max_tokens": v } if (v := claude_request.get("max_tokens")) is not None else {}),
         **({ "temperature": v } if (v := claude_request.get("temperature")) is not None else {}),
@@ -164,6 +165,7 @@ def claude_request_to_openai_payload(claude_request: Dict[str, Any]) -> Dict[str
     # logger.debug("Converted Claude request to OpenAI payload.") # Keep this simple
     return openai_payload
 
+# ... (openai_response_to_claude_response and stream_openai_response_to_claude_events remain the same) ...
 def openai_response_to_claude_response(openai_response: Dict[str, Any], claude_request_id: str) -> Dict[str, Any]:
     """Converts a non-streaming OpenAI response to Claude API format."""
     try:
@@ -204,8 +206,6 @@ def openai_response_to_claude_response(openai_response: Dict[str, Any], claude_r
         return claude_response
     except (KeyError, IndexError, TypeError) as e:
         logger.error(f"[{claude_request_id}] Error converting non-streaming OpenAI response: {e}")
-        # Avoid logging the full original response here as it might be large/sensitive
-        # logger.error(f"Original response snippet: {str(openai_response)[:200]}...") # Optional: log a snippet
         raise ValueError(f"Failed to parse OpenAI response: {e}")
 
 async def stream_openai_response_to_claude_events(openai_response: httpx.Response, claude_request_id: str, requested_model: str) -> AsyncGenerator[str, None]:
@@ -304,7 +304,6 @@ async def stream_openai_response_to_claude_events(openai_response: httpx.Respons
         yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': 0})}\n\n"
 
         # 6. Send message_delta with final usage and stop reason
-        # Note: Usage might be estimated if not provided by OpenAI at the end
         final_delta = {
             'type': 'message_delta',
             'delta': {
@@ -312,8 +311,7 @@ async def stream_openai_response_to_claude_events(openai_response: httpx.Respons
                 'stop_sequence': None # OpenAI doesn't provide this
             },
             'usage': {
-                # Use accumulated length as a rough proxy if output_tokens not updated
-                'output_tokens': output_tokens if output_tokens > 0 else (accumulated_content_len // 4) # Very rough estimate
+                'output_tokens': output_tokens if output_tokens > 0 else (accumulated_content_len // 4) # Rough estimate
             }
         }
         yield f"event: message_delta\ndata: {json.dumps(final_delta)}\n\n"
@@ -337,8 +335,8 @@ def create_error_response(status_code: int, error_type: str, message: str) -> JS
     )
 
 # --- Main Proxy Endpoint ---
-@app.post("/v1/messages", dependencies=[Depends(get_api_key)])
-async def proxy_claude_to_openai(request: Request):
+@app.post("/v1/messages") # Removed dependency here, will get API key manually
+async def proxy_claude_to_openai(request: Request, client_api_key: str = Depends(get_api_key)): # Get API key via Depends
     """
     Receives a Claude-formatted request, proxies it to OpenAI,
     and returns a Claude-formatted response.
@@ -347,15 +345,15 @@ async def proxy_claude_to_openai(request: Request):
     request_id = f"msg_{uuid.uuid4().hex[:24]}" # Generate a unique ID for logging/tracking
     try:
         claude_request_data = await request.json()
-        # Avoid logging potentially large/sensitive request data
         logger.info(f"[{request_id}] Received request. Stream: {claude_request_data.get('stream', False)}. Model: {claude_request_data.get('model')}")
     except json.JSONDecodeError:
         logger.error(f"[{request_id}] Invalid JSON received in request body.")
         return create_error_response(400, "invalid_request_error", "Invalid JSON data in request body.")
 
-    # Convert request format
+    # Convert request format, passing the client's API key
     try:
-        openai_payload = claude_request_to_openai_payload(claude_request_data)
+        # Pass the validated client_api_key to the conversion function
+        openai_payload = claude_request_to_openai_payload(claude_request_data, client_api_key)
     except Exception as e:
         logger.error(f"[{request_id}] Failed to convert Claude request to OpenAI format: {e}")
         return create_error_response(400, "invalid_request_error", f"Failed to process request data: {e}")
@@ -364,33 +362,54 @@ async def proxy_claude_to_openai(request: Request):
     requested_model = openai_payload.get("model", "unknown_model") # For response generation
 
     # Prepare headers for the target OpenAI endpoint
-    headers = {
-        "Content-Type": "application/json",
-        # Add other headers if needed
-    }
+    # --- Start Detailed Logging (DEBUG Level) ---
+    target_headers = { "Content-Type": "application/json" }
     if OPENAI_API_KEY:
-        # Avoid logging API key
-        # logger.debug(f"[{request_id}] Adding Authorization header to upstream request.") # Log presence, not value
-        headers["Authorization"] = f"Bearer {OPENAI_API_KEY}"
+        target_headers["Authorization"] = f"Bearer {OPENAI_API_KEY}"
+        # Avoid logging the actual key value even in debug
+        logger.debug(f"[{request_id}] Added Authorization header to upstream request.")
+    else:
+        logger.debug(f"[{request_id}] No OPENAI_API_KEY configured for upstream request.")
+
+    # Log headers and payload if log level is DEBUG
+    if logger.level("DEBUG").no >= logger.level(log_level).no: # Check if DEBUG is enabled
+        # Redact Authorization header before logging
+        logged_headers = target_headers.copy()
+        if "Authorization" in logged_headers:
+             logged_headers["Authorization"] = "Bearer [REDACTED]"
+        logger.debug(f"[{request_id}] Sending request to upstream API.")
+        logger.debug(f"[{request_id}] Upstream Headers: {json.dumps(logged_headers)}")
+        # Be cautious logging the full payload due to size/sensitivity
+        try:
+            # Log payload safely using json.dumps for proper formatting
+            payload_str = json.dumps(openai_payload, indent=2)
+            # Limit payload log size if needed
+            max_log_len = 1024 # Log max 1KB of payload
+            if len(payload_str) > max_log_len:
+                 logger.debug(f"[{request_id}] Upstream Payload (truncated): {payload_str[:max_log_len]}...")
+            else:
+                 logger.debug(f"[{request_id}] Upstream Payload: {payload_str}")
+        except Exception as log_e:
+            logger.warning(f"[{request_id}] Could not serialize or log upstream payload: {log_e}")
+    else:
+         logger.debug(f"[{request_id}] Sending request to upstream API...") # Generic message for INFO level
+    # --- End Detailed Logging ---
 
-    try:
-        # Hide full target endpoint URL and payload from logs
-        logger.debug(f"[{request_id}] Sending request to upstream API...")
 
+    try:
         # Build the request to the target endpoint
         target_request = client.build_request(
             method="POST",
             url=OPENAI_API_ENDPOINT,
-            headers=headers,
+            headers=target_headers, # Use the prepared headers
             json=openai_payload,
-            # Timeout is handled globally by the client
         )
 
         # Send the request and handle the response
         response = await client.send(target_request, stream=is_streaming)
 
         # Check for HTTP errors from the target endpoint *before* processing the body
-        response.raise_for_status()
+        response.raise_for_status() # Raises exception for 4xx/5xx errors
 
         # Process the response based on streaming or non-streaming
         if is_streaming:
@@ -406,9 +425,7 @@ async def proxy_claude_to_openai(request: Request):
             )
         else:
             logger.info(f"[{request_id}] Upstream response is non-streaming. Converting.")
-            # FIX: Remove await here
-            openai_response_data = response.json()
-            # Avoid logging full response data
+            openai_response_data = response.json() # No await needed
             logger.debug(f"[{request_id}] Received non-streaming response from upstream.")
             try:
                 claude_response_data = openai_response_to_claude_response(openai_response_data, request_id)
@@ -424,66 +441,46 @@ async def proxy_claude_to_openai(request: Request):
     # --- Error Handling for Target API Request ---
     except httpx.HTTPStatusError as e:
         status_code = e.response.status_code
-        error_detail_text = "[Could not decode error response]" # Default message
+        error_detail_text = "[Could not decode error response]"
         try:
-            # Try reading error details as JSON first
             error_detail = e.response.json()
-            error_detail_text = json.dumps(error_detail) # Convert back to string for logging snippet
+            error_detail_text = json.dumps(error_detail)
         except json.JSONDecodeError:
-            # If not JSON, read as text
             try:
-                # Read text synchronously since the body is likely already loaded for non-streaming error responses
                 error_detail_text = e.response.text
             except Exception:
                  logger.warning(f"[{request_id}] Could not read error response body as text.")
 
-
-        # Log error snippet, avoid full potentially sensitive detail
         logger.error(f"[{request_id}] HTTP error from target endpoint ({status_code}). Response snippet: {error_detail_text[:200]}...")
 
-        # Map common HTTP errors to Claude error types
-        # Use generic messages in production to avoid leaking upstream details
-        if status_code == 400:
-            err_type, msg = "invalid_request_error", f"Upstream API reported Bad Request ({status_code})."
-        elif status_code == 401:
-            err_type, msg = "authentication_error", f"Authentication failed with upstream API ({status_code})."
-        elif status_code == 403:
-            err_type, msg = "permission_error", f"Forbidden by upstream API ({status_code})."
-        elif status_code == 429:
-            err_type, msg = "rate_limit_error", f"Rate limit exceeded with upstream API ({status_code})."
-        elif status_code >= 500:
-            err_type, msg = "api_error", f"Upstream API unavailable or encountered an error ({status_code})."
-        else:
-            err_type, msg = "api_error", f"Received unexpected error from upstream API ({status_code})."
+        if status_code == 400: err_type, msg = "invalid_request_error", f"Upstream API Bad Request ({status_code})."
+        elif status_code == 401: err_type, msg = "authentication_error", f"Authentication failed with upstream API ({status_code})."
+        elif status_code == 403: err_type, msg = "permission_error", f"Forbidden by upstream API ({status_code})."
+        elif status_code == 429: err_type, msg = "rate_limit_error", f"Rate limit exceeded with upstream API ({status_code})."
+        elif status_code >= 500: err_type, msg = "api_error", f"Upstream API unavailable or encountered an error ({status_code})."
+        else: err_type, msg = "api_error", f"Received unexpected error from upstream API ({status_code})."
 
-        # Return error response immediately
         return create_error_response(status_code, err_type, msg)
 
     except httpx.TimeoutException:
-         # Hide target endpoint URL from timeout log
         logger.error(f"[{request_id}] Request to target endpoint timed out ({READ_TIMEOUT}s).")
         return create_error_response(504, "api_error", "Gateway Timeout: Request to upstream API timed out.")
     except httpx.RequestError as e:
-        # Hide target endpoint URL from request error log
-        # The exception 'e' might contain the URL, so log a generic message
         logger.error(f"[{request_id}] Network error connecting to target endpoint: {type(e).__name__}")
         return create_error_response(502, "api_error", f"Bad Gateway: Network error connecting to upstream API.")
     except Exception as e:
-        logger.exception(f"[{request_id}] Unexpected error during proxy operation: {e}") # Use logger.exception to include traceback
+        logger.exception(f"[{request_id}] Unexpected error during proxy operation: {e}")
         return create_error_response(500, "internal_server_error", f"Internal Server Error: {e}")
 
 # --- Health Check Endpoint ---
 @app.get("/health", summary="Health Check", tags=["Management"])
 async def health_check():
     """Returns the operational status of the proxy server."""
-    # Could add checks here (e.g., try connecting to target endpoint)
     return {"status": "healthy"}
 
 # --- Run with Uvicorn (for local development) ---
-# In production (Docker), Uvicorn will be started via CMD
 if __name__ == "__main__":
     import uvicorn
-    # Load .env file for local development if python-dotenv is installed
     try:
         from dotenv import load_dotenv
         load_dotenv()
@@ -497,22 +494,21 @@ if __name__ == "__main__":
         log_level = os.getenv("LOG_LEVEL", "INFO").upper()
         # Reconfigure logger if level changed
         logger.remove()
-        logger.add(sys.stderr, level=log_level)
+        logger.add(sys.stderr, level=log_level, format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>")
         logger.info(f"Log level set to: {log_level}")
         logger.info(f"Valid Proxy API Keys configured: {len(VALID_API_KEYS)}")
     except ImportError:
         logger.info("python-dotenv not installed, skipping .env file loading.")
 
-    # Use port 7860 for consistency with Dockerfile default
     port = int(os.getenv("PORT", 7860))
     host = os.getenv("HOST", "0.0.0.0")
     log_config_level = log_level.lower() if log_level in ["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG", "TRACE"] else "info"
 
     logger.info(f"Starting Uvicorn server on {host}:{port}")
     uvicorn.run(
-        "proxy_server:app", # Point to the FastAPI app object in this file
+        "proxy_server:app",
         host=host,
         port=port,
-        reload=True, # Enable auto-reload for local development
-        log_level=log_config_level # Set Uvicorn's log level
+        reload=True,
+        log_level=log_config_level
     )