Update main.py

main.py

@@ -15,7 +15,6 @@ async def lifespan(app: FastAPI):
     """
     Manages the lifecycle of the HTTPX client.
     The client is created on startup and gracefully closed on shutdown.
-
     WARNING: This client has no timeout and no explicit connection pool limits.
     """
     # timeout=None disables all client-side timeouts.
@@ -30,8 +29,8 @@ app = FastAPI(docs_url=None, redoc_url=None, lifespan=lifespan)
 # --- Reverse Proxy Logic ---
 async def _reverse_proxy(request: Request):
     """
-    Forwards a request specifically for the /chat endpoint to the target URL.
-    It injects required headers and strips any user-provided Authorization header.
+    Forwards a request specifically for the /completions endpoint to the target URL.
+    It injects required headers and allows for a user-provided Authorization header.
     """
     client: httpx.AsyncClient = request.app.state.http_client
 
@@ -42,13 +41,14 @@ async def _reverse_proxy(request: Request):
     # Start with headers from the incoming request.
     request_headers = dict(request.headers)
 
-    # 1. CRITICAL: Remove host and authorization headers.
+    # 1. CRITICAL: Remove host header.
     #    The 'host' header is managed by httpx.
-    #    Removing 'authorization' prevents the user's key from reaching the backend.
     request_headers.pop("host", None)
-    request_headers.pop("authorization", None)
 
-    # 2. Set the specific, required headers for the target API.
+    # 2. Get the user's authorization key from the incoming request.
+    authorization_header = request.headers.get("authorization")
+
+    # 3. Set the specific, required headers for the target API.
     #    This will overwrite any conflicting headers from the original request.
     specific_headers = {
         "accept": "application/json, text/plain, */*",
@@ -67,6 +67,10 @@ async def _reverse_proxy(request: Request):
     }
     request_headers.update(specific_headers)
 
+    # 4. Add the user's authorization key to the headers if it exists.
+    if authorization_header:
+        request_headers["authorization"] = authorization_header
+
     # Build the final request to the target service.
     rp_req = client.build_request(
         method=request.method,
@@ -108,4 +112,4 @@ async def health_check():
     """Provides a basic health check endpoint."""
     return {"status": "ok", "proxying_endpoint": "/completions", "target": "TypeGPT"}
 
-# Any request to a path other than "/chat" or "/" will result in a 404 Not Found.
+# Any request to a path other than "/completions" or "/" will result in a 404 Not Found.