apipilot/tests/conftest.py

from __future__ import annotations

from typing import Any

from fastapi import APIRouter, FastAPI, HTTPException, Query

ITEMS: dict[int, dict[str, Any]] = {
    1: {"id": 1, "name": "Item A", "category": "active", "price": 9.99},
    2: {"id": 2, "name": "Item B", "category": "active", "price": 19.99},
    3: {"id": 3, "name": "Item C", "category": "archived", "price": 29.99},
}

USERS = {
    "alice-token": {"id": "alice", "role": "admin"},
    "bob-token": {"id": "bob", "role": "user"},
}


def get_current_user_from_header(authorization: str | None = None) -> dict[str, Any]:
    if not authorization:
        raise HTTPException(status_code=401, detail="Missing authorization")
    prefix = "Bearer "
    token = authorization[len(prefix):] if authorization.startswith(prefix) else authorization
    user = USERS.get(token)
    if not user:
        raise HTTPException(status_code=401, detail="Invalid token")
    return user


def build_sample_app(
    *,
    expose_http_endpoints: bool = False,
    add_composite: bool = False,
    add_deny: bool = False,
):
    from apipilot.fastapi import install_apipilot
    from apipilot.tool import CompositeTool

    app = FastAPI(title="SampleApp", version="1.0")
    router = APIRouter()

    @router.get(
        "/items",
        summary="Elenca gli elementi",
        openapi_extra={
            "x-apipilot": {
                "enabled": True,
                "name": "list_items",
                "description": (
                    "Cerca gli elementi disponibili. Usalo quando "
                    "l'utente chiede un elenco o cerca elementi di una categoria."
                ),
            }
        },
    )
    async def list_items(
        category: str | None = Query(default=None, description="Categoria filtro"),
        limit: int = Query(default=10, ge=1, le=100, description="Massimo elementi"),
    ):
        data = list(ITEMS.values())
        if category:
            data = [it for it in data if it["category"] == category]
        return {"items": data[:limit], "total": len(data)}

    @router.get(
        "/items/{item_id}",
        summary="Ottieni dettaglio elemento",
        openapi_extra={
            "x-apipilot": {
                "enabled": True,
                "name": "get_item",
                "description": "Restituisce i dettagli di un elemento dato il suo ID.",
            }
        },
    )
    async def get_item(item_id: int):
        it = ITEMS.get(item_id)
        if not it:
            raise HTTPException(status_code=404, detail="Item not found")
        return it

    @router.post(
        "/items",
        summary="Crea un nuovo elemento",
        openapi_extra={
            "x-apipilot": {
                "enabled": True,
                "name": "create_item",
                "description": "Crea un nuovo elemento con nome, categoria e prezzo.",
            }
        },
    )
    async def create_item(name: str, category: str = "active", price: float = 0.0):
        new_id = max(ITEMS.keys(), default=0) + 1
        ITEMS[new_id] = {"id": new_id, "name": name, "category": category, "price": price}
        return ITEMS[new_id]

    @router.get(
        "/items/{item_id}/slow",
        summary="Dettaglio lento",
        openapi_extra={
            "x-apipilot": {
                "enabled": True,
                "name": "get_slow_item",
                "description": "Endpoint che simula un timeout per testare limiti.",
                "timeout": 0.1,
            }
        },
    )
    async def get_slow_item(item_id: int):
        import asyncio

        await asyncio.sleep(2)
        return ITEMS.get(item_id)

    @router.get(
        "/items/huge",
        summary="Genera una risposta enorme",
        openapi_extra={
            "x-apipilot": {
                "enabled": True,
                "name": "get_huge_items",
                "description": "Restituisce una lista molto grande per testare il troncamento.",
                "max_response_size": 1024,
            }
        },
    )
    async def get_huge_items():
        huge = [
            {"id": i, "name": f"Item {i}", "payload": "x" * 200}
            for i in range(100)
        ]
        return {"items": huge, "total": len(huge)}

    @router.get(
        "/not-pilot",
        summary="Non abilitato",
    )
    async def not_pilot():
        return {"ok": True}

    app.include_router(router)

    def allow_tool(user: Any, record: Any) -> bool:
        if not add_deny:
            return True
        if user is None:
            return True
        role = user.get("role") if isinstance(user, dict) else "user"
        if role != "admin" and record.name == "create_item":
            return False
        return True

    pilot = install_apipilot(
        app,
        get_current_user=get_current_user_from_header,
        allow_tool=allow_tool,
        default_timeout=5.0,
        expose_http_endpoints=expose_http_endpoints,
        require_auth_for_http=False,
    )

    if add_composite:
        async def find_cheapest(args: dict[str, Any], user: Any):
            cat = args.get("category")
            r1 = await pilot.execute(user=user, name="list_items", arguments={"category": cat})
            if not r1["success"]:
                return {"error": r1["error"]}
            items = (r1["data"] or {}).get("items", [])
            if not items:
                return None
            cheapest = min(items, key=lambda it: it.get("price", float("inf")))
            return {"cheapest": cheapest}

        pilot.add_tool(
            CompositeTool(
                name="find_cheapest_item",
                description="Trova l'elemento più economico nella categoria specificata.",
                parameters_schema={
                    "type": "object",
                    "properties": {
                        "category": {
                            "type": "string",
                            "description": "Categoria degli elementi dove cercare.",
                        }
                    },
                    "required": [],
                },
                run_fn=find_cheapest,
            )
        )

    return app, pilot