proxy_server.py

"""
Reverse Proxy for CatCon Space
Routes API/inference calls to FastAPI, everything else to Streamlit
"""

import asyncio
from aiohttp import web, client

FASTAPI_URL = "http://127.0.0.1:8000"
STREAMLIT_URL = "http://127.0.0.1:8501"

API_PATHS = ('/api/', '/inference/', '/health', '/docs', '/openapi.json',
             '/admin/', '/redoc', '/swagger-ui')

async def proxy_handler(request):
    path = request.path_qs
    is_api = any(path.startswith(p) or path == p.rstrip('/') for p in API_PATHS)
    target = f"{FASTAPI_URL}{path}" if is_api else f"{STREAMLIT_URL}{path}"

    # Build headers (preserve host for Streamlit)
    headers = {k: v for k, v in request.headers.items() if k.lower() not in ('host',)}

    try:
        async with client.ClientSession() as session:
            body = await request.read()
            async with session.request(
                method=request.method,
                url=target,
                headers=headers,
                data=body if body else None
            ) as resp:
                data = await resp.read()
                response_headers = {k: v for k, v in resp.headers.items()
                                     if k.lower() not in ('content-encoding', 'transfer-encoding', 'content-length')}
                return web.Response(body=data, status=resp.status, headers=response_headers)
    except Exception as e:
        return web.Response(text=f"Proxy error: {e}", status=502)

app = web.Application()
app.router.add_route('*', '/{path:.*}', proxy_handler)

if __name__ == "__main__":
    print("[proxy] Starting on port 7860...")
    web.run_app(app, host="0.0.0.0", port=7860)