main.py

import httpx
from fastapi import FastAPI, Request, HTTPException
from starlette.responses import StreamingResponse
from starlette.background import BackgroundTask
import os
from contextlib import asynccontextmanager

# --- Configuration ---
# The target URL is configurable via an environment variable.
TARGET_URL = os.getenv("TARGET_URL", "https://api.gmi-serving.com/v1/chat")

# --- HTTPX Client Lifecycle Management ---
@asynccontextmanager
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.
    # The absence of a `limits` parameter means we rely on system defaults.
    async with httpx.AsyncClient(base_url=TARGET_URL, timeout=None) as client:
        app.state.http_client = client
        yield

# Initialize the FastAPI app with the lifespan manager and disable default docs
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 /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

    # Construct the URL for the outgoing request using the incoming path and query.
    url = httpx.URL(path=request.url.path, query=request.url.query.encode("utf-8"))

    # --- Header Processing ---
    # Start with headers from the incoming request.
    request_headers = dict(request.headers)

    # 1. CRITICAL: Remove host header.
    #    The 'host' header is managed by httpx.
    request_headers.pop("host", None)

    # 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, */*",
        "accept-language": "en-US,en;q=0.9,ru;q=0.8",
        "content-type": "application/json",
        "origin": "https://console.gmicloud.ai",
        "priority": "u=1, i",
        "referer": "https://console.gmicloud.ai/playground/llm/deepseek-r1-0528/01da5dd6-aa6a-40cb-9dbd-241467aa5cbb?tab=playground",
        "sec-ch-ua": '"Not;A=Brand";v="99", "Google Chrome";v="139", "Chromium";v="139"',
        "sec-ch-ua-mobile": "?0",
        "sec-ch-ua-platform": '"Windows"',
        "sec-fetch-dest": "empty",
        "sec-fetch-mode": "cors",
        "sec-fetch-site": "same-origin",
        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/139.0.0.0 Safari/537.36",
    }
    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,
        url=url,
        headers=request_headers,
        content=await request.body(),
    )

    try:
        # Send the request and get a streaming response.
        rp_resp = await client.send(rp_req, stream=True)
    except httpx.ConnectError as e:
        # This error occurs if the target service is down or unreachable.
        raise HTTPException(status_code=502, detail=f"Bad Gateway: Cannot connect to target service. {e}")

    # Stream the response from the target service back to the original client.
    return StreamingResponse(
        rp_resp.aiter_raw(),
        status_code=rp_resp.status_code,
        headers=rp_resp.headers,
        background=BackgroundTask(rp_resp.aclose),
    )

# --- API Endpoint ---
@app.api_route(
    "/completions",
    methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"]
)
async def chat_proxy_handler(request: Request):
    """
    This endpoint captures requests specifically for the "/completions" path
    and forwards them through the reverse proxy.
    """
    return await _reverse_proxy(request)

# A simple root endpoint for health checks.
@app.get("/")
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 "/completions" or "/" will result in a 404 Not Found.