from fastapi import APIRouter

from axon.models import (
    IndexSearchRequest,
    IndexSearchResponse,
    IngestRequest,
    IngestResponse,
)
from axon.services.index import ingest_documents, search_index

router = APIRouter()


@router.post("/ingest", response_model=IngestResponse)
async def ingest(request: IngestRequest) -> IngestResponse:
    return IngestResponse(indexed=await ingest_documents(request.documents))


@router.post("/index/search", response_model=IndexSearchResponse)
async def index_search(request: IndexSearchRequest) -> IndexSearchResponse:
    return IndexSearchResponse(
        query=request.query,
        results=await search_index(request.query, request.max_results),
    )