| import asyncio |
| import os |
| import json |
| import logging |
| from contextlib import asynccontextmanager |
| from fastapi import FastAPI, Request, Response |
| from fastapi.responses import StreamingResponse, JSONResponse |
| import aiohttp |
|
|
| |
| TARGET_BASE_URL = "https://api.deepinfra.com/v1/openai" |
| REAL_API_KEY = os.environ["REAL_API_KEY"] |
| FAKE_KEY = os.environ.get("FAKE_KEY", "123456") |
|
|
| |
| LOCAL_PROXY = None |
|
|
| |
|
|
| logging.basicConfig(level=logging.INFO) |
| logger = logging.getLogger("proxy") |
|
|
| @asynccontextmanager |
| async def lifespan(app: FastAPI): |
| app.state.session = aiohttp.ClientSession() |
| yield |
| if app.state.session: |
| await app.state.session.close() |
|
|
| app = FastAPI(lifespan=lifespan) |
|
|
| @app.api_route("/v1", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"]) |
| @app.api_route("/v1/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"]) |
| async def proxy_endpoint(path: str = "", request: Request = None): |
| if request.method == "OPTIONS": |
| return Response(status_code=200, headers={ |
| "Access-Control-Allow-Origin": "*", |
| "Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, OPTIONS", |
| "Access-Control-Allow-Headers": "Authorization, Content-Type" |
| }) |
|
|
| |
| auth_header = request.headers.get("authorization", "") |
| if auth_header.lower() != f"bearer {FAKE_KEY}": |
| return JSONResponse( |
| status_code=401, |
| content={"error": {"message": "Invalid local API key", "type": "invalid_request_error"}} |
| ) |
|
|
| target_url = f"{TARGET_BASE_URL.rstrip('/')}/{path}" if path else TARGET_BASE_URL.rstrip('/') |
|
|
| |
| headers = {} |
| for key, value in request.headers.items(): |
| if key.lower() in ["host", "content-length", "connection", "authorization", "transfer-encoding"]: |
| continue |
| headers[key] = value |
| headers["Authorization"] = f"Bearer {REAL_API_KEY}" |
|
|
| body = await request.body() |
| query_params = dict(request.query_params) |
|
|
| |
| is_stream = False |
| if body: |
| try: |
| is_stream = json.loads(body).get("stream", False) |
| except json.JSONDecodeError: |
| pass |
|
|
| session = request.app.state.session |
| timeout = aiohttp.ClientTimeout(total=None, connect=10, sock_read=60) |
|
|
| |
| resp = await session.request( |
| method=request.method, |
| url=target_url, |
| headers=headers, |
| data=body, |
| params=query_params, |
| timeout=timeout, |
| proxy=LOCAL_PROXY |
| ) |
|
|
| |
| if is_stream or "text/event-stream" in resp.headers.get("content-type", ""): |
| async def stream_generator(): |
| try: |
| async for chunk in resp.content.iter_any(): |
| yield chunk |
| except (aiohttp.ClientConnectionError, aiohttp.ServerDisconnectedError, |
| aiohttp.ClientPayloadError, TimeoutError, asyncio.TimeoutError, |
| asyncio.CancelledError): |
| pass |
| finally: |
| resp.close() |
|
|
| resp_headers = { |
| "content-type": resp.headers.get("content-type", "text/event-stream"), |
| "cache-control": "no-cache", |
| "connection": "keep-alive", |
| "access-control-allow-origin": "*" |
| } |
| return StreamingResponse( |
| stream_generator(), |
| media_type=resp_headers["content-type"], |
| headers=resp_headers |
| ) |
|
|
| |
| else: |
| content = await resp.read() |
| resp.close() |
|
|
| |
| logger.warning(f"UPSTREAM STATUS: {resp.status}") |
| logger.warning(f"UPSTREAM BODY LENGTH: {len(content)}") |
| if content: |
| logger.warning(f"UPSTREAM BODY PREVIEW: {content[:200]}") |
|
|
| resp_headers = { |
| "content-type": resp.headers.get("content-type", "application/json"), |
| "access-control-allow-origin": "*" |
| } |
| |
| for k in ["content-encoding", "transfer-encoding"]: |
| resp_headers.pop(k, None) |
|
|
| return Response(content=content, status_code=resp.status, headers=resp_headers) |
