Create app.py

app.py

@@ -10,6 +10,7 @@ import zipfile
 
 from fastapi import Request
 from fastapi.responses import StreamingResponse
+from fastapi.responses import Response
 from gradio import Server
 import httpx
 
@@ -30,11 +31,35 @@ def whygpu():
 # being called again (module re-import, a reload, a duplicate call, etc.).
 _llama_server_started = False
 
+LLAMA_HOST = "127.0.0.1"
+LLAMA_PORT = 8000
+LLAMA_BASE_URL = f"http://{LLAMA_HOST}:{LLAMA_PORT}"
+GRADIO_PREFIX = "/gradio_api"
+
+# Hop-by-hop headers (RFC 7230 section 6.1) describe a single connection and
+# must never be forwarded by a proxy in either direction.
+HOP_BY_HOP = {
+    "connection",
+    "keep-alive",
+    "proxy-authenticate",
+    "proxy-authorization",
+    "te",
+    "trailer",
+    "trailers",
+    "transfer-encoding",
+    "upgrade",
+}
+# Request side: also drop Host (httpx sets the right one for 127.0.0.1) and
+# Accept-Encoding / Content-Length (httpx recomputes them).
+REQUEST_DROP = HOP_BY_HOP | {"host", "content-length", "accept-encoding"}
+# Response side: Content-Length is recomputed by StreamingResponse.
+RESPONSE_DROP = HOP_BY_HOP | {"content-length"}
+
 
 def _is_llama_server_up(timeout: float = 1.0) -> bool:
     """Return True if something is already listening on the llama-server port."""
     try:
-        with socket.create_connection(("127.0.0.1", 8000), timeout=timeout):
+        with socket.create_connection((LLAMA_HOST, LLAMA_PORT), timeout=timeout):
             return True
     except OSError:
         return False
@@ -124,8 +149,8 @@ def start_llama_server():
         binary_path,
         "-m", model_path,
         "-mm", mmproj_path,
-        "--port","8000",
-        "--host", "127.0.0.1",
+        "--port", str(LLAMA_PORT),
+        "--host", LLAMA_HOST,
         "-t", "16",
         "-tb", "16",
         "-fa", "on",
@@ -153,54 +178,89 @@ def start_llama_server():
     return process
 
 
+# One shared client for the whole process: reuses keep-alive connections to
+# llama-server instead of opening a new client (and TCP pool) per request.
+_proxy_client: httpx.AsyncClient | None = None
+
+
+def _get_proxy_client() -> httpx.AsyncClient:
+    global _proxy_client
+    if _proxy_client is None or _proxy_client.is_closed:
+        _proxy_client = httpx.AsyncClient(
+            base_url=LLAMA_BASE_URL,
+            timeout=httpx.Timeout(600.0, connect=10.0),
+            limits=httpx.Limits(max_connections=64, max_keepalive_connections=16),
+            follow_redirects=False,
+        )
+    return _proxy_client
+
+
 @app.middleware("http")
 async def proxy_middleware(request: Request, call_next):
     path = request.url.path
-    if path.startswith("/gradio_api"):
+    if path.startswith(GRADIO_PREFIX):
         return await call_next(request)
 
-    url = f"http://127.0.0.1:8000{path}"
+    # Keep the raw query string so encoding and repeated keys are preserved.
+    target = path
+    if request.url.query:
+        target = f"{path}?{request.url.query}"
 
-    headers = {
-        k: v for k, v in request.headers.items()
-        if k.lower() not in ("host", "content-length", "accept-encoding")
-    }
+    # Use a list of tuples so duplicate request headers are not collapsed.
+    headers = [
+        (k, v) for k, v in request.headers.items()
+        if k.lower() not in REQUEST_DROP
+    ]
     body = await request.body()
 
-    client = httpx.AsyncClient(timeout=httpx.Timeout(600.0, connect=120.0))
+    client = _get_proxy_client()
     req = client.build_request(
         method=request.method,
-        url=url,
+        url=target,
         headers=headers,
-        params=request.query_params,
         content=body,
     )
 
     try:
-        response = await client.send(req, stream=True)
-    except Exception as e:
-        await client.aclose()
-        return StreamingResponse(iter([f"Proxy Error: {e}".encode()]), status_code=502)
+        upstream = await client.send(req, stream=True)
+    except httpx.ConnectError:
+        # llama-server is not running / not listening yet.
+        return Response(
+            content=f"Proxy Error: llama-server is not reachable at {LLAMA_BASE_URL}",
+            status_code=503,
+            media_type="text/plain",
+        )
+    except httpx.HTTPError as e:
+        return Response(
+            content=f"Proxy Error: {e}",
+            status_code=502,
+            media_type="text/plain",
+        )
 
-    async def stream_and_close():
+    async def stream_body():
         try:
-            async for chunk in response.aiter_raw():
+            async for chunk in upstream.aiter_raw():
                 yield chunk
+        except httpx.HTTPError as e:
+            # Upstream dropped mid-stream; we cannot change the status now.
+            print(f"[proxy] upstream stream error for {target}: {e}")
         finally:
-            await response.aclose()
-            await client.aclose()
-
-    res_headers = {
-        k: v for k, v in response.headers.items()
-        if k.lower() not in ("content-length", "transfer-encoding", "connection")
-    }
-
-    return StreamingResponse(
-        stream_and_close(),
-        status_code=response.status_code,
-        headers=res_headers
+            await upstream.aclose()
+
+    response = StreamingResponse(
+        stream_body(),
+        status_code=upstream.status_code,
     )
 
+    # Append headers one by one (not via a dict) so multiple values, such as
+    # several Set-Cookie headers, are all forwarded.
+    for k, v in upstream.headers.multi_items():
+        if k.lower() in RESPONSE_DROP:
+            continue
+        response.raw_headers.append((k.lower().encode("latin-1"), v.encode("latin-1")))
+
+    return response
+
 
 if __name__ == "__main__":
     start_llama_server()