diff --git a/axon/main.py b/axon/main.py new file mode 100644 index 0000000..c92eb52 --- /dev/null +++ b/axon/main.py @@ -0,0 +1,65 @@ +from contextlib import asynccontextmanager + +from fastapi import FastAPI +from fastapi.middleware.cors import CORSMiddleware + +from axon.middleware.auth import APIKeyMiddleware +from axon.middleware.rate_limit import RateLimitMiddleware +from axon.middleware.request_context import RequestContextMiddleware +from axon.models import MetricsResponse +from axon.routers import crawl, extract, index, jobs, map, research, search, tavily, validate +from axon.services.cache import RedisCache, get_cache +from axon.services.metrics import metrics +from axon.services.reranker import get_reranker, load_reranker +from axon.services.index import init_index + + +@asynccontextmanager +async def lifespan(app: FastAPI): + await init_index() + load_reranker() + yield + + +app = FastAPI(title="Axon Search API", version="1.0.0", lifespan=lifespan) + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_methods=["*"], + allow_headers=["*"], +) +app.add_middleware(RequestContextMiddleware) +app.add_middleware(RateLimitMiddleware) +app.add_middleware(APIKeyMiddleware) + +app.include_router(search.router) +app.include_router(extract.router) +app.include_router(crawl.router) +app.include_router(map.router) +app.include_router(index.router) +app.include_router(jobs.router) +app.include_router(research.router) +app.include_router(tavily.router) +app.include_router(validate.router) + + +@app.get("/health") +async def health() -> dict[str, str]: + cache = await get_cache() + redis_status = ( + "connected" + if isinstance(cache, RedisCache) and await cache.ping() + else "in-memory fallback" + ) + reranker_status = "loaded" if get_reranker() is not None else "not loaded" + return { + "status": "ok", + "redis": redis_status, + "reranker": reranker_status, + } + + +@app.get("/metrics", response_model=MetricsResponse) +async def metrics_endpoint() -> MetricsResponse: + return MetricsResponse(**await metrics.snapshot())