| import asyncio |
| import json |
| import os |
| import time |
| from contextlib import asynccontextmanager |
|
|
| import httpx |
| from fastapi import FastAPI, Request, Response |
| from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse |
|
|
| LLAMA_HOST = os.getenv("LLAMA_HOST", "127.0.0.1") |
| LLAMA_PORT = int(os.getenv("LLAMA_PORT", "8080")) |
| LLAMA_URL = f"http://{LLAMA_HOST}:{LLAMA_PORT}" |
|
|
| HOP_BY_HOP = { |
| "content-length","transfer-encoding","content-encoding","connection", |
| "keep-alive","proxy-authenticate","proxy-authorization","te","trailers","upgrade", |
| } |
|
|
| def clean_headers(headers): |
| return {k: v for k, v in headers.items() if k.lower() not in HOP_BY_HOP} |
|
|
| async def wait_for_llama(timeout: float = 600.0): |
| start = time.time() |
| async with httpx.AsyncClient() as client: |
| while time.time() - start < timeout: |
| try: |
| r = await client.get(f"{LLAMA_URL}/health", timeout=2) |
| if r.status_code == 200: |
| return True |
| except Exception: |
| pass |
| await asyncio.sleep(1) |
| return False |
|
|
| @asynccontextmanager |
| async def lifespan(app: FastAPI): |
| await wait_for_llama() |
| yield |
|
|
| app = FastAPI(lifespan=lifespan) |
| http_client = httpx.AsyncClient(base_url=LLAMA_URL, timeout=None) |
|
|
| CHAT_HTML_PATH = os.path.join(os.path.dirname(__file__), "chat.html") |
| try: |
| with open(CHAT_HTML_PATH, "r", encoding="utf-8") as _f: |
| CHAT_HTML = _f.read() |
| except Exception: |
| CHAT_HTML = "<h1>Chat UI not found</h1>" |
|
|
| @app.get("/health") |
| async def health(): |
| return {"status": "ok"} |
|
|
| @app.get("/", response_class=HTMLResponse) |
| async def root(): |
| return HTMLResponse(CHAT_HTML) |
|
|
| @app.get("/api-info") |
| async def api_info(): |
| return JSONResponse({"status": "ok", "llama_server": LLAMA_URL}) |
|
|
| @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "HEAD", "PATCH"]) |
| async def proxy(request: Request, path: str): |
| url = httpx.URL(path="/" + path, query=request.url.query.encode("utf-8")) |
| headers = clean_headers(dict(request.headers)) |
| headers.pop("host", None) |
| body = await request.body() |
| is_stream = False |
| if request.method == "POST" and path.startswith("v1/"): |
| try: |
| payload = json.loads(body) |
| payload.pop("model", None) |
| is_stream = bool(payload.get("stream", False)) |
| body = json.dumps(payload).encode() |
| except Exception: |
| pass |
| if is_stream: |
| async def event_stream(): |
| async with http_client.stream(request.method, url, headers=headers, content=body) as upstream: |
| async for chunk in upstream.aiter_raw(): |
| yield chunk |
| return StreamingResponse(event_stream(), media_type="text/event-stream") |
| upstream = await http_client.request(method=request.method, url=url, headers=headers, content=body) |
| return Response(content=upstream.content, status_code=upstream.status_code, |
| headers=clean_headers(dict(upstream.headers)), media_type=upstream.headers.get("content-type")) |
|
|
