Update proxy_server.py

proxy_server.py

@@ -12,6 +12,7 @@ from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from loguru import logger
 from typing import AsyncGenerator, Set, Optional, Dict, Any, List
+from urllib.parse import urlparse # Import để phân tích URL
 
 # --- Logging Configuration ---
 logger.remove()
@@ -49,9 +50,19 @@ 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
 
-    logger.info(f"Initializing httpx client. Target Endpoint: {OPENAI_API_ENDPOINT}")
+    # --- FIX: Hide full target endpoint URL from logs ---
+    # Parse the URL to get only the hostname for logging
+    try:
+        parsed_url = urlparse(OPENAI_API_ENDPOINT)
+        target_host = parsed_url.netloc # e.g., api.openai.com
+    except Exception:
+        target_host = "[Invalid Target URL]" # Handle potential parsing errors
+
+    logger.info(f"Initializing httpx client. Target Host: {target_host}") # Log only the host
+    # --- End Fix ---
+
     if proxy_config:
-        logger.info(f"Using outbound proxy: {HTTP_PROXY}")
+        logger.info(f"Using outbound proxy: {HTTP_PROXY}") # Proxy URL might still be sensitive depending on config
     if not OPENAI_API_KEY:
         logger.warning("OPENAI_API_KEY is not set. Requests to the target endpoint might fail if it requires authentication.")
     if not VALID_API_KEYS:
@@ -101,9 +112,13 @@ async def get_api_key(key: Optional[str] = Security(api_key_header)) -> str:
         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:
-        logger.warning(f"Invalid API key received: '{key[:4]}...'")
+        # --- 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")
-    logger.debug(f"Valid API key received: '{key[:4]}...'")
+    # --- FIX: Avoid logging the valid key directly ---
+    logger.debug(f"Valid API key received (length: {len(key)}).")
+    # --- End Fix ---
     return key
 
 # --- Format Conversion Logic ---
@@ -143,7 +158,9 @@ def claude_request_to_openai_payload(claude_request: Dict[str, Any]) -> Dict[str
         # Add other relevant parameter mappings here (e.g., presence_penalty, frequency_penalty)
     }
 
-    logger.debug("Converted Claude request to OpenAI payload.")
+    # --- FIX: Avoid logging potentially large/sensitive payload ---
+    # logger.debug("Converted Claude request to OpenAI payload.") # Keep this simple
+    # --- End Fix ---
     return openai_payload
 
 def openai_response_to_claude_response(openai_response: Dict[str, Any], claude_request_id: str) -> Dict[str, Any]:
@@ -182,23 +199,24 @@ def openai_response_to_claude_response(openai_response: Dict[str, Any], claude_r
                 "output_tokens": completion_tokens,
             },
         }
-        logger.debug("Converted non-streaming OpenAI response to Claude format.")
+        logger.debug(f"[{claude_request_id}] Converted non-streaming OpenAI response to Claude format.")
         return claude_response
     except (KeyError, IndexError, TypeError) as e:
-        logger.error(f"Error converting non-streaming OpenAI response: {e}\nOriginal response: {openai_response}")
-        # Re-raise or return a structured error for the client
+        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]:
     """Converts an OpenAI SSE stream to Claude API SSE format."""
     message_id = claude_request_id # Use the original request ID for consistency
-    accumulated_content = ""
+    accumulated_content_len = 0 # Track length instead of full content
     openai_finish_reason = None
     input_tokens = 0 # Will be updated if usage info is sent
     output_tokens = 0 # Will be updated if usage info is sent
     last_ping_time = time.time()
 
-    logger.debug(f"Starting Claude SSE stream conversion for request ID: {message_id}")
+    logger.debug(f"[{message_id}] Starting Claude SSE stream conversion.")
 
     # 1. Send message_start event
     yield f"event: message_start\ndata: {json.dumps({'type': 'message_start', 'message': {'id': message_id, 'type': 'message', 'role': 'assistant', 'content': [], 'model': requested_model, 'stop_reason': None, 'stop_sequence': None, 'usage': {'input_tokens': 0, 'output_tokens': 0}}})}\n\n"
@@ -216,7 +234,7 @@ async def stream_openai_response_to_claude_events(openai_response: httpx.Respons
             if line.startswith("data:"):
                 data_str = line[len("data: "):].strip()
                 if data_str == "[DONE]":
-                    logger.debug("Received [DONE] marker from OpenAI stream.")
+                    logger.debug(f"[{message_id}] Received [DONE] marker from OpenAI stream.")
                     break # End of OpenAI stream
 
                 try:
@@ -231,25 +249,25 @@ async def stream_openai_response_to_claude_events(openai_response: httpx.Respons
                     # Check for finish reason in the chunk
                     if choices[0].get("finish_reason"):
                         openai_finish_reason = choices[0].get("finish_reason")
-                        logger.debug(f"Received OpenAI finish_reason: {openai_finish_reason}")
+                        logger.debug(f"[{message_id}] Received OpenAI finish_reason: {openai_finish_reason}")
 
                     # Check for usage update (some models send it at the end)
                     usage_update = data.get("usage")
                     if usage_update:
                         input_tokens = usage_update.get("prompt_tokens", input_tokens)
                         output_tokens = usage_update.get("completion_tokens", output_tokens)
-                        logger.debug(f"Received usage update: input={input_tokens}, output={output_tokens}")
+                        logger.debug(f"[{message_id}] Received usage update: input={input_tokens}, output={output_tokens}")
 
                     if content_chunk:
-                        accumulated_content += content_chunk
+                        accumulated_content_len += len(content_chunk)
                         # 4. Send content_block_delta for the text chunk
                         yield f"event: content_block_delta\ndata: {json.dumps({'type': 'content_block_delta', 'index': 0, 'delta': {'type': 'text_delta', 'text': content_chunk}})}\n\n"
 
                 except json.JSONDecodeError:
-                    logger.warning(f"Could not decode JSON from stream line: {data_str}")
+                    logger.warning(f"[{message_id}] Could not decode JSON from stream line: {data_str}")
                     continue
                 except Exception as e:
-                    logger.error(f"Error processing stream data chunk: {e}")
+                    logger.error(f"[{message_id}] Error processing stream data chunk: {e}")
                     continue # Skip this chunk
 
             # Send periodic pings
@@ -259,11 +277,11 @@ async def stream_openai_response_to_claude_events(openai_response: httpx.Respons
                 last_ping_time = current_time
 
     except httpx.ReadTimeout:
-        logger.error("Timeout reading from OpenAI stream.")
+        logger.error(f"[{message_id}] Timeout reading from OpenAI stream.")
         openai_finish_reason = "error_timeout" # Custom reason
         yield f"event: error\ndata: {json.dumps({'type': 'error', 'error': {'type': 'overloaded_error', 'message': 'Proxy timed out waiting for OpenAI stream'}})}\n\n"
     except Exception as e:
-        logger.exception(f"Unexpected error during stream processing: {e}")
+        logger.exception(f"[{message_id}] Unexpected error during stream processing: {e}")
         openai_finish_reason = "error_exception" # Custom reason
         yield f"event: error\ndata: {json.dumps({'type': 'error', 'error': {'type': 'internal_server_error', 'message': f'Proxy stream processing error: {e}'}})}\n\n"
     finally:
@@ -279,7 +297,7 @@ async def stream_openai_response_to_claude_events(openai_response: httpx.Respons
         }
         claude_stop_reason = stop_reason_map.get(openai_finish_reason, "stop_sequence") # Default
 
-        logger.debug(f"Stream finished. OpenAI finish reason: {openai_finish_reason}, mapped Claude stop reason: {claude_stop_reason}")
+        logger.debug(f"[{message_id}] Stream finished. OpenAI finish reason: {openai_finish_reason}, mapped Claude stop reason: {claude_stop_reason}")
 
         # 5. Send content_block_stop
         yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': 0})}\n\n"
@@ -293,14 +311,15 @@ async def stream_openai_response_to_claude_events(openai_response: httpx.Respons
                 'stop_sequence': None # OpenAI doesn't provide this
             },
             'usage': {
-                'output_tokens': output_tokens if output_tokens > 0 else len(accumulated_content.split()) # Rough estimate if needed
+                # 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
             }
         }
         yield f"event: message_delta\ndata: {json.dumps(final_delta)}\n\n"
 
         # 7. Send message_stop
         yield f"event: message_stop\ndata: {json.dumps({'type': 'message_stop'})}\n\n"
-        logger.info(f"Completed sending Claude SSE stream for request ID: {message_id}")
+        logger.info(f"[{message_id}] Completed sending Claude SSE stream.")
 
 
 def create_error_response(status_code: int, error_type: str, message: str) -> JSONResponse:
@@ -327,7 +346,9 @@ 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()
+        # --- FIX: 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')}")
+        # --- End Fix ---
     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.")
@@ -348,10 +369,16 @@ async def proxy_claude_to_openai(request: Request):
         # Add other headers if needed
     }
     if OPENAI_API_KEY:
+        # --- FIX: 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}"
+        # --- End Fix ---
 
     try:
-        logger.debug(f"[{request_id}] Sending request to OpenAI endpoint: {OPENAI_API_ENDPOINT}")
+        # --- FIX: Hide full target endpoint URL and payload from logs ---
+        logger.debug(f"[{request_id}] Sending request to upstream API...")
+        # --- End Fix ---
+
         # Build the request to the target endpoint
         target_request = client.build_request(
             method="POST",
@@ -370,7 +397,7 @@ async def proxy_claude_to_openai(request: Request):
 
         # Process the response based on streaming or non-streaming
         if is_streaming:
-            logger.info(f"[{request_id}] OpenAI response is streaming. Starting SSE conversion.")
+            logger.info(f"[{request_id}] Upstream response is streaming. Starting SSE conversion.")
             return StreamingResponse(
                 stream_openai_response_to_claude_events(response, request_id, requested_model),
                 media_type="text/event-stream",
@@ -381,15 +408,18 @@ async def proxy_claude_to_openai(request: Request):
                 }
             )
         else:
-            logger.info(f"[{request_id}] OpenAI response is non-streaming. Converting.")
-            # --- FIX: Remove await here ---
+            logger.info(f"[{request_id}] Upstream response is non-streaming. Converting.")
+            # FIX: Remove await here
             openai_response_data = response.json()
-            logger.debug(f"[{request_id}] Non-streaming response from OpenAI: {json.dumps(openai_response_data)}")
+            # --- FIX: Avoid logging full response data ---
+            # logger.debug(f"[{request_id}] Non-streaming response from OpenAI: {json.dumps(openai_response_data)}")
+            logger.debug(f"[{request_id}] Received non-streaming response from upstream.")
+            # --- End Fix ---
             try:
                 claude_response_data = openai_response_to_claude_response(openai_response_data, request_id)
                 return JSONResponse(content=claude_response_data)
             except ValueError as e:
-                 logger.error(f"[{request_id}] Failed to convert OpenAI non-streaming response: {e}")
+                 logger.error(f"[{request_id}] Failed to convert upstream non-streaming response: {e}")
                  return create_error_response(500, "api_error", f"Error processing response from upstream API: {e}")
             except Exception as e:
                  logger.exception(f"[{request_id}] Unexpected error converting non-streaming response: {e}")
@@ -399,39 +429,48 @@ 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
         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
         except json.JSONDecodeError:
             # If not JSON, read as text
-            error_detail = e.response.text # Use .text instead of await .aread() as body is likely read
-        logger.error(f"[{request_id}] HTTP error from target endpoint ({status_code}): {error_detail}")
+            error_detail_text = e.response.text # Use .text instead of await .aread() as body is likely read
+
+        # --- FIX: 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]}...")
+        # --- End Fix ---
 
         # 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 Bad Request: {error_detail}"
+            err_type, msg = "invalid_request_error", f"Upstream API reported Bad Request ({status_code})."
         elif status_code == 401:
-            err_type, msg = "authentication_error", "Authentication failed with upstream API (check OPENAI_API_KEY)."
+            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: {error_detail}"
+            err_type, msg = "permission_error", f"Forbidden by upstream API ({status_code})."
         elif status_code == 429:
-            err_type, msg = "rate_limit_error", "Rate limit exceeded with upstream API."
-        elif status_code == 500:
-            err_type, msg = "api_error", "Upstream API Internal Server Error."
-        elif status_code == 503:
-             err_type, msg = "overloaded_error", "Upstream API is overloaded or unavailable."
+            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"Upstream API error ({status_code}): {error_detail}"
+            err_type, msg = "api_error", f"Received unexpected error from upstream API ({status_code})."
 
-        # Return error response immediately, even if original request was for streaming
+        # Return error response immediately
         return create_error_response(status_code, err_type, msg)
 
     except httpx.TimeoutException:
+         # --- FIX: Hide target endpoint URL from timeout log ---
         logger.error(f"[{request_id}] Request to target endpoint timed out ({READ_TIMEOUT}s).")
+        # --- End Fix ---
         return create_error_response(504, "api_error", "Gateway Timeout: Request to upstream API timed out.")
     except httpx.RequestError as e:
-        logger.error(f"[{request_id}] Network error connecting to target endpoint: {e}")
-        return create_error_response(502, "api_error", f"Bad Gateway: Network error connecting to upstream API: {e}")
+        # --- FIX: 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__}")
+        # --- End Fix ---
+        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
         return create_error_response(500, "internal_server_error", f"Internal Server Error: {e}")