File size: 2,747 Bytes
bc0b11f 9642bd4 bc0b11f 9642bd4 bc0b11f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | import asyncio
import os
import signal
import subprocess
import sys
from contextlib import asynccontextmanager
import httpx
import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import Response, StreamingResponse
MANAGER_URL = "http://127.0.0.1:7894"
API_URL = "http://127.0.0.1:7895"
processes = []
def start_process(command):
env = os.environ.copy()
env.setdefault("DB_PATH", "/data/privy_manager.db")
return subprocess.Popen(command, env=env)
@asynccontextmanager
async def lifespan(app: FastAPI):
processes.append(start_process([sys.executable, "src/privy_manager.py"]))
processes.append(start_process([sys.executable, "src/api_server.py"]))
await asyncio.sleep(3)
try:
yield
finally:
for proc in processes:
if proc.poll() is None:
proc.send_signal(signal.SIGTERM)
for proc in processes:
try:
proc.wait(timeout=5)
except subprocess.TimeoutExpired:
proc.kill()
app = FastAPI(lifespan=lifespan)
async def proxy(request: Request, upstream: str, path: str):
url = f"{upstream}/{path}" if path else upstream
if request.url.query:
url = f"{url}?{request.url.query}"
headers = dict(request.headers)
headers.pop("host", None)
body = await request.body()
client = httpx.AsyncClient(timeout=None, follow_redirects=False)
req = client.build_request(
request.method,
url,
headers=headers,
content=body,
)
upstream_resp = await client.send(req, stream=True)
excluded = {"content-encoding", "transfer-encoding", "connection"}
response_headers = {k: v for k, v in upstream_resp.headers.items() if k.lower() not in excluded}
async def stream_response():
try:
async for chunk in upstream_resp.aiter_raw():
yield chunk
finally:
await upstream_resp.aclose()
await client.aclose()
return StreamingResponse(
stream_response(),
status_code=upstream_resp.status_code,
headers=response_headers,
media_type=upstream_resp.headers.get("content-type"),
)
@app.api_route("/v1/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"])
async def api_proxy(request: Request, path: str):
return await proxy(request, API_URL, f"v1/{path}")
@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"])
async def manager_proxy(request: Request, path: str):
return await proxy(request, MANAGER_URL, path)
if __name__ == "__main__":
port = int(os.environ.get("PORT", "7860"))
uvicorn.run(app, host="0.0.0.0", port=port)
|