|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import json |
| 4 | +from collections.abc import Callable |
| 5 | +from urllib.parse import parse_qs |
| 6 | + |
| 7 | +from starlette.applications import Starlette |
| 8 | +from starlette.requests import Request |
| 9 | +from starlette.responses import Response, StreamingResponse |
| 10 | +from starlette.routing import Route |
| 11 | + |
| 12 | +from .runtime import RuntimeServices, default_workspace_root |
| 13 | +from .service import ApiResponse, dispatch_get, dispatch_post, json_response |
| 14 | + |
| 15 | + |
| 16 | +def _starlette_response(response: ApiResponse) -> Response: |
| 17 | + if response.stream is not None: |
| 18 | + return StreamingResponse(response.stream, status_code=response.status, media_type=response.content_type, headers=response.headers) |
| 19 | + payload = response.body or b"" |
| 20 | + return Response(payload, status_code=response.status, media_type=response.content_type, headers=response.headers) |
| 21 | + |
| 22 | + |
| 23 | +async def _dispatch(request: Request) -> Response: |
| 24 | + runtime: RuntimeServices = request.app.state.runtime |
| 25 | + model_name: str = request.app.state.model_name |
| 26 | + get_dispatcher: Callable[..., ApiResponse] = getattr(request.app.state, "get_dispatcher", dispatch_get) |
| 27 | + post_dispatcher: Callable[..., ApiResponse] = getattr(request.app.state, "post_dispatcher", dispatch_post) |
| 28 | + workspace_root_provider: Callable[[], str] = getattr( |
| 29 | + request.app.state, |
| 30 | + "workspace_root_provider", |
| 31 | + default_workspace_root, |
| 32 | + ) |
| 33 | + if request.method == "GET": |
| 34 | + response = get_dispatcher( |
| 35 | + path=request.url.path, |
| 36 | + query=parse_qs(request.url.query), |
| 37 | + runtime=runtime, |
| 38 | + model_name=model_name, |
| 39 | + ) |
| 40 | + return _starlette_response(response) |
| 41 | + if request.method == "POST": |
| 42 | + raw_body = await request.body() |
| 43 | + if not raw_body: |
| 44 | + return _starlette_response(json_response(400, {"error": "empty body"})) |
| 45 | + try: |
| 46 | + body = json.loads(raw_body) |
| 47 | + except json.JSONDecodeError: |
| 48 | + return _starlette_response(json_response(400, {"error": "invalid JSON"})) |
| 49 | + response = post_dispatcher( |
| 50 | + path=request.url.path, |
| 51 | + body=body, |
| 52 | + headers=dict(request.headers.items()), |
| 53 | + runtime=runtime, |
| 54 | + model_name=model_name, |
| 55 | + workspace_root_provider=workspace_root_provider, |
| 56 | + ) |
| 57 | + return _starlette_response(response) |
| 58 | + return _starlette_response(json_response(404, {"error": "not found"})) |
| 59 | + |
| 60 | + |
| 61 | +def create_api_app( |
| 62 | + *, |
| 63 | + runtime: RuntimeServices, |
| 64 | + model_name: str, |
| 65 | + get_dispatcher: Callable[..., ApiResponse] = dispatch_get, |
| 66 | + post_dispatcher: Callable[..., ApiResponse] = dispatch_post, |
| 67 | + workspace_root_provider: Callable[[], str] = default_workspace_root, |
| 68 | +) -> Starlette: |
| 69 | + app = Starlette( |
| 70 | + debug=False, |
| 71 | + routes=[ |
| 72 | + Route("/", _dispatch, methods=["GET", "POST"]), |
| 73 | + Route("/{path:path}", _dispatch, methods=["GET", "POST"]), |
| 74 | + ], |
| 75 | + ) |
| 76 | + app.state.runtime = runtime |
| 77 | + app.state.model_name = model_name |
| 78 | + app.state.get_dispatcher = get_dispatcher |
| 79 | + app.state.post_dispatcher = post_dispatcher |
| 80 | + app.state.workspace_root_provider = workspace_root_provider |
| 81 | + return app |
0 commit comments