mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
import os
|
|
from collections.abc import AsyncGenerator, AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
from typing import Final
|
|
|
|
import httpx
|
|
from fastapi import FastAPI, Request
|
|
from headers import clean_headers
|
|
from starlette.responses import Response, StreamingResponse
|
|
|
|
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
|
|
|
BACKEND: Final = os.environ.get("BENCH_BACKEND", "http://127.0.0.1:19400")
|
|
RUST: Final = os.environ.get("BENCH_RUST", "http://127.0.0.1:19401")
|
|
DISPATCH: Final = os.environ.get("BENCH_DISPATCH") == "1"
|
|
STREAM_REQUESTS: Final = os.environ.get("BENCH_STREAM_REQUESTS") == "1"
|
|
|
|
|
|
def create_app() -> FastAPI:
|
|
handler: Final = AsyncHTTPHandler(timeout=httpx.Timeout(30.0, connect=1.0), follow_redirects=False)
|
|
client: Final = handler.client
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI) -> AsyncGenerator[None]:
|
|
yield
|
|
await client.aclose()
|
|
|
|
app: Final = FastAPI(lifespan=lifespan)
|
|
|
|
async def forward(request: Request, path: str) -> Response:
|
|
data: Final = request.stream() if STREAM_REQUESTS else await request.body()
|
|
if isinstance(data, bytes) and len(data) > 8 * 1024 * 1024:
|
|
return Response(status_code=413)
|
|
if int(request.headers.get("content-length", "0")) > 8 * 1024 * 1024:
|
|
return Response(status_code=413)
|
|
target: Final = (
|
|
RUST if DISPATCH and path == "v1/messages" and "x-bench-python" not in request.headers else BACKEND
|
|
)
|
|
url: Final = f"{target}/{path}" + (f"?{request.url.query}" if request.url.query else "")
|
|
upstream_request: Final = client.build_request(
|
|
request.method,
|
|
url,
|
|
headers=clean_headers(httpx.Headers(request.headers)),
|
|
content=data,
|
|
)
|
|
try:
|
|
upstream: Final = await client.send(upstream_request, stream=True)
|
|
except httpx.HTTPError:
|
|
return Response(status_code=502)
|
|
|
|
async def chunks() -> AsyncIterator[bytes]:
|
|
try:
|
|
async for chunk in upstream.aiter_raw():
|
|
yield chunk
|
|
finally:
|
|
await upstream.aclose()
|
|
|
|
return StreamingResponse(
|
|
chunks(),
|
|
status_code=upstream.status_code,
|
|
headers={
|
|
key: value for key, value in clean_headers(upstream.headers).items() if key not in ("date", "server")
|
|
},
|
|
)
|
|
|
|
app.add_api_route("/{path:path}", forward, methods=["GET", "POST"])
|
|
return app
|
|
|
|
|
|
app: Final = create_app()
|