nucliadb-agentic-api 1.0.0.post156__py3-none-any.whl

nucliadb_agentic_api/v1/agentic_configs.py

@@ -0,0 +1,117 @@
+from typing import TYPE_CHECKING, cast
+
+from fastapi import Header, HTTPException, Request, Response
+from nucliadb_models.resource import NucliaDBRoles
+from nucliadb_utils.authentication import requires
+
+from nucliadb_agentic_api import exceptions
+from nucliadb_agentic_api.models import AgenticConfigSchema
+from nucliadb_agentic_api.v1.router import router
+
+if TYPE_CHECKING:
+    from nucliadb_agentic_api.app import HTTPApplication
+
+
+@router.post(
+    "/api/v1/kb/{kbid}/agentic_configs/{agentic_id}",
+    status_code=201,
+    summary="Create agentic configuration",
+    tags=["Agentic configs"],
+)
+@requires(NucliaDBRoles.OWNER)
+async def create_agentic_config_endpoint(
+    request: Request,
+    kbid: str,
+    agentic_id: str,
+    item: AgenticConfigSchema,
+    x_nucliadb_account: str = Header(default="", include_in_schema=False),
+) -> Response:
+    app = cast("HTTPApplication", request.app)
+    try:
+        await app.agent_manager.create_agentic_config(
+            account=x_nucliadb_account,
+            kbid=kbid,
+            agentic_id=agentic_id,
+            config=item,
+        )
+    except exceptions.Conflict as exc:
+        raise HTTPException(status_code=409, detail=str(exc)) from exc
+    except exceptions.InvalidReference as exc:
+        raise HTTPException(status_code=422, detail=str(exc)) from exc
+    return Response(status_code=201)
+
+
+@router.get(
+    "/api/v1/kb/{kbid}/agentic_configs/{agentic_id}",
+    status_code=200,
+    summary="Get agentic configuration",
+    tags=["Agentic configs"],
+    response_model=AgenticConfigSchema,
+    response_model_exclude_none=True,
+)
+@requires(NucliaDBRoles.READER)
+async def get_agentic_config_endpoint(
+    request: Request,
+    kbid: str,
+    agentic_id: str,
+    x_nucliadb_account: str = Header(default="", include_in_schema=False),
+) -> AgenticConfigSchema:
+    app = cast("HTTPApplication", request.app)
+    try:
+        return await app.agent_manager.get_agentic_config(
+            account=x_nucliadb_account,
+            kbid=kbid,
+            agentic_id=agentic_id,
+        )
+    except exceptions.NotFound as exc:
+        raise HTTPException(status_code=404, detail=str(exc)) from exc
+
+
+@router.get(
+    "/api/v1/kb/{kbid}/agentic_configs",
+    status_code=200,
+    summary="List agentic configurations",
+    tags=["Agentic configs"],
+    response_model=dict[str, AgenticConfigSchema],
+    response_model_exclude_none=True,
+)
+@requires(NucliaDBRoles.READER)
+async def list_agentic_configs_endpoint(
+    request: Request,
+    kbid: str,
+    x_nucliadb_account: str = Header(default="", include_in_schema=False),
+) -> dict[str, AgenticConfigSchema]:
+    app = cast("HTTPApplication", request.app)
+    return await app.agent_manager.list_agentic_configs(
+        account=x_nucliadb_account,
+        kbid=kbid,
+    )
+
+
+@router.patch(
+    "/api/v1/kb/{kbid}/agentic_configs/{agentic_id}",
+    status_code=204,
+    summary="Update agentic configuration",
+    tags=["Agentic configs"],
+)
+@requires(NucliaDBRoles.OWNER)
+async def patch_agentic_config_endpoint(
+    request: Request,
+    kbid: str,
+    agentic_id: str,
+    item: AgenticConfigSchema,
+    x_nucliadb_account: str = Header(default="", include_in_schema=False),
+) -> Response:
+    app = cast("HTTPApplication", request.app)
+    try:
+        await app.agent_manager.patch_agentic_config(
+            account=x_nucliadb_account,
+            kbid=kbid,
+            agentic_id=agentic_id,
+            config=item,
+        )
+    except exceptions.NotFound as exc:
+        raise HTTPException(status_code=404, detail=str(exc)) from exc
+    except exceptions.InvalidReference as exc:
+        raise HTTPException(status_code=422, detail=str(exc)) from exc
+    return Response(status_code=204)

nucliadb_agentic_api/v1/ask.py

@@ -0,0 +1,303 @@
+import json
+from uuid import UUID
+
+from fastapi import Header, Request, Response
+from hyperforge_nucliadb_agentic.ask.exceptions import (
+    AnswerJsonSchemaTooLong,
+)
+from hyperforge_nucliadb_agentic.ask.model import (
+    AskRequest,
+    SyncAskResponse,
+    parse_max_tokens,
+)
+from hyperforge_nucliadb_agentic.ask.search import rpc
+from hyperforge_nucliadb_agentic.ask.search.ask import (
+    AskResult,
+    ask,
+    handled_ask_exceptions,
+)
+from hyperforge_nucliadb_agentic.ask.utils.responses import (
+    HTTPClientError,
+)
+from nucliadb_models.configuration import AskConfig
+from nucliadb_models.resource import NucliaDBRoles
+from nucliadb_models.search import (
+    NucliaDBClientType,
+)
+from nucliadb_models.security import RequestSecurity
+from nucliadb_sdk.v2.exceptions import PreconditionFailed, UnprocessableEntity
+from nucliadb_utils.authentication import NucliaUser, requires
+from pydantic import ValidationError
+from starlette.responses import StreamingResponse
+
+from nucliadb_agentic_api.agentic.ask_handler import create_agentic_response
+from nucliadb_agentic_api.v1.router import router
+
+
+@router.post(
+    "/api/v1/kb/{kbid}/ask",
+    status_code=200,
+    summary="Ask Knowledge Box",
+    description="Ask questions on a Knowledge Box",
+    tags=["Search"],
+    response_model=SyncAskResponse,
+)
+@requires(NucliaDBRoles.READER)
+async def ask_knowledgebox_endpoint(
+    request: Request,
+    kbid: str,
+    item: AskRequest,
+    x_nucliadb_account: str = Header(default="", include_in_schema=False),
+    x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
+    x_show_consumption: bool = Header(default=False),
+    x_nucliadb_user: str = Header(""),
+    x_forwarded_for: str = Header(""),
+    x_synchronous: bool = Header(
+        default=False,
+        description="When set to true, outputs response as JSON in a non-streaming way. "
+        "This is slower and requires waiting for entire answer to be ready.",
+    ),
+) -> StreamingResponse | HTTPClientError | Response:
+    current_user: NucliaUser = request.user
+    # If present, security groups from AuthorizationBackend overrides any
+    # security group of the payload
+    if current_user.security_groups:
+        if item.security is None:
+            item.security = RequestSecurity(groups=current_user.security_groups)
+        else:
+            item.security.groups = current_user.security_groups
+
+    if item.search_configuration is not None:
+        search_config = await rpc.get_search_configuration(
+            rpc.get_sdk("reader"), kbid, name=item.search_configuration
+        )
+        if search_config is None:
+            return HTTPClientError(
+                status_code=400, detail="Search configuration not found"
+            )
+
+        if not isinstance(search_config.config, AskConfig):
+            return HTTPClientError(
+                status_code=400,
+                detail="This search configuration is not valid for `ask`",
+            )
+
+        try:
+            item = AskRequest.model_validate(
+                search_config.config.model_dump(exclude_unset=True)
+                | item.model_dump(exclude_unset=True)
+            )
+        except ValidationError as e:
+            detail = json.loads(e.json())
+            return HTTPClientError(status_code=422, detail=detail)
+
+    if item.agentic_config_id is not None:
+        return await create_agentic_response(
+            app=request.app,
+            agentic_config_id=item.agentic_config_id,
+            kbid=kbid,
+            ask_request=item,
+            user_id=x_nucliadb_user,
+            account=x_nucliadb_account,
+            client_type=x_ndb_client,
+            origin=x_forwarded_for,
+            x_synchronous=x_synchronous,
+            extra_predict_headers={
+                "X-Show-Consumption": str(x_show_consumption).lower()
+            },
+        )
+
+    return await create_ask_response(
+        kbid=kbid,
+        ask_request=item,
+        user_id=x_nucliadb_user,
+        client_type=x_ndb_client,
+        origin=x_forwarded_for,
+        x_synchronous=x_synchronous,
+        extra_predict_headers={"X-Show-Consumption": str(x_show_consumption).lower()},
+    )
+
+
+@router.post(
+    "/api/v1/kb/{kbid}/resource/{rid}/ask",
+    status_code=200,
+    summary="Ask a resource (by id)",
+    description="Ask questions to a resource",
+    tags=["Search"],
+    response_model=SyncAskResponse,
+)
+@requires(NucliaDBRoles.READER)
+async def resource_ask_endpoint_by_uuid(
+    request: Request,
+    kbid: str,
+    rid: UUID,
+    item: AskRequest,
+    x_show_consumption: bool = Header(default=False),
+    x_nucliadb_account: str = Header(default="", include_in_schema=False),
+    x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
+    x_nucliadb_user: str = Header(""),
+    x_forwarded_for: str = Header(""),
+    x_synchronous: bool = Header(
+        False,
+        description="When set to true, outputs response as JSON in a non-streaming way. "
+        "This is slower and requires waiting for entire answer to be ready.",
+    ),
+) -> StreamingResponse | HTTPClientError | Response:
+    current_user: NucliaUser = request.user
+    # If present, security groups from AuthorizationBackend overrides any
+    # security group of the payload
+    if current_user.security_groups:
+        if item.security is None:
+            item.security = RequestSecurity(groups=current_user.security_groups)
+        else:
+            item.security.groups = current_user.security_groups
+
+    if item.agentic_config_id is not None:
+        return await create_agentic_response(
+            app=request.app,
+            agentic_config_id=item.agentic_config_id,
+            kbid=kbid,
+            account=x_nucliadb_account,
+            ask_request=item,
+            user_id=x_nucliadb_user,
+            client_type=x_ndb_client,
+            origin=x_forwarded_for,
+            x_synchronous=x_synchronous,
+            extra_predict_headers={
+                "X-Show-Consumption": str(x_show_consumption).lower()
+            },
+        )
+
+    return await create_ask_response(
+        kbid=kbid,
+        ask_request=item,
+        user_id=x_nucliadb_user,
+        client_type=x_ndb_client,
+        origin=x_forwarded_for,
+        x_synchronous=x_synchronous,
+        resource=str(rid),
+        extra_predict_headers={"X-Show-Consumption": str(x_show_consumption).lower()},
+    )
+
+
+@router.post(
+    "/api/v1/kb/{kbid}/slug/{slug}/ask",
+    status_code=200,
+    summary="Ask a resource (by slug)",
+    description="Ask questions to a resource",
+    tags=["Search"],
+    response_model=SyncAskResponse,
+)
+@requires(NucliaDBRoles.READER)
+async def resource_ask_endpoint_by_slug(
+    request: Request,
+    kbid: str,
+    slug: str,
+    item: AskRequest,
+    x_show_consumption: bool = Header(default=False),
+    x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
+    x_nucliadb_user: str = Header(""),
+    x_nucliadb_account: str = Header(default="", include_in_schema=False),
+    x_forwarded_for: str = Header(""),
+    x_synchronous: bool = Header(
+        False,
+        description="When set to true, outputs response as JSON in a non-streaming way. "
+        "This is slower and requires waiting for entire answer to be ready.",
+    ),
+) -> StreamingResponse | HTTPClientError | Response:
+    resource_id = await rpc.get_resource_uuid_from_slug(
+        rpc.get_sdk("reader"), kbid, slug
+    )
+    if resource_id is None:
+        return HTTPClientError(status_code=404, detail="Resource not found")
+
+    current_user: NucliaUser = request.user
+    # If present, security groups from AuthorizationBackend overrides any
+    # security group of the payload
+    if current_user.security_groups:
+        if item.security is None:
+            item.security = RequestSecurity(groups=current_user.security_groups)
+        else:
+            item.security.groups = current_user.security_groups
+
+    if item.agentic_config_id is not None:
+        return await create_agentic_response(
+            app=request.app,
+            agentic_config_id=item.agentic_config_id,
+            kbid=kbid,
+            account=x_nucliadb_account,
+            ask_request=item,
+            user_id=x_nucliadb_user,
+            client_type=x_ndb_client,
+            origin=x_forwarded_for,
+            x_synchronous=x_synchronous,
+            resource=str(resource_id),
+            extra_predict_headers={
+                "X-Show-Consumption": str(x_show_consumption).lower()
+            },
+        )
+
+    return await create_ask_response(
+        kbid=kbid,
+        ask_request=item,
+        user_id=x_nucliadb_user,
+        client_type=x_ndb_client,
+        origin=x_forwarded_for,
+        x_synchronous=x_synchronous,
+        resource=resource_id,
+        extra_predict_headers={"X-Show-Consumption": str(x_show_consumption).lower()},
+    )
+
+
+@handled_ask_exceptions
+async def create_ask_response(
+    kbid: str,
+    ask_request: AskRequest,
+    user_id: str,
+    client_type: NucliaDBClientType,
+    origin: str,
+    x_synchronous: bool,
+    resource: str | None = None,
+    extra_predict_headers: dict[str, str] | None = None,
+) -> Response:
+    ask_request.max_tokens = parse_max_tokens(ask_request.max_tokens)
+    try:
+        ask_result: AskResult = await ask(
+            search_sdk=rpc.get_sdk("search"),
+            reader_sdk=rpc.get_sdk("reader"),
+            kbid=kbid,
+            ask_request=ask_request,
+            user_id=user_id,
+            client_type=client_type,
+            origin=origin,
+            resource=resource,
+            extra_predict_headers=extra_predict_headers,
+        )
+
+    except AnswerJsonSchemaTooLong as err:
+        return HTTPClientError(status_code=400, detail=str(err))
+
+    # forward 412 and 422 from nucliadb to the client
+    except PreconditionFailed as err:
+        return HTTPClientError(status_code=412, detail=err.message)
+    except UnprocessableEntity as err:
+        return HTTPClientError(status_code=422, detail=err.message)
+
+    headers = {
+        "NUCLIA-LEARNING-ID": ask_result.nuclia_learning_id or "unknown",
+        "Access-Control-Expose-Headers": "NUCLIA-LEARNING-ID",
+    }
+    if x_synchronous:
+        return Response(
+            content=await ask_result.json(),
+            status_code=200,
+            headers=headers,
+            media_type="application/json",
+        )
+    else:
+        return StreamingResponse(
+            content=ask_result.ndjson_stream(),
+            status_code=200,
+            headers=headers,
+            media_type="application/x-ndjson",
+        )

nucliadb_agentic_api/v1/ask_websocket.py

@@ -0,0 +1,141 @@
+import asyncio
+import json
+
+from fastapi import Header, Query, WebSocket, WebSocketDisconnect
+from hyperforge.api.authentication import requires
+from hyperforge.api.v1.interaction import WebsocketReceiver, stream_response
+from hyperforge.interaction import AnswerOperation, AragAnswer, ARAGException
+from hyperforge_nucliadb_agentic.ask.model import (
+    AskRequest,
+)
+from hyperforge_nucliadb_agentic.ask.search import rpc
+from hyperforge_nucliadb_agentic.ask.utils.responses import (
+    HTTPClientError,
+)
+from nucliadb_models.configuration import AskConfig
+from nucliadb_models.resource import NucliaDBRoles
+from nucliadb_models.search import (
+    NucliaDBClientType,
+)
+from nucliadb_models.security import RequestSecurity
+from nucliadb_utils.authentication import NucliaUser
+from pydantic import ValidationError
+
+from nucliadb_agentic_api.v1.router import router
+
+
+@router.websocket("/api/v1/kb/{kbid}/ask")
+@requires(NucliaDBRoles.READER)
+async def websocket_endpoint(
+    websocket: WebSocket,
+    kbid: str,
+    agentic_config_id: str = Query(
+        ..., description="ID of the agentic configuration to use for this session"
+    ),
+    search_configuration: str | None = Query(
+        default=None,
+        description="Optional search configuration to use for this session",
+    ),
+    groups: list[str] = Query(
+        default=[],
+        description="List of group ids to do the request with. Overrides any security group from the Authorization header.",
+    ),
+    keep_open: bool = Query(
+        default=False,
+        description="Whether to keep the websocket open after the first question",
+    ),
+    x_stf_user: str = Header(..., include_in_schema=False),
+    x_stf_account: str = Header(..., include_in_schema=False),
+    x_stf_account_type: str = Header(..., include_in_schema=False),
+    x_nucliadb_account: str = Header(default="", include_in_schema=False),
+    x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
+    x_show_consumption: bool = Header(default=False),
+    x_nucliadb_user: str = Header(""),
+    x_forwarded_for: str = Header(""),
+):
+    await websocket.accept()
+    receiver = WebsocketReceiver(websocket)
+    task = asyncio.create_task(receiver.run())
+
+    current_user: NucliaUser = websocket.user
+    # If present, security groups from AuthorizationBackend overrides any
+    # security group of the payload
+    if current_user.security_groups:
+        security = RequestSecurity(groups=current_user.security_groups)
+    else:
+        security = RequestSecurity(groups=groups)
+
+    item = AskRequest(query="")
+    item.security = security
+
+    if search_configuration is not None:
+        search_config = await rpc.get_search_configuration(
+            rpc.get_sdk("reader"), kbid, name=search_configuration
+        )
+        if search_config is None:
+            return HTTPClientError(
+                status_code=400, detail="Search configuration not found"
+            )
+
+        if not isinstance(search_config.config, AskConfig):
+            return HTTPClientError(
+                status_code=400,
+                detail="This search configuration is not valid for `ask`",
+            )
+
+        try:
+            item = AskRequest.model_validate(
+                search_config.config.model_dump(exclude_unset=True)
+                | item.model_dump(exclude_unset=True)
+            )
+        except ValidationError as e:
+            detail = json.loads(e.json())
+            return HTTPClientError(status_code=422, detail=detail)
+
+    # Wait for questions
+    first_question = True
+    while True:
+        if not keep_open and not first_question:
+            break
+        try:
+            interaction = await receiver.receive_question()
+        except WebSocketDisconnect:
+            break
+        except ValueError as e:
+            # Wrong message type received
+            await websocket.send_json(
+                AragAnswer(
+                    exception=ARAGException(detail=f"Unexpected message: {str(e)}"),
+                    operation=AnswerOperation.ERROR,
+                ).model_dump()
+            )
+            break
+
+        first_question = False
+        for header, header_value in websocket.headers.items():
+            interaction.headers[header] = header_value
+
+        item.query = interaction.question
+        interaction.arguments["ask_request"] = item.model_dump_json()
+
+        async for msg in stream_response(
+            websocket.app,
+            receiver,
+            x_stf_account,
+            kbid,
+            "ephemeral",
+            interaction,
+            workflow_id=agentic_config_id,
+        ):
+            try:
+                await websocket.send_text(msg.model_dump_json())
+            except (RuntimeError, WebSocketDisconnect):
+                # WebSocket already closed
+                pass
+
+    try:
+        task.cancel()
+        await websocket.close()
+    except RuntimeError:
+        # WebSocket already closed
+        pass