Files
geo-setp/main.py
2026-04-15 14:05:33 +08:00

104 lines
3.8 KiB
Python

import os
from secrets import compare_digest
import httpx
from dotenv import load_dotenv
from fastapi import Depends, FastAPI, HTTPException, Security
from openai import APITimeoutError
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security.api_key import APIKeyHeader
from schemas import LoadPlanRequest, LoadPlanResult, RoutePlanRequest, RoutePlanResult
from agent import (
ConfigurationError,
GuardrailError,
LoadPlanConfigurationError,
LoadPlanGuardrailError,
LoadPlanToolConfigurationError,
LoadPlanToolRequestError,
run_load_plan,
run_route_plan,
)
load_dotenv()
api_key_header = APIKeyHeader(name="Authorization", auto_error=False)
def _csv_env(name: str, default: str) -> list[str]:
raw_value = os.getenv(name, default)
return [item.strip() for item in raw_value.split(",") if item.strip()]
def _bool_env(name: str, default: bool) -> bool:
raw_value = os.getenv(name, "true" if default else "false").strip().lower()
return raw_value in {"1", "true", "yes", "on"}
def _required_env(name: str) -> str:
value = os.getenv(name, "").strip()
if not value:
raise HTTPException(status_code=503, detail=f"Missing required environment variable: {name}")
return value
def require_authorization(api_key: str | None = Security(api_key_header)) -> None:
expected_key = _required_env("AGENT_HTTP_AUTH_KEY")
if api_key is None or not compare_digest(api_key, expected_key):
raise HTTPException(status_code=401, detail="Unauthorized")
app = FastAPI(title="Geo Route Agent", version="0.1.0")
app.add_middleware(
CORSMiddleware,
allow_origins=_csv_env("CORS_ALLOW_ORIGINS", "http://localhost,http://127.0.0.1"),
allow_origin_regex=os.getenv(
"CORS_ALLOW_ORIGIN_REGEX",
r"https?://(localhost|127\.0\.0\.1)(:\d+)?$",
),
allow_credentials=_bool_env("CORS_ALLOW_CREDENTIALS", False),
allow_methods=_csv_env("CORS_ALLOW_METHODS", "GET,POST,OPTIONS"),
allow_headers=_csv_env("CORS_ALLOW_HEADERS", "*"),
)
@app.post("/route/plan", response_model=RoutePlanResult, dependencies=[Depends(require_authorization)])
async def route_plan(request: RoutePlanRequest) -> RoutePlanResult:
try:
return await run_route_plan(request)
except ConfigurationError as exc:
raise HTTPException(status_code=503, detail=str(exc)) from exc
except GuardrailError as exc:
raise HTTPException(status_code=422, detail=str(exc)) from exc
except (httpx.TimeoutException, APITimeoutError, TimeoutError) as exc:
raise HTTPException(status_code=504, detail=f"Upstream request timed out: {exc}") from exc
except Exception as exc:
raise HTTPException(status_code=500, detail=str(exc)) from exc
@app.post("/load/plan", response_model=LoadPlanResult, dependencies=[Depends(require_authorization)])
async def load_plan(request: LoadPlanRequest) -> LoadPlanResult:
try:
return await run_load_plan(request)
except (LoadPlanConfigurationError, LoadPlanToolConfigurationError) as exc:
raise HTTPException(status_code=503, detail=str(exc)) from exc
except LoadPlanGuardrailError as exc:
raise HTTPException(status_code=422, detail=str(exc)) from exc
except LoadPlanToolRequestError as exc:
raise HTTPException(status_code=502, detail=str(exc)) from exc
except (httpx.TimeoutException, APITimeoutError, TimeoutError) as exc:
raise HTTPException(status_code=504, detail=f"Upstream request timed out: {exc}") from exc
except Exception as exc:
raise HTTPException(status_code=500, detail=str(exc)) from exc
@app.get("/healthz")
async def healthz() -> dict[str, str]:
return {"status": "ok"}
if __name__ == "__main__":
import uvicorn
uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)