10xgraph-api 0.7.0__py3-none-any.whl

tenxgraph_api/src/app/routers/ag_ui/service.py

@@ -0,0 +1,193 @@
+"""Run the graph for one AG-UI request and stream the result as AG-UI events."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from typing import Any
+
+from ag_ui.core import RunAgentInput
+from ag_ui.encoder import EventEncoder
+from tenxgraph.utils import ResponseGranularity
+
+from tenxgraph_api.src.app.core import logger
+from tenxgraph_api.src.app.core.auth.request_config import tool_result_ids
+from tenxgraph_api.src.app.routers.graph.schemas.graph_schemas import GraphInputSchema
+from tenxgraph_api.src.app.routers.graph.services.graph_service import GraphService
+
+from .converter import select_new_messages, to_tenxgraph_messages
+from .event_mapper import AgUiEventMapper
+
+
+# Graph config key holding the AG-UI request extras (frontend context, tools, forwarded props),
+# so graph nodes and tools can read them from ``config["ag_ui"]``.
+AG_UI_CONFIG_KEY = "ag_ui"
+
+# Client state keys that would overwrite the conversation or execution metadata.
+_RESERVED_STATE_KEYS = frozenset({"context", "context_summary", "execution_meta"})
+
+
+async def stream_ag_ui(
+    service: GraphService,
+    run_input: RunAgentInput,
+    user: dict[str, Any],
+    encoder: EventEncoder,
+    allow_client_tools: bool = True,
+) -> AsyncIterator[str]:
+    """Stream one AG-UI run as encoded events.
+
+    With ``allow_client_tools`` the tools the client sends are offered to the model for this run
+    (see ``ag_ui.allow_client_tools``); without it only ``remote_tools`` from 10xgraph.json are.
+
+    The thread's checkpoint is the record of the conversation: only client messages it has not
+    seen are sent to the graph. A run with nothing new (a client syncing state) still gets a
+    complete ``RUN_STARTED`` / ``RUN_FINISHED`` pair, with the thread's state.
+    """
+    initial_state = _client_state(run_input.state)
+    mapper = AgUiEventMapper(run_input.thread_id, run_input.run_id, initial_state=initial_state)
+    for event in mapper.start():
+        yield encoder.encode(event)
+
+    try:
+        async for event in _run_events(
+            service, run_input, user, mapper, initial_state, allow_client_tools
+        ):
+            yield encoder.encode(event)
+    except ValueError as exc:
+        logger.warning("AG-UI run rejected: %s", exc)
+        for event in mapper.fail(str(exc), "INVALID_INPUT"):
+            yield encoder.encode(event)
+    except Exception as exc:
+        # The response has already started, so the failure has to travel in the stream.
+        logger.error("AG-UI run failed: %s", exc)
+        for event in mapper.fail("AG-UI run failed", "AG_UI_ERROR"):
+            yield encoder.encode(event)
+
+    for event in mapper.finish():
+        yield encoder.encode(event)
+
+
+async def _run_events(
+    service: GraphService,
+    run_input: RunAgentInput,
+    user: dict[str, Any],
+    mapper: AgUiEventMapper,
+    initial_state: dict[str, Any] | None,
+    allow_client_tools: bool,
+) -> AsyncIterator[Any]:
+    """The events between ``RUN_STARTED`` and the terminal event of one run."""
+    checkpoint_state = await _checkpoint_state(service, run_input.thread_id, user)
+    context = list(getattr(checkpoint_state, "context", None) or [])
+    new_messages = select_new_messages(
+        run_input.messages,
+        known_ids={str(m.message_id) for m in context},
+        answered_call_ids=set(tool_result_ids(context)),
+    )
+    messages = to_tenxgraph_messages(new_messages)
+    paused_at = _pending_interrupt(checkpoint_state)
+    resume = _resume_value(paused_at, run_input.resume or [])
+
+    if paused_at is not None and resume is _NO_RESUME:
+        # A thread paused at interrupt() only moves on with an answer; report the pause again
+        # instead of running the graph.
+        mapper.report_interrupt(paused_at.model_dump(mode="json"))
+        return
+    if not messages and resume is _NO_RESUME:
+        if checkpoint_state is not None:
+            for event in mapper.map_state(checkpoint_state):
+                yield event
+        return
+
+    graph_input, server_config = _graph_request(run_input, messages, initial_state, resume)
+    if not allow_client_tools:
+        server_config = {}
+    async for chunk in service.stream_chunks(graph_input, user, server_config):
+        for event in mapper.map(chunk):
+            yield event
+
+
+def _client_state(state: Any) -> dict[str, Any] | None:
+    if not isinstance(state, dict):
+        return None
+    cleaned = {k: v for k, v in state.items() if k not in _RESERVED_STATE_KEYS}
+    return cleaned or None
+
+
+async def _checkpoint_state(service: GraphService, thread_id: str, user: dict[str, Any]) -> Any:
+    if service.checkpointer is None:
+        return None
+    config = {
+        "thread_id": thread_id,
+        "user": user,
+        "user_id": user.get("user_id", "anonymous"),
+    }
+    return await service.checkpointer.aget_state(config)
+
+
+# Marks "the client sent no answer for the open interrupt" (a None answer means cancelled).
+_NO_RESUME: Any = object()
+
+# ``remote_tools`` run-config key read by the tenxgraph ToolNode for per-run client tools.
+RUN_REMOTE_TOOLS_KEY = "remote_tools"
+
+
+def _pending_interrupt(state: Any) -> Any:
+    """The interrupt() the thread is paused at, or ``None`` (and on cores without interrupts)."""
+    if state is None:
+        return None
+    try:
+        from tenxgraph.utils.interrupt import pending_interrupt
+    except ImportError:  # tenxgraph core older than interrupt()
+        return None
+    return pending_interrupt(state)
+
+
+def _resume_value(paused_at: Any, entries: list[Any]) -> Any:
+    """The graph resume value for the AG-UI ``resume`` entries.
+
+    Returns ``_NO_RESUME`` when nothing answers the open interrupt. A ``cancelled`` entry resumes
+    with ``None``.
+
+    Raises:
+        ValueError: ``resume`` was sent but the thread is not paused at that interrupt.
+    """
+    if not entries:
+        return _NO_RESUME
+    if paused_at is None:
+        raise ValueError("Nothing to resume: this thread is not paused at an interrupt")
+    for entry in entries:
+        if entry.interrupt_id == paused_at.id:
+            return entry.payload if entry.status == "resolved" else None
+    raise ValueError(f"resume does not answer the open interrupt '{paused_at.id}'")
+
+
+def _graph_request(
+    run_input: RunAgentInput,
+    messages: list[Any],
+    initial_state: dict[str, Any] | None,
+    resume: Any,
+) -> tuple[GraphInputSchema, dict[str, Any]]:
+    """The graph input for one AG-UI run, and the server-owned run config that goes with it."""
+    extra: dict[str, Any] = {} if resume is _NO_RESUME else {"resume": resume}
+    graph_input = GraphInputSchema(
+        messages=messages,
+        initial_state=initial_state,
+        config={
+            "thread_id": run_input.thread_id,
+            "run_id": run_input.run_id,
+            AG_UI_CONFIG_KEY: {
+                "context": [c.model_dump(mode="json") for c in run_input.context or []],
+                "tools": [t.model_dump(mode="json") for t in run_input.tools or []],
+                "forwarded_props": run_input.forwarded_props,
+            },
+        },
+        response_granularity=ResponseGranularity.FULL,
+        **extra,
+    )
+    # The client's own tools, offered to the model for this run only.
+    server_config = {
+        RUN_REMOTE_TOOLS_KEY: [
+            {"name": t.name, "description": t.description, "parameters": t.parameters}
+            for t in run_input.tools or []
+        ]
+    }
+    return graph_input, server_config

tenxgraph_api/src/app/routers/checkpointer/__init__.py

@@ -0,0 +1 @@
+# Checkpointer module

tenxgraph_api/src/app/routers/checkpointer/router.py

@@ -0,0 +1,482 @@
+"""Checkpointer router module."""
+
+from typing import Any
+
+from fastapi import APIRouter, Depends, HTTPException, Request, status
+from injectq.integrations import InjectAPI
+from tenxgraph.core.state import Message
+
+from tenxgraph_api.src.app.core.auth.permissions import RequirePermission
+from tenxgraph_api.src.app.core.auth.request_config import client_config
+from tenxgraph_api.src.app.utils.response_helper import success_response
+from tenxgraph_api.src.app.utils.swagger_helper import generate_swagger_responses
+
+from .schemas.checkpointer_schemas import (
+    ConfigSchema,
+    MessagesListResponseSchema,
+    PutMessagesSchema,
+    ResponseSchema,
+    StateResponseSchema,
+    StateSchema,
+    ThreadResponseSchema,
+    ThreadsListResponseSchema,
+)
+from .services.checkpointer_service import CheckpointerService
+
+
+router = APIRouter(tags=["checkpointer"])
+
+# Bound pagination so an unbounded/huge ``limit`` (or ``limit=None``) can never ask the
+# backend to load a whole table into memory. Applied server-side regardless of client input.
+DEFAULT_PAGE_LIMIT = 100
+MAX_PAGE_LIMIT = 1000
+
+
+def _clamp_limit(limit: int | None) -> int:
+    """Return a bounded page size: default when unset, capped at MAX_PAGE_LIMIT."""
+    if limit is None:
+        return DEFAULT_PAGE_LIMIT
+    return min(limit, MAX_PAGE_LIMIT)
+
+
+def validate_thread_id(thread_id: int | str) -> None:
+    if isinstance(thread_id, str):
+        if not thread_id.strip():
+            raise HTTPException(status_code=422, detail="thread_id cannot be empty or whitespace")
+    elif isinstance(thread_id, int):
+        if thread_id < 1:
+            raise HTTPException(status_code=422, detail="thread_id must be a non-negative integer")
+    else:
+        raise HTTPException(status_code=422, detail="thread_id must be a string or integer")
+
+
+@router.get(
+    "/v1/threads/{thread_id}/state",
+    status_code=status.HTTP_200_OK,
+    responses=generate_swagger_responses(StateResponseSchema),
+    summary="Get state from checkpointer",
+    description="Retrieve state data from the checkpointer using configuration.",
+)
+async def get_state(
+    request: Request,
+    thread_id: int | str,
+    service: CheckpointerService = InjectAPI(CheckpointerService),
+    user: dict[str, Any] = Depends(RequirePermission("checkpointer", "read")),
+):
+    """Get state from checkpointer.
+
+    Args:
+        request: State schema with configuration
+        checkpointer: Injected checkpointer instance
+
+    Returns:
+        State response with state data or error
+    """
+    validate_thread_id(thread_id)
+
+    config = {"thread_id": thread_id}
+
+    result = await service.get_state(
+        config,
+        user,
+    )
+
+    return success_response(
+        result,
+        request,
+    )
+
+
+@router.put(
+    "/v1/threads/{thread_id}/state",
+    status_code=status.HTTP_200_OK,
+    responses=generate_swagger_responses(StateResponseSchema),
+    summary="Put state to checkpointer",
+    description="Store state data in the checkpointer using configuration.",
+)
+async def put_state(
+    request: Request,
+    thread_id: str | int,
+    payload: StateSchema,
+    service: CheckpointerService = InjectAPI(CheckpointerService),
+    user: dict[str, Any] = Depends(RequirePermission("checkpointer", "write")),
+):
+    """Put state to checkpointer.
+
+    Args:
+        request: Request object
+        payload: Put state schema with configuration and state data
+        service: Injected checkpointer service
+        user: Current authenticated user
+
+    Returns:
+        Success response or error
+    """
+    validate_thread_id(thread_id)
+    # The path thread_id is the one RequirePermission checked; it is applied last so the
+    # body config cannot redirect the call to another thread.
+    config = {**client_config(payload.config), "thread_id": thread_id}
+
+    res = await service.put_state(
+        config,
+        user,
+        payload.state,
+    )
+
+    return success_response(
+        res,
+        request,
+    )
+
+
+@router.delete(
+    "/v1/threads/{thread_id}/state",
+    status_code=status.HTTP_200_OK,
+    responses=generate_swagger_responses(ResponseSchema),
+    summary="Clear state from checkpointer",
+    description="Clear state data from the checkpointer using configuration.",
+)
+async def clear_state(
+    request: Request,
+    thread_id: int | str,
+    service: CheckpointerService = InjectAPI(CheckpointerService),
+    user: dict[str, Any] = Depends(RequirePermission("checkpointer", "delete")),
+):
+    """Clear state from checkpointer.
+
+    Args:
+        request: Request object
+        payload: Clear state schema with configuration
+        service: Injected checkpointer service
+        user: Current authenticated user
+
+    Returns:
+        Success response or error
+    """
+    validate_thread_id(thread_id)
+    config = {"thread_id": thread_id}
+
+    res = await service.clear_state(
+        config,
+        user,
+    )
+
+    return success_response(
+        res,
+        request,
+    )
+
+
+# Now Handle Messages
+
+
+@router.post(
+    "/v1/threads/{thread_id}/messages",
+    status_code=status.HTTP_200_OK,
+    responses=generate_swagger_responses(ResponseSchema),
+    summary="Put messages to checkpointer",
+    description="Store messages in the checkpointer using configuration.",
+)
+async def put_messages(
+    request: Request,
+    thread_id: str | int,
+    payload: PutMessagesSchema,
+    service: CheckpointerService = InjectAPI(CheckpointerService),
+    user: dict[str, Any] = Depends(RequirePermission("checkpointer", "write")),
+):
+    """Put messages to checkpointer.
+
+    Args:
+        request: Request object
+        payload: Put messages schema with configuration and messages
+        service: Injected checkpointer service
+        user: Current authenticated user
+
+    Returns:
+        Success response or error
+    """
+    validate_thread_id(thread_id)
+    if not payload.messages:
+        raise HTTPException(status_code=422, detail="messages must not be empty")
+
+    # The path thread_id is the one RequirePermission checked; it is applied last so the
+    # body config cannot redirect the call to another thread.
+    config = {**client_config(payload.config), "thread_id": thread_id}
+
+    res = await service.put_messages(
+        config,
+        user,
+        payload.messages,
+        payload.metadata,
+    )
+
+    return success_response(
+        res,
+        request,
+    )
+
+
+@router.get(
+    "/v1/threads/{thread_id}/messages/{message_id}",
+    status_code=status.HTTP_200_OK,
+    responses=generate_swagger_responses(Message),
+    summary="Get message from checkpointer",
+    description=(
+        "Retrieve a specific message from the checkpointer using configuration and message ID."
+    ),
+)
+async def get_message(
+    request: Request,
+    thread_id: str | int,
+    message_id: str | int,
+    service: CheckpointerService = InjectAPI(CheckpointerService),
+    user: dict[str, Any] = Depends(RequirePermission("checkpointer", "read")),
+):
+    """Get message from checkpointer.
+
+    Args:
+        request: Request object
+        payload: Get message schema with configuration and message ID
+        service: Injected checkpointer service
+        user: Current authenticated user
+
+    Returns:
+        Message response with message data or error
+    """
+    validate_thread_id(thread_id)
+    if not message_id or (isinstance(message_id, str) and not str(message_id).strip()):
+        raise HTTPException(status_code=422, detail="message_id is required and cannot be empty")
+
+    config = {"thread_id": thread_id}
+
+    result = await service.get_message(
+        config,
+        user,
+        message_id,
+    )
+
+    return success_response(
+        result,
+        request,
+    )
+
+
+@router.get(
+    "/v1/threads/{thread_id}/messages",
+    status_code=status.HTTP_200_OK,
+    responses=generate_swagger_responses(MessagesListResponseSchema),
+    summary="List messages from checkpointer",
+    description="Retrieve a list of messages from the checkpointer using configuration "
+    "and optional filters.",
+)
+async def list_messages(
+    request: Request,
+    thread_id: int | str,
+    search: str | None = None,
+    offset: int | None = None,
+    limit: int | None = None,
+    service: CheckpointerService = InjectAPI(CheckpointerService),
+    user: dict[str, Any] = Depends(RequirePermission("checkpointer", "read")),
+):
+    """List messages from checkpointer.
+
+    Args:
+        request: Request object
+        payload: List messages schema with configuration and optional filters
+        service: Injected checkpointer service
+        user: Current authenticated user
+
+    Returns:
+        Messages list response with messages data or error
+    """
+    validate_thread_id(thread_id)
+    if offset is not None and offset < 0:
+        raise HTTPException(status_code=422, detail="offset must be >= 0")
+    if limit is not None and limit <= 0:
+        raise HTTPException(status_code=422, detail="limit must be > 0")
+
+    config = {"thread_id": thread_id}
+
+    result = await service.get_messages(
+        config,
+        user,
+        search,
+        offset,
+        _clamp_limit(limit),
+    )
+
+    return success_response(
+        result,
+        request,
+    )
+
+
+@router.delete(
+    "/v1/threads/{thread_id}/messages/{message_id}",
+    status_code=status.HTTP_200_OK,
+    responses=generate_swagger_responses(ResponseSchema),
+    summary="Delete message from checkpointer",
+    description="Delete a specific message from the checkpointer using configuration and ID.",
+)
+async def delete_message(
+    request: Request,
+    message_id: str | int,
+    thread_id: str | int,
+    payload: ConfigSchema,
+    service: CheckpointerService = InjectAPI(CheckpointerService),
+    user: dict[str, Any] = Depends(RequirePermission("checkpointer", "delete")),
+):
+    """Delete message from checkpointer.
+
+    Args:
+        request: Request object
+        payload: Delete message schema with configuration and message ID
+        service: Injected checkpointer service
+        user: Current authenticated user
+
+    Returns:
+        Success response or error
+    """
+    validate_thread_id(thread_id)
+    if not message_id or (isinstance(message_id, str) and not str(message_id).strip()):
+        raise HTTPException(status_code=422, detail="message_id is required and cannot be empty")
+
+    # The path thread_id is the one RequirePermission checked; it is applied last so the
+    # body config cannot redirect the call to another thread.
+    config = {**client_config(payload.config), "thread_id": thread_id}
+
+    await service.delete_message(
+        config,
+        user,
+        message_id,
+    )
+
+    return success_response(
+        {"success": True, "message": "Message deleted successfully"},
+        request,
+    )
+
+
+# Handle Threads
+
+
+@router.get(
+    "/v1/threads/{thread_id}",
+    status_code=status.HTTP_200_OK,
+    responses=generate_swagger_responses(ThreadResponseSchema),
+    summary="Get thread from checkpointer",
+    description="Retrieve a specific thread from the checkpointer using configuration.",
+)
+async def get_thread(
+    request: Request,
+    thread_id: str | int,
+    service: CheckpointerService = InjectAPI(CheckpointerService),
+    user: dict[str, Any] = Depends(RequirePermission("checkpointer", "read")),
+):
+    """Get thread from checkpointer.
+
+    Args:
+        request: Request object
+        payload: Get thread schema with configuration and thread ID
+        service: Injected checkpointer service
+        user: Current authenticated user
+
+    Returns:
+        Thread response with thread data or error
+    """
+    validate_thread_id(thread_id)
+    result = await service.get_thread(
+        {"thread_id": thread_id},
+        user,
+    )
+
+    return success_response(
+        {"thread_data": result},
+        request,
+    )
+
+
+@router.get(
+    "/v1/threads",
+    status_code=status.HTTP_200_OK,
+    responses=generate_swagger_responses(ThreadsListResponseSchema),
+    summary="List threads from checkpointer",
+    description="Retrieve a list of threads from the checkpointer with optional filters.",
+)
+async def list_threads(
+    request: Request,
+    search: str | None = None,
+    offset: int | None = None,
+    limit: int | None = None,
+    service: CheckpointerService = InjectAPI(CheckpointerService),
+    user: dict[str, Any] = Depends(RequirePermission("checkpointer", "read")),
+):
+    """List threads from checkpointer.
+
+    Args:
+        request: Request object
+        payload: List threads schema with configuration and optional filters
+        service: Injected checkpointer service
+        user: Current authenticated user
+
+    Returns:
+        Threads list response with threads data or error
+    """
+    if limit is not None and limit <= 0:
+        raise HTTPException(status_code=422, detail="limit must be > 0")
+    if offset is not None and offset < 0:
+        raise HTTPException(status_code=422, detail="offset must be >= 0")
+
+    result = await service.list_threads(
+        user,
+        search,
+        offset,
+        _clamp_limit(limit),
+    )
+
+    return success_response(
+        result,
+        request,
+    )
+
+
+@router.delete(
+    "/v1/threads/{thread_id}",
+    status_code=status.HTTP_200_OK,
+    responses=generate_swagger_responses(ResponseSchema),
+    summary="Delete thread from checkpointer",
+    description="Delete a specific thread from the checkpointer using configuration and thread ID.",
+)
+async def delete_thread(
+    request: Request,
+    thread_id: str | int,
+    payload: ConfigSchema,
+    service: CheckpointerService = InjectAPI(CheckpointerService),
+    user: dict[str, Any] = Depends(RequirePermission("checkpointer", "delete")),
+):
+    """Delete thread from checkpointer.
+
+    Args:
+        request: Request object
+        payload: Delete thread schema with configuration and thread ID
+        service: Injected checkpointer service
+        user: Current authenticated user
+
+    Returns:
+        Success response or error
+    """
+    validate_thread_id(thread_id)
+    # The path thread_id is the one RequirePermission checked; it is applied last so the
+    # body config cannot redirect the call to another thread.
+    config = {**client_config(payload.config), "thread_id": thread_id}
+
+    res = await service.delete_thread(
+        config,
+        user,
+        thread_id,
+    )
+
+    return success_response(
+        res,
+        request,
+    )

tenxgraph_api/src/app/routers/checkpointer/schemas/__init__.py

@@ -0,0 +1 @@
+# Checkpointer schemas