backstory/src/tests/test-context-routing.py

140 lines
5.0 KiB
Python

from fastapi import FastAPI, Request, Depends, Query
from fastapi.responses import RedirectResponse, JSONResponse
from uuid import UUID, uuid4
import logging
import traceback
from typing import Callable, Optional
from anyio.to_thread import run_sync
logger = logging.getLogger(__name__)
class RedirectToContext(Exception):
def __init__(self, url: str):
self.url = url
logger.info(f"Redirect to Context: {url}")
super().__init__(f"Redirect to Context: {url}")
class ContextRouteManager:
def __init__(self, app: FastAPI):
self.app = app
self.setup_handlers()
def setup_handlers(self):
@self.app.exception_handler(RedirectToContext)
async def handle_context_redirect(request: Request, exc: RedirectToContext):
logger.info(f"Handling redirect to {exc.url}")
return RedirectResponse(url=exc.url, status_code=307)
def ensure_context(
self, route_name: str = "context_id"
) -> Callable[[Request], Optional[UUID]]:
logger.info(f"Setting up context dependency for route parameter: {route_name}")
async def _ensure_context_dependency(request: Request) -> Optional[UUID]:
logger.info(
f"Entering ensure_context_dependency, Request URL: {request.url}"
)
logger.info(f"Path params: {request.path_params}")
path_params = request.path_params
route_value = path_params.get(route_name)
logger.info(f"route_value: {route_value!r}, type: {type(route_value)}")
if (
route_value is None
or not isinstance(route_value, str)
or not route_value.strip()
):
logger.info(f"route_value is invalid, generating new UUID")
path = request.url.path.rstrip("/")
new_context = await run_sync(uuid4)
redirect_url = f"{path}/{new_context}"
logger.info(f"Redirecting to {redirect_url}")
raise RedirectToContext(redirect_url)
logger.info(f"Attempting to parse route_value as UUID: {route_value}")
try:
route_context = await run_sync(UUID, route_value)
logger.info(f"Successfully parsed UUID: {route_context}")
return route_context
except ValueError as e:
logger.error(
f"Failed to parse UUID from route_value: {route_value!r}, error: {str(e)}"
)
path = request.url.path.rstrip("/")
new_context = await run_sync(uuid4)
redirect_url = f"{path}/{new_context}"
logger.info(f"Invalid UUID, redirecting to {redirect_url}")
raise RedirectToContext(redirect_url)
return _ensure_context_dependency
def route_pattern(self, path: str, *dependencies, **kwargs):
logger.info(f"Registering route: {path}")
ensure_context = self.ensure_context()
def decorator(func):
all_dependencies = list(dependencies)
all_dependencies.append(Depends(ensure_context))
logger.info(
f"Route {path} registered with dependencies: {all_dependencies}"
)
return self.app.get(path, dependencies=all_dependencies, **kwargs)(func)
return decorator
app = FastAPI(redirect_slashes=True)
@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
logger.error(f"Unhandled exception: {str(exc)}")
logger.error(f"Request URL: {request.url}, Path params: {request.path_params}")
logger.error(f"Stack trace: {''.join(traceback.format_tb(exc.__traceback__))}")
return JSONResponse(
status_code=500, content={"error": "Internal server error", "detail": str(exc)}
)
@app.middleware("http")
async def log_requests(request: Request, call_next):
logger.info(
f"Incoming request: {request.method} {request.url}, Path params: {request.path_params}"
)
response = await call_next(request)
return response
context_router = ContextRouteManager(app)
@context_router.route_pattern("/api/history/{context_id}")
async def get_history(
request: Request,
context_id: UUID = Depends(context_router.ensure_context()),
agent_type: str = Query(..., description="Type of agent to retrieve history for"),
):
logger.info(f"{request.method} {request.url.path} with context_id: {context_id}")
return {"context_id": str(context_id), "agent_type": agent_type}
@app.get("/api/history")
async def redirect_history(
request: Request,
agent_type: str = Query(..., description="Type of agent to retrieve history for"),
):
path = request.url.path.rstrip("/")
new_context = uuid4()
redirect_url = f"{path}/{new_context}?agent_type={agent_type}"
logger.info(f"Redirecting from /api/history to {redirect_url}")
return RedirectResponse(url=redirect_url, status_code=307)
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8900)