diff --git a/src/agentscope/agent/_config.py b/src/agentscope/agent/_config.py new file mode 100644 index 0000000000000000000000000000000000000000..82a4a3949a48e99c0b1d6d9a5344cfaac38944dc --- /dev/null +++ b/src/agentscope/agent/_config.py @@ -0,0 +1,172 @@ +# -*- coding: utf-8 -*- +"""The agent config classes.""" + +from pydantic import BaseModel, Field + +from ..model import ChatModelBase + + +class SummarySchema(BaseModel): + """The compressed memory model, used to generate summary of old memories""" + + task_overview: str = Field( + description=( + "The user's core request and success criteria.\n" + "Any clarifications or constraints they specified" + ), + ) + current_state: str = Field( + description=( + "What has been completed so far.\n" + "File created, modified, or analyzed (with paths if relevant).\n" + "Key outputs or artifacts produced." + ), + ) + important_discoveries: str = Field( + description=( + "Technical constraints or requirements uncovered.\n" + "Decisions made and their rationale.\n" + "Errors encountered and how they were resolved.\n" + "What approaches were tried that didn't work (and why)" + ), + ) + next_steps: str = Field( + description=( + "Specific actions needed to complete the task.\n" + "Any blockers or open questions to resolve.\n" + "Priority order if multiple steps remain" + ), + ) + context_to_preserve: str = Field( + description=( + "User preferences or style requirements.\n" + "Domain-specific details that aren't obvious.\n" + "Any promises made to the user" + ), + ) + """Whether to execute multiple tool calls in parallel within one + reasoning step.""" + + +class ContextConfig(BaseModel): + """The context related configuration in AgentScope""" + + model_config = {"arbitrary_types_allowed": True} + """Allow arbitrary types in the pydantic model.""" + + trigger_ratio: float = Field(default=0.8, gt=0, lt=0.9) + """When the token exceeds this ratio of the maximum context length, the + context will be compressed. To reserve the context for context compression, + the maximum ratio is 0.9.""" + + reserve_ratio: float = Field(default=0.1, gt=0, lt=0.9) + """The ratio of the tokens to reserve in context compression, which should + be smaller than the trigger ratio.""" + + compression_prompt: str = Field( + default=( + "You have been working on the task described above " + "but have not yet completed it. " + "Now write a continuation summary that will allow you to resume " + "work efficiently in a future context window where the " + "conversation history will be replaced with this summary. " + "Your summary should be structured, concise, and actionable." + "" + ), + # ``format: textarea`` is a hint for schema-driven UI renderers + # to use a multi-line input. Plain JSON Schema doesn't natively + # express this, so we piggy-back on ``json_schema_extra``. + json_schema_extra={"format": "textarea"}, + ) + """The prompt used to guide the compression model to generate the + compressed summary, which will be wrapped into a user message and + attach to the end of the current memory.""" + + summary_template: str = Field( + default=( + "Here is a summary of your previous work\n" + "# Task Overview\n" + "{task_overview}\n\n" + "# Current State\n" + "{current_state}\n\n" + "# Important Discoveries\n" + "{important_discoveries}\n\n" + "# Next Steps\n" + "{next_steps}\n\n" + "# Context to Preserve\n" + "{context_to_preserve}" + "" + ), + json_schema_extra={"format": "textarea"}, + ) + """The string template to present the compressed summary to the agent, + which will be formatted with the fields from the + `compression_summary_model`.""" + + summary_schema: dict = Field( + default_factory=SummarySchema.model_json_schema, + ) + """The structured model used to guide the agent to generate the + structured compressed summary.""" + + tool_result_limit: int = Field( + title="Tool Result Limit", + default=50000, + description=( + "The maximum length of the tool results in tokens. " + "If exceeded, the tool result will be truncated." + ), + ) + """The tool result limit to avoid tool result bursting.""" + + +class ReActConfig(BaseModel): + """The reasoning related configuration""" + + max_iters: int = Field( + title="Max Iterations", + default=20, + description="The maximum number of reasoning-acting iterations in " + "one reply", + ) + """The maximum number of iterations for the reasoning-acting loop.""" + + stop_on_reject: bool = Field( + title="Rejection Handling", + default=False, + description="Whether to stop replying when being rejected to " + "execute tools.", + ) + """If stop reasoning when tool call(s) are rejected. If `True`, the agent + won't continue reasoning and wait for outside interaction from the user. + """ + + +class ModelConfig(BaseModel): + """The model related configuration.""" + + # TODO: remove this line after PR #1564 is merged, where the ChatModel + # will be child class of BaseModel + model_config = {"arbitrary_types_allowed": True} + + max_retries: int = Field( + default=0, + ge=0, + description=( + "Number of retries on top of the initial call before falling " + "over to the fallback model. ``0`` means call the model exactly " + "once and immediately move to the fallback on failure. Same " + "semantics as ``ChatModelBase.max_retries``. Defaults to 0 to " + "avoid compounding with the model's own inner retry loop." + ), + ) + """Number of retries on top of the initial call before falling over to + the fallback model. ``0`` means a single attempt with no retries. + Mirrors the semantics of ``ChatModelBase.max_retries``.""" + + fallback_model: ChatModelBase | None = Field( + default=None, + description="The fallback model used when the main model fails.", + ) + """The fallback model used when the main model fails. Also supports the + max_retries logic.""" diff --git a/src/agentscope/agent/_utils.py b/src/agentscope/agent/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f2449e442a6dae04f96f49f34b1ff7979e95997f --- /dev/null +++ b/src/agentscope/agent/_utils.py @@ -0,0 +1,17 @@ +# -*- coding: utf-8 -*- +"""The utility classes used in building the agent class.""" +from dataclasses import dataclass +from typing import Literal + +from ..message import ToolCallBlock + + +@dataclass +class _ToolCallBatch: + """A batch of tool calls that execute either sequentially or + concurrently.""" + + type: Literal["sequential", "concurrent"] + """The batch type""" + tool_calls: list[ToolCallBlock] + """The list of tool calls in the batch.""" diff --git a/src/agentscope/app/__init__.py b/src/agentscope/app/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..99e435d3cf7d49411630069074e793530c01d9f4 --- /dev/null +++ b/src/agentscope/app/__init__.py @@ -0,0 +1,12 @@ +# -*- coding: utf-8 -*- +"""The FastAPI based agent service module, which contains all service-related +components and a configurable FastAPI app factory. +""" + +from ._app import create_app +from ._types import SubAgentTemplate + +__all__ = [ + "create_app", + "SubAgentTemplate", +] diff --git a/src/agentscope/app/_app.py b/src/agentscope/app/_app.py new file mode 100644 index 0000000000000000000000000000000000000000..a3a70bc35effa5204dbfc78910b8f28361980960 --- /dev/null +++ b/src/agentscope/app/_app.py @@ -0,0 +1,263 @@ +# -*- coding: utf-8 -*- +"""AgentScope app factory.""" +from typing import Type, TYPE_CHECKING, Any + +from ._lifespan import lifespan +from .rag.blob_store import BlobStoreBase, LocalBlobStore +from .rag.knowledge_base_manager import KnowledgeBaseManagerBase +from .workspace_manager import WorkspaceManagerBase +from ._router import ( + agent_router, + chat_router, + credential_router, + knowledge_base_router, + model_router, + tts_model_router, + schedule_router, + session_router, + workspace_router, +) +from ._types import AgentMiddlewareFactory, AgentToolFactory, SubAgentTemplate +from .message_bus import MessageBus +from .storage import StorageBase +from ..agent import Agent +from ..credential import CredentialFactory, CredentialBase +from ..rag import ( + ApproxTokenChunker, + ChunkerBase, + ParserBase, + TextParser, +) +from .._version import __version__ + + +if TYPE_CHECKING: + from fastapi import FastAPI + from fastapi.middleware import Middleware as FastAPIMiddleware +else: + FastAPI = Any + FastAPIMiddleware = Any + + +def create_app( + storage: StorageBase, + message_bus: MessageBus, + workspace_manager: WorkspaceManagerBase, + knowledge_base_manager: KnowledgeBaseManagerBase | None = None, + knowledge_parsers: list[ParserBase] | dict[str, ParserBase] | None = None, + knowledge_chunker: ChunkerBase | None = None, + blob_store: BlobStoreBase | None = None, + enable_index_worker: bool = True, + *, + extra_credentials: list[Type[CredentialBase]] | None = None, + extra_middlewares: list[FastAPIMiddleware] | None = None, + extra_agent_middlewares: AgentMiddlewareFactory | None = None, + extra_agent_tools: AgentToolFactory | None = None, + custom_subagent_templates: list[SubAgentTemplate] | None = None, + custom_agent_cls: Type[Agent] | None = None, + title: str = "AgentScope", + version: str = __version__, +) -> FastAPI: + """Create and configure a FastAPI application. + + This is the primary entry point for embedding AgentScope into an existing + service or running it standalone. All built-in routers are registered + automatically; pass ``extra_middlewares`` to add your own. + + Usage — standalone:: + + app = create_app( + storage=RedisStorage(), + message_bus=RedisMessageBus(), + workspace_manager=LocalWorkspaceManager(), + ) + uvicorn.run(app, host="0.0.0.0", port=8000) + + Usage — mount onto an existing app:: + + root = FastAPI() + agentscope_app = create_app( + storage=RedisStorage(), + message_bus=RedisMessageBus(), + workspace_manager=LocalWorkspaceManager(), + ) + root.mount("/agentscope", agentscope_app) + + Args: + storage (`StorageBase`): + The storage backend. Its lifecycle (``__aenter__`` / + ``__aexit__``) is managed by the app lifespan. + message_bus (`MessageBus`): + The live message bus used for cross-session inbox delivery + and idle-session triggers. Required — the bus is intentionally + decoupled from ``storage`` so the persistence backend (e.g. + SQL) can differ from the transport backend (Redis). Its + lifecycle is also managed by the app lifespan. + workspace_manager (`WorkspaceManagerBase`): + The workspace manager. Required — every chat run and every + ``/workspace`` endpoint depends on it. Its lifecycle ( + ``__aenter__`` / ``__aexit__``) is managed by the app + lifespan. Pass a :class:`~agentscope.app._manager. + LocalWorkspaceManager` for local-directory workspaces. + knowledge_base_manager (`KnowledgeBaseManagerBase | None`, \ + optional): + The knowledge base manager that owns knowledge base + lifecycle and serves + :class:`~agentscope.rag.KnowledgeBase` + runtime handles to both HTTP service and agent code. + The manager carries its own vector store instance — its + ``__aenter__`` / ``__aexit__`` enter and release that + vector store, so the caller does not pass the vector + store separately. ``None`` disables knowledge base + endpoints entirely. + knowledge_parsers (`list[ParserBase] | dict[str, ParserBase] | \ + None`, optional): + Parsers registered for knowledge base document uploads. + Pass a **list** to have the service route by each parser's + ``supported_media_types`` (later entries override earlier + ones for overlapping types, with a warning); pass a + **dict** ``media_type → parser`` for explicit routing + (one parser bound to multiple types, type aliases, ...). + Defaults to ``[TextParser()]`` when + ``knowledge_base_manager`` is set. + knowledge_chunker (`ChunkerBase | None`, optional): + The chunker shared across every knowledge base. Defaults + to :class:`~agentscope.rag.ApproxTokenChunker()` when + ``knowledge_base_manager`` is set. + blob_store (`BlobStoreBase | None`, optional): + Backend storing uploaded document bytes between the + upload endpoint and the indexing worker. Required when + ``knowledge_base_manager`` is set; defaults to + :class:`~agentscope.app.rag.blob_store.LocalBlobStore` + rooted at ``./blobs``. Its lifecycle (``__aenter__`` / + ``__aexit__``) is managed by the app lifespan. + enable_index_worker (`bool`, defaults to ``True``): + When ``True`` (embedded deployment) the API process starts + an :class:`~agentscope.app._service.IndexWorker` and an + :class:`~agentscope.app._service.IndexSweeper` in its + lifespan, and dispatches indexing tasks via an + in-process queue. When ``False`` (dedicated deployment) + the API process performs no indexing — a separate worker + process is expected to consume tasks from the message + bus. No effect when ``knowledge_base_manager`` is + ``None``. + extra_credentials (`list[Type[CredentialBase]] | None`, optional): + Additional :class:`~agentscope.credential.CredentialBase` + subclasses to register before the app starts. Equivalent to + calling :func:`~agentscope.credential.CredentialFactory. + register_credential` for each class. + extra_middlewares (`list[Middleware] | None`, optional): + Additional ASGI middlewares to add to the application. + extra_agent_middlewares (`AgentMiddlewareFactory | None`, optional): + An async factory ``(user_id, agent_id, session_id) -> awaitable + of list[MiddlewareBase]`` that produces extra + :class:`~agentscope.middleware.MiddlewareBase` instances to + attach to the agent on each invocation. Called once per agent + assembly (i.e. per chat turn / scheduled trigger), so it can + return user/session-specific middleware (auth, audit logging, + tenant isolation, etc.). The returned middlewares are appended + to the framework-supplied ones (e.g. ``ToolOffloadMiddleware``). + extra_agent_tools (`AgentToolFactory | None`, optional): + An async factory ``(user_id, agent_id, session_id) -> awaitable + of list[ToolBase]`` that produces extra + :class:`~agentscope.tool.ToolBase` instances to register in the + agent's toolkit on each invocation. Useful when tool + availability depends on the caller (per-tenant integrations, + user-specific credentials). The returned tools are added to + the workspace-derived tools in the toolkit's ``"basic"`` group. + custom_subagent_templates (`list[SubAgentTemplate] | None`, optional): + Reusable blueprints for sub-agent creation within teams. + Each template defines a sub-agent *type* (e.g. ``"researcher"``, + ``"coder"``) with pre-configured system prompt, context config, + ReAct config, permission context, and task context. When + registered, the ``AgentCreate`` tool exposes a + ``subagent_type`` parameter so the leader agent can route to + the appropriate template. See + :class:`~agentscope.app._types.SubAgentTemplate` for details. + custom_agent_cls (`Type[Agent] | None`, optional): + A custom :class:`~agentscope.agent.Agent` subclass to use + when assembling agents. When ``None`` (default), the + built-in :class:`~agentscope.agent.Agent` is used. + title (`str`, defaults to ``"AgentScope"``): + OpenAPI title shown in the docs UI. + version (`str`, defaults to the package version): + API version shown in the docs UI. + + Returns: + `FastAPI`: A fully configured application ready to serve requests. + """ + from fastapi import FastAPI + + # Register any user-supplied credential types before the app starts + for cls in extra_credentials or []: + CredentialFactory.register_credential(cls) + + app = FastAPI(title=title, version=version, lifespan=lifespan) + + # Attach shared state that lifespan and dependencies read from app.state + app.state.storage = storage + app.state.message_bus = message_bus + app.state.workspace_manager = workspace_manager + app.state.knowledge_base_manager = knowledge_base_manager + app.state.extra_agent_middlewares = extra_agent_middlewares + app.state.extra_agent_tools = extra_agent_tools + app.state.custom_agent_cls = custom_agent_cls + + # Parser / chunker / blob-store defaults only make sense when the + # KB feature is actually enabled. When ``knowledge_base_manager`` is + # ``None`` every KB endpoint is disabled, so leaving these as ``None`` + # avoids unused imports being eagerly constructed at app startup. + if knowledge_base_manager is not None: + app.state.knowledge_parsers = ( + knowledge_parsers + if knowledge_parsers is not None + else [TextParser()] + ) + app.state.knowledge_chunker = knowledge_chunker or ApproxTokenChunker() + app.state.blob_store = ( + blob_store + if blob_store is not None + else LocalBlobStore(root_dir="./blobs") + ) + else: + app.state.knowledge_parsers = knowledge_parsers + app.state.knowledge_chunker = knowledge_chunker + app.state.blob_store = blob_store + app.state.enable_index_worker = ( + enable_index_worker and knowledge_base_manager is not None + ) + + # Validate custom sub-agent templates for duplicate types and store in + # app.state + templates = custom_subagent_templates or [] + seen_types: set[str] = set() + duplicates: set[str] = set() + for t in templates: + if t.type in seen_types: + duplicates.add(t.type) + seen_types.add(t.type) + if duplicates: + raise ValueError( + f"Duplicate sub_agent_template type(s): {duplicates}", + ) + app.state.custom_subagent_templates = {t.type: t for t in templates} + + # Built-in routers + for router in ( + agent_router, + chat_router, + credential_router, + knowledge_base_router, + schedule_router, + session_router, + workspace_router, + model_router, + tts_model_router, + ): + app.include_router(router) + + # Optional extra middlewares + for middleware in extra_middlewares or []: + app.add_middleware(middleware.cls, **middleware.kwargs) + + return app diff --git a/src/agentscope/app/_bus_ops.py b/src/agentscope/app/_bus_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..d360a2717096fd17f8dbc6757831a90bae630712 --- /dev/null +++ b/src/agentscope/app/_bus_ops.py @@ -0,0 +1,167 @@ +# -*- coding: utf-8 -*- +"""Business-level operations built on top of MessageBus primitives. + +These helpers compose generic bus primitives (``log_append``, ``publish``, +``queue_push``) with domain-specific key layouts from ``MessageBusKeys``. +They live here — between the transport layer (``message_bus``) and the +service layer (``_service``) — so that neither layer needs to know about the +other's internals. + +.. list-table:: + :widths: 30 70 + + * - :func:`publish_session_event` + - Append an event to the session replay log and fan it out live. + * - :func:`enqueue_run_trigger` + - Enqueue a typed run trigger and signal dispatchers. + * - :func:`enqueue_index_task` + - Enqueue a knowledge-document indexing task and signal consumers. +""" +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal + +from .message_bus._keys import MessageBusKeys + +if TYPE_CHECKING: + from .message_bus._base import MessageBus + + from agentscope.event import ( + ExternalExecutionResultEvent, + UserConfirmResultEvent, + ) + + +# ── publish_session_event ────────────────────────────────────────────── + + +async def publish_session_event( + bus: "MessageBus", + session_id: str, + event: dict, +) -> str: + """Append event to replay log + fan out live. + + Args: + bus (`MessageBus`): + The application message bus. + session_id (`str`): + The session this event belongs to. + event (`dict`): + JSON-serializable event payload. + + Returns: + `str`: + The replay-log entry id assigned by the backend. + """ + key = MessageBusKeys.session_events(session_id) + entry_id = await bus.log_append( + key, + event, + max_len=MessageBusKeys.SESSION_REPLAY_MAX_LEN, + ) + await bus.publish(key, {**event, "_entry_id": entry_id}) + return entry_id + + +# ── enqueue_run_trigger ──────────────────────────────────────────────── + + +async def enqueue_run_trigger( + bus: "MessageBus", + user_id: str, + session_id: str, + agent_id: str, + *, + kind: Literal["wake", "resume"] = MessageBusKeys.WAKEUP_KIND_WAKE, + inputs: UserConfirmResultEvent + | ExternalExecutionResultEvent + | None = None, +) -> None: + """Enqueue a typed run trigger and signal dispatchers. + + ``kind`` selects how the dispatcher handles the entry: + + - ``wake`` — idle-session wake-up. The dispatcher skips the entry + when the session is already running (the live run drains the inbox + itself). ``inputs`` must be ``None``. + - ``resume`` — resume a HITL-parked session with a user confirmation + or external execution result. The dispatcher waits (with backoff) + until the parked run releases its lock, then spawns with + ``input_msg`` set to the deserialised event. + + The payload is serialised to a plain dict before being pushed to the + wakeup queue; the ``MessageBus`` transport layer never sees event + types. + + Args: + bus (`MessageBus`): + The application message bus. + user_id (`str`): + The owning user id. + session_id (`str`): + The session to trigger a run for. + agent_id (`str`): + The agent id that owns the session. + kind: + Trigger kind. Defaults to ``"wake"``. + inputs: + The input event for ``resume`` triggers. Ignored (and + should be ``None``) for ``wake``. The function calls + ``model_dump(mode="json")`` internally — callers pass the + event object, not a pre-serialised dict. + """ + await bus.queue_push( + MessageBusKeys.wakeup_queue(), + { + "user_id": user_id, + "session_id": session_id, + "agent_id": agent_id, + "kind": kind, + "input": inputs.model_dump(mode="json") if inputs else None, + }, + ) + await bus.publish(MessageBusKeys.wakeup_signal(), {}) + + +# ── enqueue_index_task ───────────────────────────────────────────────── + + +async def enqueue_index_task( + bus: "MessageBus", + user_id: str, + knowledge_base_id: str, + document_id: str, +) -> None: + """Enqueue a knowledge-document indexing task and signal consumers. + + Pushes a structured payload onto the durable index-task queue and + publishes a signal so any subscribed + :class:`~agentscope.app._service.IndexTaskConsumer` drains it within + one ``subscribe`` round-trip. + + The push happens *before* the publish so a worker woken by the + signal is guaranteed to find the entry on its drain. Re-enqueuing + the same document is safe — the worker's lease CAS rejects + duplicates — so the queue may legitimately hold multiple entries + for the same document (one from upload, one from sweeper). + + Args: + bus (`MessageBus`): + The application message bus. + user_id (`str`): + The owning user id. + knowledge_base_id (`str`): + The parent knowledge base id. + document_id (`str`): + The document id to index. + """ + await bus.queue_push( + MessageBusKeys.index_tasks_queue(), + { + "user_id": user_id, + "knowledge_base_id": knowledge_base_id, + "document_id": document_id, + }, + ) + await bus.publish(MessageBusKeys.index_tasks_signal(), {}) diff --git a/src/agentscope/app/_lifespan.py b/src/agentscope/app/_lifespan.py new file mode 100644 index 0000000000000000000000000000000000000000..70a081ceb1e41ba8ea28ed8251588de28de4b0f1 --- /dev/null +++ b/src/agentscope/app/_lifespan.py @@ -0,0 +1,180 @@ +# -*- coding: utf-8 -*- +"""The lifespan of the agent service.""" +import socket +import uuid +from contextlib import AsyncExitStack, asynccontextmanager +from typing import TYPE_CHECKING, Any, AsyncIterator + +from ._manager import ( + BackgroundTaskManager, + CancelDispatcher, + ChatRunRegistry, + SchedulerManager, + WakeupDispatcher, +) +from ._service import ( + ChatService, + IndexSweeper, + IndexTaskConsumer, + IndexWorker, + KnowledgeBaseService, + SessionService, +) + +if TYPE_CHECKING: + from fastapi import FastAPI +else: + FastAPI = Any + + +@asynccontextmanager +async def lifespan(app: FastAPI) -> AsyncIterator[None]: + """Manage startup and shutdown of all application-wide resources. + + Every resource with a lifecycle is an async context manager and is + entered through a single :class:`AsyncExitStack`. The stack tears + them down in reverse order on shutdown — including when an entry + later in the sequence raises during startup, so no resource leaks + on partial failure. + + Service-layer ``ChatService`` and ``SessionService`` have no + lifecycle of their own and are constructed inline. + """ + storage = app.state.storage + message_bus = app.state.message_bus + workspace_manager = app.state.workspace_manager + knowledge_base_manager = app.state.knowledge_base_manager + blob_store = app.state.blob_store + enable_index_worker = app.state.enable_index_worker + + async with AsyncExitStack() as stack: + await stack.enter_async_context(storage) + await stack.enter_async_context(message_bus) + await stack.enter_async_context(workspace_manager) + if knowledge_base_manager is not None: + # ``KnowledgeBaseManagerBase.__aenter__`` enters the bound + # vector store too, so a single context covers both. + await stack.enter_async_context(knowledge_base_manager) + if blob_store is not None: + await stack.enter_async_context(blob_store) + + bg_manager = await stack.enter_async_context( + BackgroundTaskManager(message_bus=message_bus), + ) + app.state.background_task_manager = bg_manager + + # Per-process registry of in-flight chat-run asyncio tasks. + # Entered before the wake-up + cancel dispatchers so they can + # share the same registry; exited last so its shutdown can + # cancel any leftover runs after the dispatchers stop. + chat_run_registry = await stack.enter_async_context(ChatRunRegistry()) + app.state.chat_run_registry = chat_run_registry + + # Scheduler is independent of ChatService now (its fire path + # pushes to inbox + enqueues wakeup via the bus), so we build it + # before ChatService and inject it via the constructor. + scheduler = await stack.enter_async_context( + SchedulerManager( + storage=storage, + message_bus=message_bus, + ), + ) + app.state.scheduler_manager = scheduler + + chat_service = ChatService( + storage=storage, + workspace_manager=workspace_manager, + scheduler_manager=scheduler, + background_task_manager=bg_manager, + message_bus=message_bus, + knowledge_base_manager=knowledge_base_manager, + extra_agent_middlewares=app.state.extra_agent_middlewares, + extra_agent_tools=app.state.extra_agent_tools, + custom_subagent_templates=app.state.custom_subagent_templates, + custom_agent_cls=app.state.custom_agent_cls, + ) + app.state.chat_service = chat_service + + app.state.session_service = SessionService( + storage=storage, + message_bus=message_bus, + ) + + # ---------------- Knowledge-base wiring ---------------- + knowledge_base_service = None + if knowledge_base_manager is not None: + # Indexing is uniformly driven by the message bus: the + # service publishes an index-task entry, and a consumer + # (in-process or in a dedicated worker process) drains it. + # + # * Embedded — ``enable_index_worker=True``: this lifespan + # additionally starts an :class:`IndexWorker` plus an + # :class:`IndexTaskConsumer` that subscribes to the same + # channel. The ``InMemoryMessageBus`` makes the round + # trip near-free; everything runs in one binary. + # + # * Dedicated — ``enable_index_worker=False``: this + # lifespan does NOT start a worker. One or more separate + # processes (``python -m agentscope.app.rag.index_worker``) + # run their own consumer + worker pair subscribed to the + # same channel. + # + # The sweeper STILL runs in the API process either way, + # because the API is the only resource guaranteed to be + # live whenever uploads happen — if the publish ever races + # a worker restart the durable queue catches the task, and + # if the queue write itself failed the sweeper eventually + # re-enqueues from storage. + if enable_index_worker: + node_id = f"{socket.gethostname()}:{uuid.uuid4().hex[:8]}" + worker = IndexWorker( + storage=storage, + blob_store=blob_store, + knowledge_base_manager=knowledge_base_manager, + parsers=app.state.knowledge_parsers, + chunker=app.state.knowledge_chunker, + node_id=node_id, + ) + await stack.enter_async_context( + IndexTaskConsumer( + message_bus=message_bus, + worker=worker, + ), + ) + + sweeper = IndexSweeper( + storage=storage, + message_bus=message_bus, + ) + await sweeper.start() + stack.push_async_callback(sweeper.stop) + + knowledge_base_service = KnowledgeBaseService( + storage=storage, + knowledge_base_manager=knowledge_base_manager, + blob_store=blob_store, + message_bus=message_bus, + ) + + app.state.knowledge_base_service = knowledge_base_service + + # Dispatchers need live references somewhere, or they would be + # garbage-collected; the AsyncExitStack holds those references + # for us, so we don't need local bindings or app.state slots. + await stack.enter_async_context( + WakeupDispatcher( + message_bus=message_bus, + storage=storage, + chat_service=chat_service, + chat_run_registry=chat_run_registry, + ), + ) + await stack.enter_async_context( + CancelDispatcher( + message_bus=message_bus, + registry=chat_run_registry, + bg_manager=bg_manager, + ), + ) + + yield diff --git a/src/agentscope/app/_manager/__init__.py b/src/agentscope/app/_manager/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..82b418e814168340d059f60cd3ce9137f7186c2c --- /dev/null +++ b/src/agentscope/app/_manager/__init__.py @@ -0,0 +1,17 @@ +# -*- coding: utf-8 -*- +"""The agent service managers, used in FastAPI lifespan to manage +application-wide resources.""" + +from ._scheduler import SchedulerManager +from ._wakeup_dispatcher import WakeupDispatcher +from ._cancel_dispatcher import CancelDispatcher +from ._chat_run_registry import ChatRunRegistry +from ._background_task_manager import BackgroundTaskManager + +__all__ = [ + "BackgroundTaskManager", + "CancelDispatcher", + "ChatRunRegistry", + "SchedulerManager", + "WakeupDispatcher", +] diff --git a/src/agentscope/app/_manager/_background_task_manager.py b/src/agentscope/app/_manager/_background_task_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..3cc8e486ba3b99e55bedc1955b7768039f54cfec --- /dev/null +++ b/src/agentscope/app/_manager/_background_task_manager.py @@ -0,0 +1,469 @@ +# -*- coding: utf-8 -*- +"""The background task manager.""" +import asyncio +import json +import time +from collections import OrderedDict +from dataclasses import dataclass, field +from typing import Any, Self, TYPE_CHECKING + +import shortuuid +from pydantic import BaseModel, Field + +from agentscope.message import TextBlock, ToolResultState +from agentscope.permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, +) +from agentscope.tool import ToolBase, ToolChunk +from agentscope._logging import logger +from ..message_bus import MessageBusKeys + +if TYPE_CHECKING: + from ..message_bus import MessageBus + + +@dataclass +class BackgroundTask: + """Metadata for a single background task. + + Attributes: + asyncio_task (`asyncio.Task`): + The running asyncio task. + session_id (`str`): + The session id of the originating request. + agent_id (`str`): + The name of the agent that created the task. + user_id (`str`): + The user id of the originating request. + tool_name (`str`): + The name of the tool that was offloaded. + id (`str`): + Auto-generated unique task identifier. + """ + + asyncio_task: asyncio.Task + """The running asyncio task.""" + + session_id: str + """The session id of the background task.""" + + agent_id: str + """The agent that created the background task.""" + + user_id: str + """The user id of the originating request.""" + + tool_name: str + """The name of the offloaded tool.""" + + id: str = field(default_factory=shortuuid.uuid) + """The background task id.""" + + +class _ToolStopParams(BaseModel): + """The params of the stop tool.""" + + task_id: str = Field( + description="The task id of the background tool to stop.", + ) + + +class ToolStop(ToolBase): + """A tool to stop a running background tool execution.""" + + name: str = "ToolStop" + """The tool name.""" + + description: str = ( + "Stop a background tool execution by its task id. " + "Use this when you want to cancel a previously offloaded tool " + "that is still running in the background." + ) + """The tool description.""" + + input_schema: dict = _ToolStopParams.model_json_schema() + """The input schema.""" + + is_concurrency_safe: bool = True + is_read_only: bool = False + is_state_injected: bool = False + is_external_tool: bool = False + is_mcp: bool = False + mcp_name: str | None = None + + def __init__( + self, + background_tasks: dict[str, BackgroundTask], + message_bus: "MessageBus", + session_id: str, + ) -> None: + """Initialize the ToolStop tool. + + Args: + background_tasks (`dict[str, BackgroundTask]`): + A reference to the local background tasks managed by + the :class:`BackgroundTaskManager`. + message_bus (`MessageBus`): + The application message bus, used to check the global + registry and broadcast cross-worker cancel requests. + session_id (`str`): + The current session id, used to scope Redis registry + lookups. + """ + self.background_tasks = background_tasks + self._message_bus = message_bus + self._session_id = session_id + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permission for the tool usage. + + Args: + tool_input (`dict[str, Any]`): + The tool input parameters. + context (`PermissionContext`): + The permission context. + + Returns: + `PermissionDecision`: + Always returns ALLOW. + """ + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message=f"{self.name} is always allowed to be called.", + ) + + async def __call__(self, task_id: str) -> ToolChunk: + """Stop the background task. + + Args: + task_id (`str`): + The task id. + + Returns: + `ToolChunk`: + The tool chunk. + """ + # Path 1: task is on this worker — cancel directly. + # Only cancel when the task belongs to the same session as this + # ToolStop instance, so a leaked/guessed task_id from another + # session cannot trigger cross-session cancellation on a shared + # worker. + local_task = self.background_tasks.get(task_id) + if ( + local_task is not None + and local_task.session_id == self._session_id + ): + self.background_tasks.pop(task_id, None) + local_task.asyncio_task.cancel() + logger.info( + "Background task stopped via ToolStop (local): task_id=%s, " + "session_id=%s, agent_id=%s", + task_id, + local_task.session_id, + local_task.agent_id, + ) + return ToolChunk( + content=[ + TextBlock(text=f"Task {task_id} stopped successfully."), + ], + state=ToolResultState.SUCCESS, + ) + + # Path 2: task exists in the global registry (another worker, or + # a different session on this worker). + if await self._message_bus.registry_exists( + MessageBusKeys.bg_tasks(self._session_id), + task_id, + ): + await self._message_bus.publish( + MessageBusKeys.task_cancel_channel(), + {"task_id": task_id}, + ) + logger.info( + "Background task cancel broadcast via ToolStop (remote): " + "task_id=%s, session_id=%s", + task_id, + self._session_id, + ) + return ToolChunk( + content=[ + TextBlock( + text=f"Cancel request sent for task {task_id}. " + f"The owning worker will stop it shortly.", + ), + ], + state=ToolResultState.SUCCESS, + ) + + # Path 3: task not found anywhere. + return ToolChunk( + content=[ + TextBlock( + text=f"TaskNotFoundError: The task {task_id} " + f"does not exist.", + ), + ], + state=ToolResultState.ERROR, + ) + + +class BackgroundTaskManager: + """Tracks background asyncio task lifecycle within the agent service. + + Responsibilities: + + - **Global registry**: register/unregister tasks in Redis so any + process can query which tasks are alive for a session. + - **Local handle cache**: hold ``asyncio.Task`` references for + cancel and shutdown. + - **Task scheduling**: convenience method for creating a task from + a plain coroutine with a done callback that cleans up both sides. + + Completion results are delivered via the :class:`MessageBus` inbox + + wakeup path (same as team messages), so any process's + :class:`WakeupDispatcher` can pick up the result. + """ + + def __init__(self, message_bus: "MessageBus") -> None: + """Initialise the background task manager. + + Args: + message_bus (`MessageBus`): + The application message bus; used for the global BG + task registry (Redis Hash) and task-level cancel + broadcasts. + """ + self._message_bus = message_bus + self.tasks: OrderedDict[str, BackgroundTask] = OrderedDict() + + # ------------------------------------------------------------------ + # Task registration + # ------------------------------------------------------------------ + + async def register_task( + self, + asyncio_task: asyncio.Task, + session_id: str, + agent_id: str, + user_id: str, + tool_name: str = "", + ) -> str: + """Register an already-running asyncio task. + + Writes to both the local handle cache and the global Redis + registry. The task auto-removes from both when it finishes + (via ``add_done_callback``). + + Args: + asyncio_task (`asyncio.Task`): + The already-running task to register. + session_id (`str`): + The originating session id. + agent_id (`str`): + The agent record id that owns the task. + user_id (`str`): + The user id of the originating request. + tool_name (`str`, optional): + The name of the offloaded tool. + + Returns: + `str`: + The generated task id. + """ + bg_task = BackgroundTask( + asyncio_task=asyncio_task, + session_id=session_id, + agent_id=agent_id, + user_id=user_id, + tool_name=tool_name, + ) + task_id = bg_task.id + self.tasks[task_id] = bg_task + + # Register in the global Redis registry. + metadata = json.dumps( + { + "tool_name": tool_name, + "agent_id": agent_id, + "started_at": time.time(), + }, + ) + await self._message_bus.registry_set( + MessageBusKeys.bg_tasks(session_id), + task_id, + metadata, + ttl_secs=MessageBusKeys.BG_TASKS_TTL_SECS, + ) + + logger.info( + "Background task registered: task_id=%s, session_id=%s, " + "agent_id=%s, tool_name=%s", + task_id, + session_id, + agent_id, + tool_name, + ) + + def _on_done(_t: asyncio.Task) -> None: + self.tasks.pop(task_id, None) + # Schedule async Redis cleanup (fire-and-forget). Wrap in a + # coroutine that logs failures so the bus error (e.g. Redis + # connection drop) does not surface as + # ``Task exception was never retrieved``. + try: + asyncio.ensure_future( + self._safe_bg_task_unregister(session_id, task_id), + ) + except RuntimeError: + # Event loop already closed during shutdown. + pass + + asyncio_task.add_done_callback(_on_done) + return task_id + + async def _safe_bg_task_unregister( + self, + session_id: str, + task_id: str, + ) -> None: + """Unregister a finished background task, logging any failure. + + Args: + session_id (`str`): + The session id of the finished task. + task_id (`str`): + The task id to unregister from the global registry. + """ + try: + await self._message_bus.registry_del( + MessageBusKeys.bg_tasks(session_id), + task_id, + ) + except Exception as e: # pylint: disable=broad-except + logger.exception( + "Failed to unregister background task from the global " + "registry: task_id=%s, session_id=%s, error=%s", + task_id, + session_id, + str(e), + ) + + # ------------------------------------------------------------------ + # Tool listing + # ------------------------------------------------------------------ + + async def list_tools(self, session_id: str) -> list[ToolBase]: + """List the background task tools for a given session. + + Args: + session_id (`str`): + The current session id (for ToolStop's registry + lookups). + + Returns: + `list[ToolBase]`: + A list containing the :class:`ToolStop` tool. + """ + return [ToolStop(self.tasks, self._message_bus, session_id)] + + # ------------------------------------------------------------------ + # Session-scoped cancel + # ------------------------------------------------------------------ + + def cancel_session_tasks(self, session_id: str) -> int: + """Cancel every locally-tracked task whose owner session matches. + + Called by :class:`CancelDispatcher` on each incoming session + cancel broadcast. Returns the number of tasks cancelled on this + process. + + Args: + session_id (`str`): + The session whose tasks should be cancelled. + + Returns: + `int`: + Number of tasks cancelled locally. + """ + cancelled = 0 + for bg_task in list(self.tasks.values()): + if bg_task.session_id != session_id: + continue + logger.info( + "Cancelling background task for session cancel: " + "task_id=%s, session_id=%s, agent_id=%s", + bg_task.id, + bg_task.session_id, + bg_task.agent_id, + ) + bg_task.asyncio_task.cancel() + cancelled += 1 + return cancelled + + # ------------------------------------------------------------------ + # Single-task cancel (called by CancelDispatcher on bus signal) + # ------------------------------------------------------------------ + + def cancel_task(self, task_id: str) -> bool: + """Cancel a single locally-tracked task by its id. + + Called by :class:`CancelDispatcher` when a task-level cancel + broadcast arrives. Returns whether the task was found and + cancelled on this process. + + Args: + task_id (`str`): + The task to cancel. + + Returns: + `bool`: + ``True`` if the task was found locally and cancelled. + """ + bg_task = self.tasks.get(task_id) + if bg_task is None: + return False + logger.info( + "Cancelling background task via bus signal: " + "task_id=%s, session_id=%s, agent_id=%s", + task_id, + bg_task.session_id, + bg_task.agent_id, + ) + bg_task.asyncio_task.cancel() + return True + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def __aenter__(self) -> Self: + """Enter the async context. No setup required. + + Returns: + `Self`: This manager instance. + """ + return self + + async def __aexit__(self, *exc: object) -> None: + """Cancel all running background tasks on context exit.""" + count = len(self.tasks) + logger.info( + "Shutting down BackgroundTaskManager: cancelling %d task(s).", + count, + ) + for bg_task in list(self.tasks.values()): + logger.info( + "Cancelling background task on shutdown: task_id=%s, " + "session_id=%s, agent_id=%s", + bg_task.id, + bg_task.session_id, + bg_task.agent_id, + ) + bg_task.asyncio_task.cancel() + self.tasks.clear() diff --git a/src/agentscope/app/_manager/_cancel_dispatcher.py b/src/agentscope/app/_manager/_cancel_dispatcher.py new file mode 100644 index 0000000000000000000000000000000000000000..e659b8202db7c4283843684d54a39046fedf86c2 --- /dev/null +++ b/src/agentscope/app/_manager/_cancel_dispatcher.py @@ -0,0 +1,186 @@ +# -*- coding: utf-8 -*- +"""Single per-process dispatcher for cross-process cancels. + +Subscribes to two bus channels: + +1. **Session cancel** — cancel all local work for a session (chat run + + all BG tasks). Triggered by session deletion or explicit abort. +2. **Task cancel** — cancel a single BG task by task_id. Triggered by + the :class:`ToolStop` agent tool when the target task lives on a + different worker. + +Processes whose registry / BG-manager do not hold the targeted session +or task simply do no work — the publisher does not need to know which +worker holds what; it broadcasts and lets each holder self-select. +""" +import asyncio +from typing import TYPE_CHECKING, Self + +from ..._logging import logger +from ..message_bus import MessageBusKeys + +if TYPE_CHECKING: + from ..message_bus import MessageBus + from ._background_task_manager import BackgroundTaskManager + from ._chat_run_registry import ChatRunRegistry + + +class CancelDispatcher: + """Subscribes to bus cancel channels and cancels matching local + tasks. + + Args: + message_bus (`MessageBus`): + Application message bus. + registry (`ChatRunRegistry`): + The per-process chat-run registry whose tasks may be + cancelled. + bg_manager (`BackgroundTaskManager`): + The per-process background task manager. + """ + + def __init__( + self, + message_bus: "MessageBus", + registry: "ChatRunRegistry", + bg_manager: "BackgroundTaskManager", + ) -> None: + """Bind dependencies. + + Args: + message_bus (`MessageBus`): + Application message bus. + registry (`ChatRunRegistry`): + The per-process chat-run registry. + bg_manager (`BackgroundTaskManager`): + The per-process background task manager. + """ + self._bus = message_bus + self._registry = registry + self._bg_manager = bg_manager + self._session_task: asyncio.Task | None = None + self._task_cancel_task: asyncio.Task | None = None + + async def __aenter__(self) -> Self: + """Start both dispatcher loops and wait until their bus + subscriptions are live. + + Returns: + `Self`: This dispatcher instance. + """ + session_ready = asyncio.Event() + task_ready = asyncio.Event() + + self._session_task = asyncio.create_task( + self._session_cancel_loop(session_ready), + name="cancel-dispatcher:session", + ) + self._task_cancel_task = asyncio.create_task( + self._task_cancel_loop(task_ready), + name="cancel-dispatcher:task", + ) + + await session_ready.wait() + await task_ready.wait() + return self + + async def __aexit__(self, *exc: object) -> None: + """Cancel both dispatcher loops on context exit.""" + for task in (self._session_task, self._task_cancel_task): + if task is None: + continue + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + self._session_task = None + self._task_cancel_task = None + + # ------------------------------------------------------------------ + # Session-level cancel loop + # ------------------------------------------------------------------ + + async def _session_cancel_loop(self, ready: asyncio.Event) -> None: + """Subscribe to session cancel channel and act on each signal. + + Args: + ready (`asyncio.Event`): + Signalled after the underlying SUBSCRIBE completes. + """ + try: + async for payload in self._bus.subscribe( + MessageBusKeys.session_cancel_channel(), + on_ready=ready.set, + ): + sid = payload.get("session_id") + if isinstance(sid, str): + self._cancel_session(sid) + except Exception: # pylint: disable=broad-except + logger.exception( + "CancelDispatcher session-cancel loop crashed.", + ) + finally: + # Unblock ``__aenter__`` even if subscribe failed before + # ``on_ready`` ran, so startup cannot deadlock. + ready.set() + + def _cancel_session(self, session_id: str) -> None: + """Cancel every locally-tracked task for a session. + + Args: + session_id (`str`): + The session whose runs and BG tasks should be cancelled. + """ + task = self._registry.get(session_id) + if task is not None and not task.done(): + logger.info( + "CancelDispatcher: cancelling local chat run for " + "session %s", + session_id, + ) + task.cancel() + + bg_cancelled = self._bg_manager.cancel_session_tasks(session_id) + if bg_cancelled: + logger.info( + "CancelDispatcher: cancelled %d local BG task(s) for " + "session %s", + bg_cancelled, + session_id, + ) + + # ------------------------------------------------------------------ + # Task-level cancel loop + # ------------------------------------------------------------------ + + async def _task_cancel_loop(self, ready: asyncio.Event) -> None: + """Subscribe to the task cancel channel and act on each signal. + + Args: + ready (`asyncio.Event`): + Signalled after the underlying SUBSCRIBE completes. + """ + try: + async for payload in self._bus.subscribe( + MessageBusKeys.task_cancel_channel(), + on_ready=ready.set, + ): + task_id = payload.get("task_id") + if not isinstance(task_id, str): + continue + cancelled = self._bg_manager.cancel_task(task_id) + if cancelled: + logger.info( + "CancelDispatcher: cancelled local BG task %s " + "via task-level broadcast.", + task_id, + ) + except Exception: # pylint: disable=broad-except + logger.exception( + "CancelDispatcher task-cancel loop crashed.", + ) + finally: + # Unblock ``__aenter__`` even if subscribe failed before + # ``on_ready`` ran, so startup cannot deadlock. + ready.set() diff --git a/src/agentscope/app/_manager/_chat_run_registry.py b/src/agentscope/app/_manager/_chat_run_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..8bd641dfbdfd15158ed13f0511825f93d1626c95 --- /dev/null +++ b/src/agentscope/app/_manager/_chat_run_registry.py @@ -0,0 +1,134 @@ +# -*- coding: utf-8 -*- +"""Per-process registry of in-flight ``ChatService.run`` asyncio tasks. + +Owns the asyncio.Task handles only — it is not the public cancel +entry point. The cross-process cancel path goes through the bus's +:meth:`~agentscope.app.message_bus.MessageBus.session_publish_cancel` +broadcast, picked up locally by +:class:`~agentscope.app._manager.CancelDispatcher`, which then looks +up the task here and calls ``.cancel()`` on it. + +A given ``session_id`` can have at most one entry. Concurrent runs for +the same session are already prevented at a cluster level by +:meth:`~agentscope.app.message_bus.MessageBus.session_run` (the +distributed lock), so a second :meth:`spawn` for the same id is treated +as a programming error. +""" +import asyncio +from typing import Coroutine, Self + +from ..._logging import logger + + +class ChatRunRegistry: + """In-process index of active chat-run asyncio tasks, keyed by + session id. + + Used by :class:`~agentscope.app._manager.CancelDispatcher` to find + and cancel the local task for a given session, and by the lifespan + to cancel any leftover runs on application shutdown. + """ + + def __init__(self) -> None: + """Initialise an empty registry.""" + self._tasks: dict[str, asyncio.Task] = {} + + def spawn( + self, + coro: Coroutine, + *, + session_id: str, + name: str | None = None, + ) -> asyncio.Task: + """Create and register an asyncio task that runs ``coro``. + + The task auto-removes from the registry when it finishes (via + ``add_done_callback``). + + Args: + coro (`Coroutine`): + A coroutine — typically ``chat_service.run(...)`` — to + run as a background task. + session_id (`str`): + The session this run belongs to. Used as the registry + key for later cancel lookup. + name (`str | None`, optional): + Optional task name passed through to + :func:`asyncio.create_task`. Defaults to + ``f"chat-run:{session_id}"``. + + Returns: + `asyncio.Task`: + The created task. Callers normally do not need to keep + the reference — the registry holds it for the task's + lifetime. + + Raises: + `RuntimeError`: + When a non-finished task is already registered for + ``session_id``. Callers are expected to coordinate via + the distributed session lock before spawning. + """ + existing = self._tasks.get(session_id) + if existing is not None and not existing.done(): + raise RuntimeError( + f"Session {session_id!r} already has an active chat run " + "in this process.", + ) + + task = asyncio.create_task( + coro, + name=name or f"chat-run:{session_id}", + ) + self._tasks[session_id] = task + + def _cleanup(t: asyncio.Task) -> None: + # Only remove the entry if it still points at this task — + # a fresh spawn for the same sid may have replaced it. + if self._tasks.get(session_id) is t: + self._tasks.pop(session_id, None) + + task.add_done_callback(_cleanup) + return task + + def get(self, session_id: str) -> asyncio.Task | None: + """Return the registered task for ``session_id``, or ``None``. + + Args: + session_id (`str`): + The session whose task to look up. + + Returns: + `asyncio.Task | None`: + The task if one is currently registered for the + session, else ``None``. + """ + return self._tasks.get(session_id) + + async def __aenter__(self) -> Self: + """No-op enter; the registry has no startup work. + + Returns: + `Self`: This registry instance. + """ + return self + + async def __aexit__(self, *exc: object) -> None: + """Cancel every still-running task on application shutdown. + + Each task is cancelled and awaited so its ``finally`` blocks + and any ``async with`` cleanups (notably the bus's session + run-lock release) execute before the process exits. + """ + if not self._tasks: + return + logger.info( + "ChatRunRegistry shutdown: cancelling %d in-flight chat run(s).", + len(self._tasks), + ) + tasks = list(self._tasks.values()) + for task in tasks: + task.cancel() + # Wait for every cancel to land; swallow CancelledError per task. + await asyncio.gather(*tasks, return_exceptions=True) + self._tasks.clear() diff --git a/src/agentscope/app/_manager/_scheduler/__init__.py b/src/agentscope/app/_manager/_scheduler/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..960f6cdbfa70a433aed4ad3eca0afba97b1ac716 --- /dev/null +++ b/src/agentscope/app/_manager/_scheduler/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- +"""The scheduler related components.""" + +from ._scheduler_manager import SchedulerManager + +__all__ = [ + "SchedulerManager", +] diff --git a/src/agentscope/app/_manager/_scheduler/_scheduler_manager.py b/src/agentscope/app/_manager/_scheduler/_scheduler_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..4ec4c7f832230b09b20c616bb029d94f0390131d --- /dev/null +++ b/src/agentscope/app/_manager/_scheduler/_scheduler_manager.py @@ -0,0 +1,432 @@ +# -*- coding: utf-8 -*- +"""The cron scheduler manager class.""" +import json +from collections.abc import Callable, Coroutine + +from typing import Self + +from ....message import HintBlock +from ....permission import PermissionContext +from ....state import AgentState +from ....tool import ToolBase +from ...._logging import logger +from ._tools import ScheduleCreate, ScheduleDelete, ScheduleList, ScheduleView +from ...message_bus import MessageBus, MessageBusKeys +from ..._bus_ops import enqueue_run_trigger +from ...storage import ( + StorageBase, + ScheduleRecord, + ChatModelConfig, + SessionConfig, + SessionSource, +) + + +class SchedulerManager: + """The cron scheduler manager, responsible for managing scheduled-task + lifecycle within the agent service. + + The manager owns both the in-memory APScheduler instance and the trigger + logic that fires scheduled tasks. Triggers do not call ``ChatService`` + directly; instead they push a :class:`HintBlock` to the target session's + inbox and enqueue a wakeup, so that the application-wide + :class:`WakeupDispatcher` (running on any process) picks up the work. + This keeps the scheduler decoupled from ``ChatService`` and makes the + fire path consistent with team / background-tool result delivery. + """ + + def __init__( + self, + storage: StorageBase, + message_bus: MessageBus, + ) -> None: + """Initialize the scheduler manager. + + Args: + storage (`StorageBase`): + The storage backend used for persistence and session + creation. + message_bus (`MessageBus`): + The application message bus. Each scheduled fire pushes + a :class:`HintBlock` to the target session's inbox and + enqueues a wakeup via this bus. + """ + from apscheduler.schedulers.asyncio import AsyncIOScheduler + + self._storage = storage + self._message_bus = message_bus + self._scheduler = AsyncIOScheduler() + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def __aenter__(self) -> Self: + """Start APScheduler and re-register persisted schedules. + + Reading all schedules from storage and restoring them is the + only thing a caller would ever do right after starting this + manager, so the work lives inside the context entry — the + lifespan does not need to remember to call :meth:`restore`. + + Returns: + `Self`: This manager instance. + """ + logger.info("SchedulerManager starting APScheduler") + self._scheduler.start() + logger.info("SchedulerManager APScheduler started") + + records = await self._storage.list_all_schedules() + if records: + await self.restore(records) + + return self + + async def __aexit__(self, *exc: object) -> None: + """Shut down the underlying APScheduler on context exit.""" + logger.info("SchedulerManager shutting down APScheduler") + self._scheduler.shutdown() + logger.info("SchedulerManager APScheduler shut down") + + # ------------------------------------------------------------------ + # Trigger construction + # ------------------------------------------------------------------ + + def _build_trigger( + self, + record: ScheduleRecord, + ) -> Callable[[], Coroutine]: + """Build the zero-argument coroutine executed by APScheduler on each + trigger fire. + + The returned coroutine: + + 1. Skips execution when the schedule is disabled. + 2. Resolves or creates the target session (stateful reuses a fixed + session; non-stateful creates a fresh one on every fire). + 3. Calls :class:`~agentscope.app._service._chat.ChatService` and + drains the response stream (fire-and-forget). + 4. Catches and logs all exceptions to prevent APScheduler from + removing the job on failure. + + Args: + record (`ScheduleRecord`): + The persisted schedule record that describes what to run. + + Returns: + `Callable[[], Coroutine]`: + A zero-argument async callable suitable for APScheduler. + """ + # Closure-friendly references so APScheduler doesn't have to + # re-look these up on every fire. + storage = self._storage + message_bus = self._message_bus + + async def _trigger() -> None: + logger.info( + "[Schedule:%s(%s)] Trigger fired", + record.id, + record.data.name, + ) + + if not record.data.enabled: + logger.info( + "[Schedule:%s(%s)] Skipped — schedule disabled", + record.id, + record.data.name, + ) + return + + try: + if record.data.stateful: + stateful_session_id = f"{record.id}_stateful" + logger.info( + "[Schedule:%s(%s)] Stateful mode, " + "looking up session %s", + record.id, + record.data.name, + stateful_session_id, + ) + session = await storage.get_session( + record.user_id, + record.agent_id, + stateful_session_id, + ) + if session is None: + logger.info( + "[Schedule:%s(%s)] First fire, " + "creating stateful session", + record.id, + record.data.name, + ) + state = AgentState() + state.permission_context = PermissionContext( + mode=record.data.permission_mode, + ) + session_config = SessionConfig( + workspace_id="", + chat_model_config=record.data.chat_model_config, + ) + session = await storage.upsert_session( + user_id=record.user_id, + agent_id=record.agent_id, + config=session_config, + state=state, + session_id=stateful_session_id, + source=SessionSource.SCHEDULE, + source_schedule_id=record.id, + ) + else: + logger.info( + "[Schedule:%s(%s)] Reusing existing " + "stateful session %s", + record.id, + record.data.name, + session.id, + ) + else: + logger.info( + "[Schedule:%s(%s)] Non-stateful mode, " + "creating fresh session", + record.id, + record.data.name, + ) + state = AgentState() + state.permission_context = PermissionContext( + mode=record.data.permission_mode, + ) + session = await storage.upsert_session( + user_id=record.user_id, + agent_id=record.agent_id, + config=SessionConfig( + workspace_id="", + chat_model_config=record.data.chat_model_config, + ), + state=state, + source=SessionSource.SCHEDULE, + source_schedule_id=record.id, + ) + + logger.info( + "[Schedule:%s(%s)] Session ready: %s, " + "delivering prompt via inbox + wakeup", + record.id, + record.data.name, + session.id, + ) + + # Wrap the schedule prompt in an XML tag so the LLM + # recognises it as a system-driven trigger rather than + # a regular user turn — same shape as team / system + # notification hints. + hint = HintBlock( + hint=( + f"\n" + f"{record.data.description}\n" + f"" + ), + source=json.dumps( + { + "label": "schedule", + "sublabel": record.data.name, + }, + ensure_ascii=False, + ), + ) + await message_bus.queue_push( + MessageBusKeys.inbox(session.id), + hint.model_dump(mode="json"), + ) + await enqueue_run_trigger( + message_bus, + user_id=record.user_id, + session_id=session.id, + agent_id=record.agent_id, + ) + + logger.info( + "[Schedule:%s(%s)] Wakeup enqueued for session %s", + record.id, + record.data.name, + session.id, + ) + + except Exception: + logger.exception( + "[Schedule:%s(%s)] Trigger failed", + record.id, + record.data.name, + ) + + return _trigger + + # ------------------------------------------------------------------ + # Schedule management + # ------------------------------------------------------------------ + + async def register_schedule(self, record: ScheduleRecord) -> str: + """Persist-and-register a schedule record with APScheduler. + + Builds the trigger coroutine via :meth:`_build_trigger` and adds the + job to APScheduler. This is the single entry point used by both the + HTTP API and the :class:`ScheduleCreate` agent tool. + + Args: + record (`ScheduleRecord`): + The fully-populated record (already persisted to storage). + + Returns: + `str`: + The APScheduler job ID (equal to ``record.id``). + """ + + from apscheduler.triggers.cron import CronTrigger + + logger.info( + "Registering schedule %s(%s) cron=%s tz=%s", + record.id, + record.data.name, + record.data.cron_expression, + record.data.timezone, + ) + + # ``CronTrigger.from_crontab`` is a thin helper that only forwards + # the 5 parsed fields and ``timezone`` — it has no parameter for + # ``start_date`` / ``end_date``. Parse the expression ourselves so + # the configured activation window is honoured. + fields = record.data.cron_expression.split() + if len(fields) != 5: + raise ValueError( + "Expected a 5-field cron expression, got " + f"{record.data.cron_expression!r}", + ) + minute, hour, day, month, day_of_week = fields + + trigger = self._build_trigger(record) + job = self._scheduler.add_job( + trigger, + trigger=CronTrigger( + minute=minute, + hour=hour, + day=day, + month=month, + day_of_week=day_of_week, + timezone=record.data.timezone, + start_date=record.data.started_at, + end_date=record.data.ended_at, + ), + id=record.id, + name=record.data.name, + misfire_grace_time=300, + ) + logger.info( + "Schedule %s(%s) registered, next_run=%s", + record.id, + record.data.name, + job.next_run_time, + ) + return job.id + + async def remove_schedule(self, job_id: str) -> None: + """Remove a job from APScheduler. + + Args: + job_id (`str`): + The APScheduler job ID to remove. + """ + from apscheduler.jobstores.base import JobLookupError + + logger.info("Removing schedule job %s", job_id) + try: + self._scheduler.remove_job(job_id) + logger.info("Schedule job %s removed", job_id) + except JobLookupError: + logger.warning("Schedule job %s not found in APScheduler", job_id) + + async def restore(self, records: list[ScheduleRecord]) -> None: + """Re-register persisted schedules on service startup. + + Only enabled schedules are restored. + + Args: + records (`list[ScheduleRecord]`): + All schedule records loaded from storage on startup. + """ + enabled = [r for r in records if r.data.enabled] + logger.info( + "Restoring schedules: %d total, %d enabled", + len(records), + len(enabled), + ) + for record in enabled: + await self.register_schedule(record) + logger.info("Schedule restore complete") + + async def list_tasks(self) -> list[dict]: + """Return a summary of all currently registered APScheduler jobs. + + Returns: + `list[dict]`: + Each entry contains ``id``, ``name``, and ``next_run``. + """ + return [ + { + "id": job.id, + "name": job.name, + "next_run": job.next_run_time, + } + for job in self._scheduler.get_jobs() + ] + + # ------------------------------------------------------------------ + # Agent tools + # ------------------------------------------------------------------ + + async def list_tools( + self, + user_id: str, + agent_id: str, + chat_model_config: ChatModelConfig, + ) -> list[ToolBase]: + """Return the agent-facing tools provided by the scheduler manager. + + Args: + user_id (`str`): + The authenticated user who owns the schedules. + agent_id (`str`): + The agent that will be run by newly created schedules. + chat_model_config (`ChatModelConfig`): + Model configuration inherited from the current session and + stored on new :class:`~...ScheduleRecord` objects. + + Returns: + `list[ToolBase]`: + The four schedule tools: :class:`ScheduleCreate`, + :class:`ScheduleView`, :class:`ScheduleDelete`, and + :class:`ScheduleList`. + """ + return [ + ScheduleCreate( + user_id=user_id, + agent_id=agent_id, + chat_model_config=chat_model_config, + storage=self._storage, + scheduler_manager=self, + ), + ScheduleView( + user_id=user_id, + scheduler=self._scheduler, + storage=self._storage, + ), + ScheduleDelete( + user_id=user_id, + scheduler=self._scheduler, + storage=self._storage, + message_bus=self._message_bus, + ), + ScheduleList( + user_id=user_id, + scheduler=self._scheduler, + storage=self._storage, + ), + ] diff --git a/src/agentscope/app/_manager/_scheduler/_tools/__init__.py b/src/agentscope/app/_manager/_scheduler/_tools/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..aa8e665043704dd3f8ebd77ee83a3688b4e1cf13 --- /dev/null +++ b/src/agentscope/app/_manager/_scheduler/_tools/__init__.py @@ -0,0 +1,14 @@ +# -*- coding: utf-8 -*- +"""The schedule related tools.""" + +from ._schedule_create import ScheduleCreate +from ._schedule_delete import ScheduleDelete +from ._schedule_list import ScheduleList +from ._schedule_view import ScheduleView + +__all__ = [ + "ScheduleCreate", + "ScheduleDelete", + "ScheduleList", + "ScheduleView", +] diff --git a/src/agentscope/app/_manager/_scheduler/_tools/_schedule_create.py b/src/agentscope/app/_manager/_scheduler/_tools/_schedule_create.py new file mode 100644 index 0000000000000000000000000000000000000000..1058d53f5ebb1017f787e5f5c2a5f1a983de2d3d --- /dev/null +++ b/src/agentscope/app/_manager/_scheduler/_tools/_schedule_create.py @@ -0,0 +1,253 @@ +# -*- coding: utf-8 -*- +"""The schedule create tool.""" +from datetime import datetime +from typing import Any + +from pydantic import BaseModel, Field + +from .....message import ToolResultState, TextBlock +from .....permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, + PermissionMode, +) +from .....state import AgentState +from .....tool import ToolBase, ToolChunk +from ....storage import ( + ScheduleData, + ScheduleRecord, + ScheduleSource, + ChatModelConfig, +) + + +class _ScheduleCreateParams(BaseModel): + """The params for the schedule create tool.""" + + name: str = Field(description="Display name of the schedule.") + + description: str = Field( + default="", + description="Description of the schedule, including its purpose.", + ) + + cron_expression: str = Field( + description="Standard 5-field cron expression, e.g. '0 9 * * 1-5'.", + ) + + timezone: str = Field( + default="UTC", + description="IANA timezone name used to evaluate the cron expression, " + "e.g. 'America/New_York' or 'Asia/Shanghai'.", + ) + + enabled: bool = Field( + default=True, + description="Whether the schedule is active immediately after " + "creation. Set to False to create a disabled schedule.", + ) + + started_at: datetime | None = Field( + default=None, + description="ISO-8601 datetime at which the schedule becomes active. " + "Defaults to the current time when not specified.", + ) + + ended_at: datetime | None = Field( + default=None, + description="ISO-8601 datetime at which the schedule stops firing. " + "If not set the schedule runs indefinitely.", + ) + + stateful: bool = Field( + default=False, + description="If True, consecutive executions share the same session " + "context. If False, each execution gets a fresh session.", + ) + + permission_mode: str = Field( + default=PermissionMode.DONT_ASK.value, + description=( + "Permission mode for the agent during scheduled execution. " + f"Allowed values: {[m.value for m in PermissionMode]}. " + "Defaults to 'dont_ask' since no user is present." + ), + ) + + +class ScheduleCreate(ToolBase): + """The schedule create tool. + + Creates a new scheduled task that will execute the current agent at a + given cron interval. The record is persisted to storage and immediately + registered with the in-memory APScheduler. + + The schedule inherits the model configuration of the current session. + The agent that creates the schedule is also the agent that will be run + on each trigger. + """ + + name: str = "ScheduleCreate" + + description: str = """Create a new recurring scheduled task for yourself. \ +You will be notified in a new session each time the schedule is triggered. + +**About the cron expression:** +- Determine your current timezone first, that's very important for setting a \ +correct cron expression. Get it by bash command like `date +%z`, \ +`cat /etc/timezone` or directly ask the user. +- Determine whether the task should run once or recur at an interval, \ +then set the cron expression accordingly. +- For a one-off task, query the current time first and set the cron \ +expression to fire at that specific moment. +- Set `started_at` and `ended_at` to match the user's requirements. \ +When in doubt, ask for clarification before creating the schedule. + +**About the description field:** +- The `description` is the only context available to you when the \ +schedule fires in a new session. Include all necessary details: the goal, \ +expected output, constraints, relevant file paths, and anything else needed \ +to complete the task independently. +""" + + input_schema: dict = _ScheduleCreateParams.model_json_schema() + + is_concurrency_safe: bool = False + is_read_only: bool = False + is_state_injected: bool = True + is_external_tool: bool = False + is_mcp: bool = False + mcp_name: str | None = None + + def __init__( + self, + user_id: str, + agent_id: str, + chat_model_config: ChatModelConfig, + storage: Any, + scheduler_manager: Any, + ) -> None: + """Initialize the schedule create tool. + + Args: + user_id (`str`): + The authenticated user who owns this schedule. + agent_id (`str`): + The agent that will be executed on each trigger. + chat_model_config (`ChatModelConfig`): + Model configuration inherited from the current session. + storage (`Any`): + The storage backend used to persist the schedule record. + scheduler_manager (`Any`): + The scheduler manager used to register the APScheduler job. + Must expose a ``register_schedule(record)`` coroutine. + """ + self._user_id = user_id + self._agent_id = agent_id + self._chat_model_config = chat_model_config + self._storage = storage + self._scheduler_manager = scheduler_manager + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permission for the tool usage.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message=f"{self.name} is always allowed to be called.", + ) + + async def __call__( # type: ignore[override] + self, + name: str, + cron_expression: str, + description: str = "", + timezone: str = "UTC", + enabled: bool = True, + started_at: datetime | None = None, + ended_at: datetime | None = None, + stateful: bool = False, + permission_mode: str = PermissionMode.DONT_ASK.value, + _agent_state: AgentState | None = None, + ) -> ToolChunk: + """Create a new scheduled task. + + Args: + name (`str`): + Display name of the schedule. + cron_expression (`str`): + Standard 5-field cron expression, e.g. ``'0 9 * * 1-5'``. + description (`str`, optional): + Human-readable description of what this schedule does. + timezone (`str`, optional): + IANA timezone name, e.g. ``'Asia/Shanghai'``. + enabled (`bool`, optional): + Whether the schedule is active immediately after creation. + started_at (`datetime | None`, optional): + Datetime at which the schedule becomes active. Defaults to + the current time when not specified. + ended_at (`datetime | None`, optional): + Datetime at which the schedule stops firing. If not set the + schedule runs indefinitely. + stateful (`bool`, optional): + Whether consecutive executions share the same session context. + permission_mode (`str`, optional): + Permission mode value string. + _agent_state (`AgentState | None`, optional): + Injected agent state; provides the source session ID. + + Returns: + `ToolChunk`: + A chunk with the new schedule ID on success, or an error + description on failure. + """ + try: + perm_mode = PermissionMode(permission_mode) + except ValueError: + perm_mode = PermissionMode.DONT_ASK + + source_session_id = ( + _agent_state.session_id if _agent_state is not None else "" + ) + + record = ScheduleRecord( + user_id=self._user_id, + agent_id=self._agent_id, + data=ScheduleData( + name=name, + description=description, + enabled=enabled, + cron_expression=cron_expression, + timezone=timezone, + started_at=started_at or datetime.now(), + ended_at=ended_at, + stateful=stateful, + permission_mode=perm_mode, + source=ScheduleSource.AGENT, + source_session_id=source_session_id, + chat_model_config=self._chat_model_config, + ), + ) + + await self._storage.upsert_schedule(self._user_id, record) + await self._scheduler_manager.register_schedule(record) + + return ToolChunk( + content=[ + TextBlock( + text=( + f"Schedule {name!r} created successfully.\n" + f"Schedule ID: {record.id}\n" + f"Cron: {cron_expression} (timezone: {timezone})\n" + f"Enabled: {enabled}\n" + f"Started at: {record.data.started_at}\n" + f"Ended at: {ended_at or '(no end time)'}\n" + f"Stateful: {stateful}" + ), + ), + ], + state=ToolResultState.SUCCESS, + ) diff --git a/src/agentscope/app/_manager/_scheduler/_tools/_schedule_delete.py b/src/agentscope/app/_manager/_scheduler/_tools/_schedule_delete.py new file mode 100644 index 0000000000000000000000000000000000000000..c09c99299b3da5b9cfc3786b8b8c58c4e93a8f08 --- /dev/null +++ b/src/agentscope/app/_manager/_scheduler/_tools/_schedule_delete.py @@ -0,0 +1,154 @@ +# -*- coding: utf-8 -*- +"""Schedule delete tool – removes a job from the scheduler and storage.""" +from typing import Any + +from pydantic import BaseModel, Field +from apscheduler.jobstores.base import JobLookupError + +from .....message import ToolResultState, TextBlock +from .....permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, +) +from .....tool import ToolBase, ToolChunk +from ....message_bus import MessageBus +from ....storage._base import StorageBase + + +class _ScheduleDeleteParams(BaseModel): + """The params for the schedule delete tool.""" + + schedule_id: str = Field( + description="The schedule ID to delete (permanently remove).", + ) + + +class ScheduleDelete(ToolBase): + """The schedule delete tool. + + Permanently removes the given scheduled job from APScheduler, + storage, and the message bus. Every execution session spawned by + the schedule is cancelled (if running) and has its bus state + purged. The job cannot be recovered after removal. + """ + + name: str = "ScheduleDelete" + + description: str = ( + "Permanently delete a scheduled task by its schedule ID. " + "After this call the task will no longer be executed and its record " + "will be deleted from storage." + ) + input_schema: dict = _ScheduleDeleteParams.model_json_schema() + + is_concurrency_safe: bool = False + is_read_only: bool = False + is_state_injected: bool = False + is_external_tool: bool = False + is_mcp: bool = False + mcp_name: str | None = None + + def __init__( + self, + user_id: str, + scheduler: Any, + storage: StorageBase, + message_bus: MessageBus, + ) -> None: + """Initialize the schedule delete tool. + + Args: + user_id (`str`): + The authenticated user; used to scope the storage deletion. + scheduler (`Any`): + The ``AsyncIOScheduler`` instance whose job will be removed. + storage (`StorageBase`): + The storage backend used to delete the persisted record. + message_bus (`MessageBus`): + The message bus used to cancel in-flight chat runs for + any execution session spawned by this schedule and to + purge their per-session bus state. + """ + self._user_id = user_id + self._scheduler = scheduler + self._storage = storage + self._message_bus = message_bus + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permission for the tool usage.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message=f"{self.name} is always allowed to be called.", + ) + + async def __call__( + self, + schedule_id: str, + ) -> ToolChunk: # type: ignore[override] + """Permanently delete the scheduled task with the given ID. + + Delegates the storage + bus cascade to + :meth:`SessionService.delete_schedule`, which cancels in-flight + runs for any session this schedule spawned and purges their + bus state before dropping the schedule record. The APScheduler + job is removed separately because it lives in-process and the + service layer is bus/storage-only. + + Args: + schedule_id (`str`): + The unique identifier of the schedule to delete. + + Returns: + `ToolChunk`: + A chunk describing the result of the delete operation. + """ + + # Remove from the in-memory scheduler (best-effort; may already be + # absent if the job finished naturally or the server restarted) + try: + self._scheduler.remove_job(schedule_id) + except JobLookupError: + pass + + # Local import to avoid a circular dependency between + # ``_manager`` and ``_service`` at module load. + from ...._service import SessionService # noqa: PLC0415 + + session_service = SessionService( + storage=self._storage, + message_bus=self._message_bus, + ) + deleted = await session_service.delete_schedule( + self._user_id, + schedule_id, + ) + + if not deleted: + return ToolChunk( + content=[ + TextBlock( + text=( + f"ScheduleNotFoundError: Schedule with id " + f"{schedule_id!r} not found in storage." + ), + ), + ], + state=ToolResultState.ERROR, + ) + + return ToolChunk( + content=[ + TextBlock( + text=( + f"Schedule {schedule_id!r} has been permanently " + f"deleted." + ), + ), + ], + state=ToolResultState.SUCCESS, + ) diff --git a/src/agentscope/app/_manager/_scheduler/_tools/_schedule_list.py b/src/agentscope/app/_manager/_scheduler/_tools/_schedule_list.py new file mode 100644 index 0000000000000000000000000000000000000000..6d3142414b450e5d9480cee38fc6dcba13a8e671 --- /dev/null +++ b/src/agentscope/app/_manager/_scheduler/_tools/_schedule_list.py @@ -0,0 +1,115 @@ +# -*- coding: utf-8 -*- +"""The tool to list the scheduled jobs in the cron scheduler manager.""" +from typing import Any + +from pydantic import BaseModel + +from .....message import ToolResultState, TextBlock +from .....permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, +) +from .....tool import ToolBase, ToolChunk +from ....storage import StorageBase + + +class _ScheduleListParams(BaseModel): + """The params for the schedule list tool.""" + + +class ScheduleList(ToolBase): + """The schedule list tool. + + Lists all scheduled tasks owned by the current user. Each entry is + fetched from storage (rich :class:`ScheduleData`) and augmented with + ``next_run_time`` from the in-memory APScheduler job when available. + """ + + name: str = "ScheduleList" + + description: str = ( + "List all scheduled tasks for the current user. " + "Shows schedule ID, name, cron expression, timezone, next run time, " + "enabled/disabled status, and whether the schedule is stateful." + ) + input_schema: dict = _ScheduleListParams.model_json_schema() + + is_concurrency_safe: bool = True + is_read_only: bool = True + is_state_injected: bool = False + is_external_tool: bool = False + is_mcp: bool = False + mcp_name: str | None = None + + def __init__( + self, + user_id: str, + scheduler: Any, + storage: StorageBase, + ) -> None: + """Initialize the schedule list tool. + + Args: + user_id (`str`): + The authenticated user; used to scope the storage lookup. + scheduler (`Any`): + The ``AsyncIOScheduler`` instance for reading ``next_run_time`` + storage (`StorageBase`): + The storage backend that holds the persisted schedule records. + """ + self._user_id = user_id + self._scheduler = scheduler + self._storage = storage + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permission for the tool usage.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message=f"{self.name} is always allowed to be called.", + ) + + async def __call__(self) -> ToolChunk: # type: ignore[override] + """List all scheduled tasks for the current user. + + Returns: + `ToolChunk`: + A chunk containing a formatted list of all scheduled tasks, + or a message indicating none exist. + """ + records = await self._storage.list_schedules(self._user_id) + + if not records: + return ToolChunk( + content=[TextBlock(text="No scheduled tasks found.")], + state=ToolResultState.SUCCESS, + ) + + # Build a map of schedule_id -> next_run_time from the live scheduler + next_run_map: dict[str, str] = { + job.id: str(job.next_run_time) + for job in self._scheduler.get_jobs() + } + + lines: list[str] = [f"Found {len(records)} scheduled task(s):\n"] + for record in records: + enabled_str = "enabled" if record.data.enabled else "disabled" + next_run = next_run_map.get(record.id, "not in scheduler") + lines.append( + f"- [{enabled_str}] {record.data.name!r} (ID: {record.id})\n" + f" Cron: {record.data.cron_expression}" + f" ({record.data.timezone})\n" + f" Next run: {next_run}\n" + f" Stateful: {record.data.stateful}" + f" | Agent: {record.agent_id}\n" + f" Source: {record.data.source.value}\n", + ) + + return ToolChunk( + content=[TextBlock(text="\n".join(lines))], + state=ToolResultState.SUCCESS, + ) diff --git a/src/agentscope/app/_manager/_scheduler/_tools/_schedule_view.py b/src/agentscope/app/_manager/_scheduler/_tools/_schedule_view.py new file mode 100644 index 0000000000000000000000000000000000000000..a37bd63ed1f8d6ad74f86d8799cb2a6c37c72d3b --- /dev/null +++ b/src/agentscope/app/_manager/_scheduler/_tools/_schedule_view.py @@ -0,0 +1,138 @@ +# -*- coding: utf-8 -*- +"""The schedule view tool.""" +from typing import Any + +from pydantic import BaseModel, Field + +from .....message import ToolResultState, TextBlock +from .....permission import ( + PermissionContext, + PermissionDecision, + PermissionBehavior, +) +from .....tool import ToolBase, ToolChunk +from ....storage import StorageBase + + +class _ScheduleViewParams(BaseModel): + """The params for the schedule view tool.""" + + schedule_id: str = Field( + description="The schedule ID.", + ) + + +class ScheduleView(ToolBase): + """The schedule view tool. + + Fetches the persisted :class:`ScheduleRecord` from storage and enriches + it with the ``next_run_time`` from the in-memory APScheduler job. + """ + + name: str = "ScheduleView" + + description: str = ( + "View the full details of a scheduled task by its schedule ID, " + "including cron expression, timezone, stateful flag, permission " + "mode, and the next scheduled run time." + ) + input_schema: dict = _ScheduleViewParams.model_json_schema() + + is_concurrency_safe: bool = True + is_read_only: bool = True + is_state_injected: bool = False + is_external_tool: bool = False + is_mcp: bool = False + mcp_name: str | None = None + + def __init__( + self, + user_id: str, + scheduler: Any, + storage: StorageBase, + ) -> None: + """Initialize the schedule view tool. + + Args: + user_id (`str`): + The authenticated user; used to scope the storage lookup. + scheduler (`Any`): + The ``AsyncIOScheduler`` instance for + reading ``next_run_time``. + storage (`StorageBase`): + The storage backend that holds the persisted schedule records. + """ + self._user_id = user_id + self._scheduler = scheduler + self._storage = storage + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Check permission for the tool usage.""" + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message=f"{self.name} is always allowed to be called.", + ) + + async def __call__( + self, + schedule_id: str, + ) -> ToolChunk: # type: ignore[override] + """View the full details of a scheduled task. + + Args: + schedule_id (`str`): + The unique identifier of the schedule to view. + + Returns: + `ToolChunk`: + A chunk containing the formatted schedule details, or an + error description if the schedule is not found. + """ + record = await self._storage.get_schedule(self._user_id, schedule_id) + + if record is None: + return ToolChunk( + content=[ + TextBlock( + text=( + f"ScheduleNotFoundError: Schedule with " + f"id {schedule_id!r} not found." + ), + ), + ], + state=ToolResultState.ERROR, + ) + + job = self._scheduler.get_job(schedule_id) + next_run = ( + str(job.next_run_time) + if job is not None + else "not in scheduler (may be disabled)" + ) + enabled_str = "enabled" if record.data.enabled else "disabled" + + text = ( + f"Schedule ID: {record.id}\n" + f"Name: {record.data.name}\n" + f"Description: {record.data.description or '(none)'}\n" + f"Status: {enabled_str}\n" + f"Cron: {record.data.cron_expression}" + f" (timezone: {record.data.timezone})\n" + f"Next run: {next_run}\n" + f"Stateful: {record.data.stateful}\n" + f"Permission mode: {record.data.permission_mode.value}\n" + f"Source: {record.data.source.value}\n" + f"Source session: {record.data.source_session_id or '(none)'}\n" + f"Agent ID: {record.agent_id}\n" + f"Created at: {record.created_at}\n" + f"Updated at: {record.updated_at}\n" + ) + + return ToolChunk( + content=[TextBlock(text=text)], + state=ToolResultState.SUCCESS, + ) diff --git a/src/agentscope/app/_manager/_wakeup_dispatcher.py b/src/agentscope/app/_manager/_wakeup_dispatcher.py new file mode 100644 index 0000000000000000000000000000000000000000..4f7de1f6e4d64f92d739ff245ded8c1c3d39c162 --- /dev/null +++ b/src/agentscope/app/_manager/_wakeup_dispatcher.py @@ -0,0 +1,371 @@ +# -*- coding: utf-8 -*- +"""Single per-process dispatcher for all cross-session run triggers. + +One asyncio task per process. Subscribes to the shared trigger signal +channel and drains the durable trigger queue on each signal. It is the +**sole** site that spawns :meth:`ChatService.run` into the shared +:class:`ChatRunRegistry`, which is what makes concurrent-spawn races +(two writers contending for one session's run slot → a spurious "already +has an active chat run" 409) structurally impossible: every run trigger +funnels through this one serial consumer. + +Each queue entry carries a ``kind`` that selects how a busy session is +handled: + +- ``wake`` (idle-session wake-up, ``input_msg=None``): skipped while the + session is already running — the live run will drain the inbox. +- ``resume`` (a parked HITL run being fed its result): must *not* be + skipped while running, because the session is typically still running + the parked tail at trigger time. It is re-queued after a short backoff + until the parked run releases its session lock, then spawned with the + carried input event. + +All bus keys live on the :class:`MessageBus` base class (see +``enqueue_wakeup`` / ``enqueue_input``, ``dequeue_wakeups``, +``subscribe_wakeup_signal``, ``session_is_running``), so this file has +no hard-coded key strings. +""" +import asyncio +from typing import TYPE_CHECKING, Self + +from pydantic import TypeAdapter + +from ..._logging import logger +from ...event import UserConfirmResultEvent, ExternalExecutionResultEvent +from ..message_bus import MessageBusKeys +from .._bus_ops import enqueue_run_trigger + +if TYPE_CHECKING: + from ..message_bus import MessageBus + from ..storage import StorageBase + from .._service import ChatService + from ._chat_run_registry import ChatRunRegistry + +# Parses a queued ``resume`` input dict back into its concrete event, +# discriminated by the ``type`` field shared by both result events. +_RESUME_INPUT_ADAPTER: TypeAdapter = TypeAdapter( + UserConfirmResultEvent | ExternalExecutionResultEvent, +) + +# Delay before re-queuing a ``resume`` trigger whose target session is +# still running (the parked run is finishing and about to free its +# lock). Short enough to feel instant to the user, long enough to avoid +# a hot re-enqueue loop while the lock is held. +_RESUME_RETRY_BACKOFF_SECS = 0.1 + + +class WakeupDispatcher: + """One asyncio task per process, draining the shared trigger queue. + + Args: + message_bus (`MessageBus`): + Application message bus. Used for signal subscription, + queue drain, ``session_is_running`` checks, and re-queuing + deferred ``resume`` triggers. + storage (`StorageBase`): + Persistent storage backend. Consulted before spawning a + run so triggers whose target session has been deleted are + dropped instead of crashing :class:`ChatService.run`. + chat_service (`ChatService`): + Drives the actual chat run when a trigger fires. + chat_run_registry (`ChatRunRegistry`): + Per-process registry that holds the spawned task handle so + it can be located by :class:`CancelDispatcher`. + """ + + def __init__( + self, + message_bus: "MessageBus", + storage: "StorageBase", + chat_service: "ChatService", + chat_run_registry: "ChatRunRegistry", + ) -> None: + """Bind dependencies. + + Args: + message_bus (`MessageBus`): + Application message bus. + storage (`StorageBase`): + Persistent storage backend. + chat_service (`ChatService`): + Drives session runs via :meth:`ChatService.run`. + chat_run_registry (`ChatRunRegistry`): + Shared chat-run registry to spawn into. + """ + self._bus = message_bus + self._storage = storage + self._chat_service = chat_service + self._registry = chat_run_registry + self._task: asyncio.Task | None = None + # Detached backoff timers for deferred ``resume`` re-enqueues. + # Held so they are not garbage-collected mid-sleep and can be + # cancelled on shutdown. + self._retry_tasks: set[asyncio.Task] = set() + + async def __aenter__(self) -> Self: + """Start the dispatcher loop and wait until its bus + subscription is live. + + Also performs an initial drain right after subscription so + triggers produced while this process was down (durable in + the queue) are picked up immediately on startup. + + Returns: + `Self`: This dispatcher instance. + """ + ready = asyncio.Event() + self._task = asyncio.create_task( + self._loop(ready), + name="wakeup-dispatcher", + ) + await ready.wait() + await self._drain_and_dispatch() + return self + + async def __aexit__(self, *exc: object) -> None: + """Cancel the dispatcher loop and any pending retries.""" + retries = list(self._retry_tasks) + for retry in retries: + retry.cancel() + for retry in retries: + try: + await retry + except asyncio.CancelledError: + pass + self._retry_tasks.clear() + if self._task is None: + return + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + self._task = None + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + async def _loop(self, ready: asyncio.Event) -> None: + """Long-lived loop: subscribe to the signal channel and drain + the queue on every received signal. + + Args: + ready (`asyncio.Event`): + Signalled after the underlying SUBSCRIBE completes. + :meth:`start` blocks on this so callers can publish a + trigger immediately after start without racing. + """ + try: + async for _signal in self._bus.subscribe( + MessageBusKeys.wakeup_signal(), + on_ready=ready.set, + ): + await self._drain_and_dispatch() + except Exception: # pylint: disable=broad-except + logger.exception( + "WakeupDispatcher loop crashed; subscription ended.", + ) + + async def _drain_and_dispatch(self) -> None: + """Read up to a batch of trigger entries and dispatch each.""" + try: + raw_entries = await self._bus.queue_drain( + MessageBusKeys.wakeup_queue(), + max_count=64, + ) + entries = [payload for _entry_id, payload in raw_entries] + except Exception: # pylint: disable=broad-except + logger.exception("WakeupDispatcher: dequeue_wakeups failed.") + return + + for payload in entries: + try: + user_id = payload["user_id"] + session_id = payload["session_id"] + agent_id = payload["agent_id"] + except (KeyError, TypeError): + logger.warning( + "WakeupDispatcher: skipping malformed trigger entry %r", + payload, + ) + continue + # Entries from older producers omit ``kind`` — treat as wake. + kind = payload.get("kind", MessageBusKeys.WAKEUP_KIND_WAKE) + await self._dispatch_one( + user_id=user_id, + session_id=session_id, + agent_id=agent_id, + kind=kind, + raw_input=payload.get("input"), + ) + + async def _dispatch_one( + self, + user_id: str, + session_id: str, + agent_id: str, + kind: str, + raw_input: dict | None, + ) -> None: + """Dispatch a single trigger entry by its ``kind``. + + Args: + user_id (`str`): + The owning user id. + session_id (`str`): + The session to trigger. + agent_id (`str`): + The agent that owns the session. + kind (`str`): + Trigger kind (``wake`` / ``resume``); see module docstring. + raw_input (`dict | None`): + Serialised input event for ``resume`` triggers, else + ``None``. + """ + is_resume = kind == MessageBusKeys.WAKEUP_KIND_RESUME + + # Parse the resume input early so every downstream path + # (lock-retry, spawn-retry) receives a typed event object + # rather than a raw dict. + input_msg: UserConfirmResultEvent | ExternalExecutionResultEvent | None + input_msg = None + if is_resume: + if raw_input is None: + logger.warning( + "WakeupDispatcher: dropping resume trigger for session " + "%s — no input event carried.", + session_id, + ) + return + try: + input_msg = _RESUME_INPUT_ADAPTER.validate_python(raw_input) + except Exception: # pylint: disable=broad-except + logger.exception( + "WakeupDispatcher: dropping resume trigger for session " + "%s — input event failed to parse: %r", + session_id, + raw_input, + ) + return + + if await self._bus.is_locked( + MessageBusKeys.session_lock(session_id), + ): + if is_resume: + # The session is busy finishing its parked tail. Do NOT + # drop the resume — re-queue it after a short backoff so + # it lands once the parked run releases its lock. + self._schedule_resume_retry( + user_id, + session_id, + agent_id, + input_msg, + ) + # ``wake`` triggers are safe to drop while running — the + # live run drains the inbox itself. + return + + # Orphan guard: the queue is unaware of session lifecycle. A + # trigger enqueued before the session was deleted (e.g. by a + # BG-task completion callback or a schedule trigger) will still + # arrive here. Drop it rather than letting ChatService.run crash + # on a missing storage record. + if ( + await self._storage.get_session(user_id, agent_id, session_id) + is None + ): + logger.warning( + "WakeupDispatcher: dropping %s trigger for session %s " + "(agent %s, user %s) — session no longer exists in " + "storage; it was likely enqueued before the session was " + "deleted.", + kind, + session_id, + agent_id, + user_id, + ) + return + + try: + self._registry.spawn( + self._chat_service.run( + user_id=user_id, + session_id=session_id, + agent_id=agent_id, + input_msg=input_msg, + ), + session_id=session_id, + name=f"{kind}-run:{session_id}", + ) + except RuntimeError: + # A local run was registered between the running-check and + # the spawn. For ``wake`` that run will drain the inbox; for + # ``resume`` re-queue so the result is not lost. + if is_resume: + self._schedule_resume_retry( + user_id, + session_id, + agent_id, + input_msg, + ) + else: + logger.debug( + "WakeupDispatcher: skipping wake trigger for session " + "%s; a local run is already registered.", + session_id, + ) + + def _schedule_resume_retry( + self, + user_id: str, + session_id: str, + agent_id: str, + input_msg: UserConfirmResultEvent + | ExternalExecutionResultEvent + | None, + ) -> None: + """Re-enqueue a ``resume`` trigger after a short backoff. + + Spawns a detached timer that sleeps, then re-enqueues the resume + (which re-fires the signal, re-driving the drain). This keeps the + resume alive across the window where the parked run still holds + the session lock, without a hot re-enqueue loop. + + Args: + user_id (`str`): + The owning user id. + session_id (`str`): + The session to resume. + agent_id (`str`): + The agent that owns the session. + input_msg: + The parsed input event to redeliver. + """ + + async def _retry() -> None: + try: + await asyncio.sleep(_RESUME_RETRY_BACKOFF_SECS) + await enqueue_run_trigger( + self._bus, + user_id=user_id, + session_id=session_id, + agent_id=agent_id, + kind=MessageBusKeys.WAKEUP_KIND_RESUME, + inputs=input_msg, + ) + except asyncio.CancelledError: + pass + except Exception: # pylint: disable=broad-except + logger.exception( + "WakeupDispatcher: failed to re-enqueue resume trigger " + "for session %s.", + session_id, + ) + + task = asyncio.create_task( + _retry(), + name=f"resume-retry:{session_id}", + ) + self._retry_tasks.add(task) + task.add_done_callback(self._retry_tasks.discard) diff --git a/src/agentscope/app/_router/__init__.py b/src/agentscope/app/_router/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..22e7765f393952db42dec887140fd0b3744cbca7 --- /dev/null +++ b/src/agentscope/app/_router/__init__.py @@ -0,0 +1,23 @@ +# -*- coding: utf-8 -*- +"""App routers.""" +from ._agent import agent_router +from ._chat import chat_router +from ._credential import credential_router +from ._knowledge_base import knowledge_base_router +from ._schedule import schedule_router +from ._session import session_router +from ._model import model_router +from ._tts_model import tts_model_router +from ._workspace import workspace_router + +__all__ = [ + "agent_router", + "model_router", + "tts_model_router", + "chat_router", + "credential_router", + "knowledge_base_router", + "schedule_router", + "session_router", + "workspace_router", +] diff --git a/src/agentscope/app/_router/_agent.py b/src/agentscope/app/_router/_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..9d03666b534e2f59aa2e141d91d62044b6d6eb89 --- /dev/null +++ b/src/agentscope/app/_router/_agent.py @@ -0,0 +1,214 @@ +# -*- coding: utf-8 -*- +"""Agent router — CRUD endpoints for agent configurations.""" +from datetime import datetime + +from fastapi import APIRouter, Depends, HTTPException, status + +from ...agent import ContextConfig, ReActConfig +from ..deps import get_current_user_id, get_session_service, get_storage +from ._schema import ( + AgentSchemaResponse, + ListAgentsResponse, + CreateAgentRequest, + CreateAgentResponse, + UpdateAgentRequest, +) +from .._service import SessionService +from ..storage import StorageBase, AgentData, AgentRecord + +agent_router = APIRouter( + prefix="/agent", + tags=["agent"], + responses={404: {"description": "Not found"}}, +) + + +@agent_router.get( + "/schema", + response_model=AgentSchemaResponse, + summary="Get JSON Schema fragments for the agent form", +) +async def get_agent_schema() -> AgentSchemaResponse: + """Return the JSON Schema fragments used by the frontend to render + the agent create / edit forms. + + The frontend uses three sections — identity, context config, and + react config — so we return them as separate self-contained schemas + rather than a single ``AgentData`` schema with ``$ref``s. + + Returns: + `AgentSchemaResponse`: + Schemas for the three form sections. + """ + # Slice ``AgentData``'s schema down to the identity-relevant fields. + # Going through ``AgentData.model_json_schema()`` (rather than building + # a dict by hand) keeps Pydantic as the single source of truth for + # defaults, titles, descriptions, and the ``format: textarea`` hint. + agent_schema = AgentData.model_json_schema() + identity_keys = ("name", "system_prompt") + identity = { + "type": "object", + "title": "Identity", + "properties": { + k: v + for k, v in agent_schema.get("properties", {}).items() + if k in identity_keys + }, + "required": [ + r for r in agent_schema.get("required", []) if r in identity_keys + ], + } + + context_schema = ContextConfig.model_json_schema() + # ``summary_schema`` holds a Pydantic JSON Schema describing how the + # compression model should structure its output. The end-user is not + # expected to edit it from the form, so we hide it. + context_schema.get("properties", {}).pop("summary_schema", None) + + return AgentSchemaResponse( + identity=identity, + context_config=context_schema, + react_config=ReActConfig.model_json_schema(), + ) + + +@agent_router.get( + "/", + response_model=ListAgentsResponse, + summary="List all agents", +) +async def list_agents( + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), +) -> ListAgentsResponse: + """Return all agent records belonging to the authenticated user. + + Args: + user_id (`str`): + Injected authenticated user ID. + storage (`StorageBase`): + Injected storage backend. + + Returns: + `ListAgentsResponse`: + All agent records and their total count. + """ + agents = await storage.list_agents(user_id) + return ListAgentsResponse(agents=agents, total=len(agents)) + + +@agent_router.post( + "/", + response_model=CreateAgentResponse, + status_code=status.HTTP_201_CREATED, + summary="Create a new agent", +) +async def create_agent( + body: CreateAgentRequest, + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), +) -> CreateAgentResponse: + """Create and persist a new agent configuration. + + Args: + body (`CreateAgentRequest`): + Agent configuration to store. + user_id (`str`): + Injected authenticated user ID. + storage (`StorageBase`): + Injected storage backend. + + Returns: + `CreateAgentResponse`: + The server-assigned agent identifier. + """ + record = AgentRecord( + user_id=user_id, + data=AgentData( + name=body.name, + system_prompt=body.system_prompt, + context_config=body.context_config, + react_config=body.react_config, + ), + ) + agent_id = await storage.upsert_agent(user_id, record) + return CreateAgentResponse(agent_id=agent_id) + + +@agent_router.patch( + "/{agent_id}", + response_model=AgentRecord, + summary="Update an agent", +) +async def update_agent( + agent_id: str, + body: UpdateAgentRequest, + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), +) -> AgentRecord: + """Partially update an existing agent configuration. + + Only the fields present in the request body are updated; all other fields + keep their current values. + + Args: + agent_id (`str`): The agent to update. + body (`UpdateAgentRequest`): Fields to update. + user_id (`str`): Injected authenticated user ID. + storage (`StorageBase`): Injected storage backend. + + Returns: + `AgentRecord`: The full agent record after the update. + + Raises: + `HTTPException`: 404 if the agent does not exist or does not belong + to the authenticated user. + """ + agents = await storage.list_agents(user_id) + existing = next((a for a in agents if a.id == agent_id), None) + if existing is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Agent '{agent_id}' not found.", + ) + + updates = body.model_dump(exclude_none=True) + updated_data = existing.data.model_copy(update=updates) + updated_agent = existing.model_copy( + update={"data": updated_data, "updated_at": datetime.now()}, + ) + await storage.upsert_agent(user_id, updated_agent) + return updated_agent + + +@agent_router.delete( + "/{agent_id}", + status_code=status.HTTP_204_NO_CONTENT, + summary="Delete an agent", +) +async def delete_agent( + agent_id: str, + user_id: str = Depends(get_current_user_id), + session_service: SessionService = Depends(get_session_service), +) -> None: + """Permanently delete an agent configuration. + + Cascades through every session owned by this agent (and, for team + leaders, through every worker session) — cancelling any in-flight + chat run, removing storage records, and purging bus state. + + Args: + agent_id (`str`): The agent to delete. + user_id (`str`): Injected authenticated user ID. + session_service (`SessionService`): Injected session service. + + Raises: + `HTTPException`: 404 if the agent does not exist or does not belong + to the authenticated user. + """ + deleted = await session_service.delete_agent(user_id, agent_id) + if not deleted: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Agent '{agent_id}' not found.", + ) diff --git a/src/agentscope/app/_router/_chat.py b/src/agentscope/app/_router/_chat.py new file mode 100644 index 0000000000000000000000000000000000000000..722fa36cf4a7eea62caa1df9562d405716e3ab09 --- /dev/null +++ b/src/agentscope/app/_router/_chat.py @@ -0,0 +1,160 @@ +# -*- coding: utf-8 -*- +"""Chat router — fire-and-forget trigger for chat runs. + +The endpoint no longer returns an SSE stream. Instead, it kicks off a +chat run as a background task and returns immediately. Events produced +by the run are published to the message bus and delivered to the +frontend via the long-lived ``GET /sessions/{sid}/stream`` SSE +connection provided by the session router. + +Two trigger paths, deliberately asymmetric: + +- **New user message(s)** are spawned directly into the + :class:`ChatRunRegistry`. The registry's single-run-per-session rule + surfaces as a 409, which is exactly the desired double-submit guard. +- **HITL results** (``UserConfirmResultEvent`` / + ``ExternalExecutionResultEvent``) are *enqueued* onto the shared + run-trigger queue and drained by the single + :class:`WakeupDispatcher`. Routing the resume through the queue keeps + the dispatcher the sole spawn site, so a resume can never collide with + the worker's still-finishing parked run (the old 409 race) — the + dispatcher serialises them. +""" +from fastapi import APIRouter, Depends, HTTPException, status + +from ..deps import ( + get_chat_run_registry, + get_chat_service, + get_current_user_id, + get_message_bus, +) +from ._schema import ChatRequest, ChatTriggerResponse +from .._manager import ChatRunRegistry +from .._service import ( + ChatService, + SessionProjection, + SubagentHitlProjector, +) +from ..message_bus import MessageBus, MessageBusKeys +from .._bus_ops import enqueue_run_trigger +from ...event import UserConfirmResultEvent, ExternalExecutionResultEvent + +chat_router = APIRouter( + prefix="/chat", + tags=["chat"], + responses={404: {"description": "Not found"}}, +) + + +@chat_router.post( + "/", + response_model=ChatTriggerResponse, + summary="Trigger a chat run (fire-and-forget)", +) +async def chat( + request: ChatRequest, + user_id: str = Depends(get_current_user_id), + chat_service: ChatService = Depends(get_chat_service), + chat_run_registry: ChatRunRegistry = Depends(get_chat_run_registry), + message_bus: MessageBus = Depends(get_message_bus), +) -> ChatTriggerResponse: + """Trigger a chat run for the specified session. + + Events produced during the run are published to the message bus and + delivered to any active ``GET /sessions/{session_id}/stream`` SSE + subscriber. The caller does **not** receive events from this + endpoint's response body. + + Accepts the same ``input`` payloads as before: + + - ``Msg`` / ``list[Msg]``: new user message(s) — spawned directly. + - ``UserConfirmResultEvent`` / ``ExternalExecutionResultEvent``: + resume a paused tool call (human-in-the-loop) — routed to the + owning session and enqueued for the dispatcher. + - ``None``: continue from current state — spawned directly. + + Args: + request (`ChatRequest`): + JSON body with ``agent_id``, ``session_id``, and ``input``. + user_id (`str`): + Injected user id. + chat_service (`ChatService`): + Injected application-wide chat service. + chat_run_registry (`ChatRunRegistry`): + Injected per-process chat-run registry. + message_bus (`MessageBus`): + Injected message bus, used to resolve subagent-confirm + routing and to enqueue resume triggers. + + Returns: + `ChatTriggerResponse`: + Confirms the run was scheduled (for a resume, that it was + enqueued). + + Raises: + `HTTPException`: + 409 if a chat run for this session is already in flight in + this process (the registry enforces single-run-per-session). + Only direct-spawn paths (new messages / ``None``) can raise + this; the enqueued resume path never does. + """ + # ------------------------------------------------------------------ + # HITL resume — route to the owning session, then enqueue. + # + # A confirmation / external-result POSTed to a *leader* session may + # actually belong to a team *member*: the leader is the single front + # door clients talk to. Resolve the owning worker HERE, then enqueue + # a ``resume`` trigger for that session. The single WakeupDispatcher + # drains it — spawning under the *worker* session id, serialised + # behind any still-finishing parked run, so there is no registry + # collision (no 409) and the leader's run slot is never occupied by + # the worker's resume. + # ------------------------------------------------------------------ + if isinstance( + request.input, + (UserConfirmResultEvent, ExternalExecutionResultEvent), + ): + run_session_id = request.session_id + run_agent_id = request.agent_id + target = await SubagentHitlProjector.resolve( + SessionProjection(message_bus), + request.session_id, + request.input.reply_id, + ) + if target is not None: + run_session_id = target["worker_session_id"] + run_agent_id = target["worker_agent_id"] + + await enqueue_run_trigger( + message_bus, + user_id=user_id, + session_id=run_session_id, + agent_id=run_agent_id, + kind=MessageBusKeys.WAKEUP_KIND_RESUME, + inputs=request.input, + ) + return ChatTriggerResponse(status="started", session_id=run_session_id) + + # ------------------------------------------------------------------ + # New user message(s) / None — spawn directly. The registry's + # single-run-per-session rule is the desired double-submit guard. + # ------------------------------------------------------------------ + try: + chat_run_registry.spawn( + chat_service.run( + user_id=user_id, + session_id=request.session_id, + agent_id=request.agent_id, + input_msg=request.input, + ), + session_id=request.session_id, + ) + except RuntimeError as e: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=str(e), + ) from e + return ChatTriggerResponse( + status="started", + session_id=request.session_id, + ) diff --git a/src/agentscope/app/_router/_credential.py b/src/agentscope/app/_router/_credential.py new file mode 100644 index 0000000000000000000000000000000000000000..c3c21bfd0ee3eda517a5f3ee5441a29615dab761 --- /dev/null +++ b/src/agentscope/app/_router/_credential.py @@ -0,0 +1,164 @@ +# -*- coding: utf-8 -*- +"""Credential router — CRUD endpoints for API key credentials.""" +from fastapi import APIRouter, Depends, HTTPException, status + +from ..deps import get_current_user_id, get_storage +from ._schema import ( + CreateCredentialRequest, + CreateCredentialResponse, + ListCredentialsResponse, + ListCredentialSchemasResponse, + UpdateCredentialRequest, +) +from ..storage import StorageBase, CredentialRecord +from ...credential import CredentialFactory + +credential_router = APIRouter( + prefix="/credential", + tags=["credential"], + responses={404: {"description": "Not found"}}, +) + + +@credential_router.get( + "/schemas", + response_model=ListCredentialSchemasResponse, + summary="List JSON schemas for all credential types", +) +async def list_credential_schemas() -> ListCredentialSchemasResponse: + """Return JSON schemas for all registered credential types. + + Used by the frontend to render credential creation forms dynamically. + """ + + return ListCredentialSchemasResponse( + schemas=CredentialFactory.list_schemas(), + ) + + +@credential_router.get( + "/", + response_model=ListCredentialsResponse, + summary="List all credentials", +) +async def list_credentials( + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), +) -> ListCredentialsResponse: + """Return all credential records belonging to the authenticated user. + + Args: + user_id (`str`): + Injected authenticated user ID. + storage (`StorageBase`): + Injected storage backend. + + Returns: + `ListCredentialsResponse`: + All credential records and their total count. + """ + credentials = await storage.list_credentials(user_id) + return ListCredentialsResponse( + credentials=credentials, + total=len(credentials), + ) + + +@credential_router.post( + "/", + response_model=CreateCredentialResponse, + status_code=status.HTTP_201_CREATED, + summary="Create a new credential", +) +async def create_credential( + body: CreateCredentialRequest, + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), +) -> CreateCredentialResponse: + """Store a new credential. + + Args: + body (`CreateCredentialRequest`): Credential payload to store. + user_id (`str`): Injected authenticated user ID. + storage (`StorageBase`): Injected storage backend. + + Returns: + `CreateCredentialResponse`: The server-assigned credential identifier. + """ + credential_id = await storage.upsert_credential( + user_id, + CredentialFactory.from_dict(body.data), + ) + return CreateCredentialResponse(credential_id=credential_id) + + +@credential_router.patch( + "/{credential_id}", + response_model=CredentialRecord, + summary="Update a credential", +) +async def update_credential( + credential_id: str, + body: UpdateCredentialRequest, + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), +) -> CredentialRecord: + """Replace the payload of an existing credential. + + Args: + credential_id (`str`): The credential to update. + body (`UpdateCredentialRequest`): New credential payload. + user_id (`str`): Injected authenticated user ID. + storage (`StorageBase`): Injected storage backend. + + Returns: + `CredentialRecord`: The updated credential record. + + Raises: + `HTTPException`: 404 if the credential does not exist or does not + belong to the authenticated user. + """ + credentials = await storage.list_credentials(user_id) + existing = next((c for c in credentials if c.id == credential_id), None) + if existing is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Credential '{credential_id}' not found.", + ) + + credential = CredentialFactory.from_dict(body.data) + credential.id = credential_id + await storage.upsert_credential(user_id, credential) + # Re-fetch to return the persisted record with updated timestamps. + credentials = await storage.list_credentials(user_id) + updated = next(c for c in credentials if c.id == credential_id) + return updated + + +@credential_router.delete( + "/{credential_id}", + status_code=status.HTTP_204_NO_CONTENT, + summary="Delete a credential", +) +async def delete_credential( + credential_id: str, + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), +) -> None: + """Permanently delete a credential. + + Args: + credential_id (`str`): The credential to delete. + user_id (`str`): Injected authenticated user ID. + storage (`StorageBase`): Injected storage backend. + + Raises: + `HTTPException`: 404 if the credential does not exist or does not + belong to the authenticated user. + """ + deleted = await storage.delete_credential(user_id, credential_id) + if not deleted: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Credential '{credential_id}' not found.", + ) diff --git a/src/agentscope/app/_router/_knowledge_base.py b/src/agentscope/app/_router/_knowledge_base.py new file mode 100644 index 0000000000000000000000000000000000000000..3b8a45bef4f2ca31e5240392f91a2060608ede95 --- /dev/null +++ b/src/agentscope/app/_router/_knowledge_base.py @@ -0,0 +1,574 @@ +# -*- coding: utf-8 -*- +"""Knowledge base router — manage knowledge bases and their documents. + +A knowledge base is the user-facing concept; physically each one maps +to a single vector store collection (in the MVP isolation strategy). +The HTTP layer is intentionally thin — every endpoint translates the +request into a single :class:`~agentscope.app._service. +KnowledgeBaseService` call and returns the result. +""" +from fastapi import ( + APIRouter, + Depends, + File, + Form, + Path, + Query, + UploadFile, + status, +) + +from ..deps import ( + get_current_user_id, + get_knowledge_base_manager, + get_knowledge_base_service, + get_knowledge_parsers, + get_storage, +) +from ._schema import ( + CreateKnowledgeBaseRequest, + CreateKnowledgeBaseResponse, + KbEmbeddingProvider, + KbMiddlewareParametersSchemaResponse, + KnowledgeBaseView, + KnowledgeDocumentView, + ListKbEmbeddingModelsResponse, + ListKnowledgeBasesResponse, + ListKnowledgeDocumentsResponse, + ListKnowledgeDocumentStatusResponse, + ListSupportedContentTypesResponse, + SearchKnowledgeBaseRequest, + SearchKnowledgeBaseResponse, + UpdateKnowledgeBaseRequest, + UploadKnowledgeDocumentResponse, +) +from ...credential import CredentialFactory +from ..rag.knowledge_base_manager import KnowledgeBaseManagerBase +from ..storage import StorageBase +from .._service import KnowledgeBaseService +from ...middleware import RAGMiddleware +from ...rag import ParserBase + + +knowledge_base_router = APIRouter( + prefix="/knowledge_bases", + tags=["knowledge_bases"], + responses={404: {"description": "Not found"}}, +) + + +@knowledge_base_router.get( + "/embedding_models", + response_model=ListKbEmbeddingModelsResponse, + summary="List embedding models compatible with the KB dimension policy", +) +async def list_kb_embedding_models( + user_id: str = Depends(get_current_user_id), + storage: "StorageBase" = Depends(get_storage), + manager: "KnowledgeBaseManagerBase" = Depends( + get_knowledge_base_manager, + ), +) -> ListKbEmbeddingModelsResponse: + """List embedding models the user can pick at KB-creation time. + + Walks the caller's credentials, looks up each provider's + embedding model class, gathers its model cards, and projects + each card through the manager's :class:`DimensionPolicy`. + Incompatible cards are dropped; matryoshka cards under a + ``FIXED`` / ``LOCKED_BY_EXISTING`` policy are narrowed to the + locked dimension. Providers that end up with zero compatible + models are omitted from the response entirely. + + Args: + user_id (`str`): + Injected authenticated user ID. + storage (`StorageBase`): + Injected storage backend used to enumerate credentials. + manager (`KnowledgeBaseManagerBase`): + Injected knowledge base manager. + + Returns: + `ListKbEmbeddingModelsResponse`: + One entry per credential with at least one compatible + embedding model, plus the policy used for filtering. + """ + policy = await manager.get_dimension_policy() + credentials = await storage.list_credentials(user_id) + + providers: list[KbEmbeddingProvider] = [] + for credential in credentials: + credential_type = credential.data.get("type") + if not credential_type: + continue + credential_cls = CredentialFactory.get_credential_class( + credential_type, + ) + if credential_cls is None: + continue + embedding_cls = credential_cls.get_embedding_model_class() + if embedding_cls is None: + continue + + filtered = [] + for card in embedding_cls.list_models(): + projected = policy.filter_card(card) + if projected is not None: + filtered.append(projected) + if not filtered: + continue + providers.append( + KbEmbeddingProvider(credential=credential, models=filtered), + ) + + return ListKbEmbeddingModelsResponse(providers=providers, policy=policy) + + +@knowledge_base_router.get( + "/middleware/parameters_schema", + response_model=KbMiddlewareParametersSchemaResponse, + summary="JSON Schema for the KB middleware's tunable parameters", +) +async def get_kb_middleware_parameters_schema( + _: str = Depends(get_current_user_id), +) -> KbMiddlewareParametersSchemaResponse: + """Return the parameter schema for + :class:`agentscope.middleware.RAGMiddleware`. + + The schema is shaped like every other ``parameter_schema`` served + by this service — title / description / default / enum / minimum + / maximum — so the front-end can render the session-level KB + attachment form with the same schema-driven component used for + model parameters. + + Args: + _ (`str`): + Injected authenticated user ID; only used to gate the + endpoint behind authentication. + + Returns: + `KbMiddlewareParametersSchemaResponse`: + The JSON Schema describing the middleware's + user-tunable parameters. + """ + return KbMiddlewareParametersSchemaResponse( + parameter_schema=(RAGMiddleware.Parameters.model_json_schema()), + ) + + +@knowledge_base_router.get( + "/supported_content_types", + response_model=ListSupportedContentTypesResponse, + summary="List file types the configured parsers can ingest", +) +async def list_supported_content_types( + _: str = Depends(get_current_user_id), + parsers: list[ParserBase] + | dict[str, ParserBase] = Depends( + get_knowledge_parsers, + ), +) -> ListSupportedContentTypesResponse: + """Advertise the union of media types and filename extensions every + registered parser accepts. + + Used by the front-end to populate the document picker's ``accept`` + attribute and to reject drag-dropped files whose extension lies + outside the supported set before the upload starts. Routing on + upload still goes through the media type — this endpoint is a + capability hint, not authoritative dispatch. + + Args: + _ (`str`): + Injected authenticated user ID; only used to gate the + endpoint behind authentication. + parsers (`list[ParserBase] | dict[str, ParserBase]`): + Injected parser registry — the same value the index worker + uses to dispatch uploads. + + Returns: + `ListSupportedContentTypesResponse`: + Deduplicated, sorted unions of ``media_types`` and + ``extensions``. + """ + parser_iter = parsers.values() if isinstance(parsers, dict) else parsers + media_types: set[str] = set() + extensions: set[str] = set() + for parser in parser_iter: + media_types.update(parser.supported_media_types) + extensions.update(parser.supported_extensions()) + return ListSupportedContentTypesResponse( + media_types=sorted(media_types), + extensions=sorted(extensions), + ) + + +# ---------------------------------------------------------------------- +# Knowledge base management +# ---------------------------------------------------------------------- + + +@knowledge_base_router.post( + "/", + response_model=CreateKnowledgeBaseResponse, + status_code=status.HTTP_201_CREATED, + summary="Create a new knowledge base", +) +async def create_knowledge_base( + body: CreateKnowledgeBaseRequest, + user_id: str = Depends(get_current_user_id), + service: "KnowledgeBaseService" = Depends(get_knowledge_base_service), +) -> CreateKnowledgeBaseResponse: + """Create a new knowledge base for the authenticated user. + + Allocates a fresh vector store collection sized to the embedding + model's output dimension and persists the knowledge base record. + + Args: + body (`CreateKnowledgeBaseRequest`): + Knowledge base name, description, and embedding model + configuration. + user_id (`str`): + Injected authenticated user ID. + service (`KnowledgeBaseService`): + Injected knowledge base service. + + Returns: + `CreateKnowledgeBaseResponse`: + The server-assigned knowledge base identifier. + """ + record = await service.create_knowledge_base( + user_id=user_id, + name=body.name, + description=body.description, + embedding_model_config=body.embedding_model_config, + ) + return CreateKnowledgeBaseResponse(knowledge_base_id=record.id) + + +@knowledge_base_router.get( + "/", + response_model=ListKnowledgeBasesResponse, + summary="List the caller's knowledge bases", +) +async def list_knowledge_bases( + user_id: str = Depends(get_current_user_id), + service: "KnowledgeBaseService" = Depends(get_knowledge_base_service), +) -> ListKnowledgeBasesResponse: + """Return all knowledge bases owned by the authenticated user. + + Args: + user_id (`str`): + Injected authenticated user ID. + service (`KnowledgeBaseService`): + Injected knowledge base service. + + Returns: + `ListKnowledgeBasesResponse`: + The user's knowledge bases. + """ + records = await service.list_knowledge_bases(user_id) + views = [ + KnowledgeBaseView( + id=record.id, + name=record.name, + description=record.description, + embedding_model_config=record.embedding_model_config, + created_at=record.created_at, + updated_at=record.updated_at, + ) + for record in records + ] + return ListKnowledgeBasesResponse(knowledge_bases=views, total=len(views)) + + +@knowledge_base_router.patch( + "/{knowledge_base_id}", + response_model=KnowledgeBaseView, + summary="Update mutable fields on a knowledge base", +) +async def update_knowledge_base( + body: UpdateKnowledgeBaseRequest, + knowledge_base_id: str = Path(description="The knowledge base id."), + user_id: str = Depends(get_current_user_id), + service: "KnowledgeBaseService" = Depends(get_knowledge_base_service), +) -> KnowledgeBaseView: + """Update mutable fields on a knowledge base. + + Only ``name`` and ``description`` can be updated. The embedding + model configuration is pinned at creation time and cannot be + changed. + + Args: + body (`UpdateKnowledgeBaseRequest`): + The fields to update; omitted fields stay unchanged. + knowledge_base_id (`str`): + The knowledge base to update. + user_id (`str`): + Injected authenticated user ID. + service (`KnowledgeBaseService`): + Injected knowledge base service. + + Returns: + `KnowledgeBaseView`: + The knowledge base record after the update. + """ + record = await service.update_knowledge_base( + user_id=user_id, + knowledge_base_id=knowledge_base_id, + name=body.name, + description=body.description, + ) + return KnowledgeBaseView( + id=record.id, + name=record.name, + description=record.description, + embedding_model_config=record.embedding_model_config, + created_at=record.created_at, + updated_at=record.updated_at, + ) + + +@knowledge_base_router.delete( + "/{knowledge_base_id}", + status_code=status.HTTP_204_NO_CONTENT, + summary="Delete a knowledge base", +) +async def delete_knowledge_base( + knowledge_base_id: str = Path(description="The knowledge base id."), + user_id: str = Depends(get_current_user_id), + service: "KnowledgeBaseService" = Depends(get_knowledge_base_service), +) -> None: + """Permanently delete a knowledge base. + + Drops the underlying vector store collection together with every + associated document and the knowledge base record itself. + + Args: + knowledge_base_id (`str`): + The knowledge base to delete. + user_id (`str`): + Injected authenticated user ID. + service (`KnowledgeBaseService`): + Injected knowledge base service. + """ + await service.delete_knowledge_base(user_id, knowledge_base_id) + + +# ---------------------------------------------------------------------- +# Document management +# ---------------------------------------------------------------------- + + +@knowledge_base_router.get( + "/{knowledge_base_id}/documents", + response_model=ListKnowledgeDocumentsResponse, + summary="List documents registered in a knowledge base", +) +async def list_knowledge_documents( + knowledge_base_id: str = Path(description="The knowledge base id."), + user_id: str = Depends(get_current_user_id), + service: "KnowledgeBaseService" = Depends(get_knowledge_base_service), +) -> ListKnowledgeDocumentsResponse: + """List every document registered against a knowledge base. + + Reads from the storage backend (service-mode source of truth), so + documents in any lifecycle state — including ``pending`` / + ``parsing`` / ``error`` — are returned alongside ``ready`` ones. + + Args: + knowledge_base_id (`str`): + The target knowledge base id. + user_id (`str`): + Injected authenticated user ID. + service (`KnowledgeBaseService`): + Injected knowledge base service. + + Returns: + `ListKnowledgeDocumentsResponse`: + One view per registered document. + """ + records = await service.list_documents(user_id, knowledge_base_id) + views = [KnowledgeDocumentView.from_record(r) for r in records] + return ListKnowledgeDocumentsResponse( + documents=views, + total=len(views), + ) + + +@knowledge_base_router.get( + "/{knowledge_base_id}/documents/status", + response_model=ListKnowledgeDocumentStatusResponse, + summary="Batch-query indexing status of one or more documents", +) +async def list_knowledge_document_status( + knowledge_base_id: str = Path(description="The knowledge base id."), + ids: str = Query( + description=( + "Comma-separated list of document ids to query. " + "Missing ids are silently omitted from the response." + ), + ), + user_id: str = Depends(get_current_user_id), + service: "KnowledgeBaseService" = Depends(get_knowledge_base_service), +) -> ListKnowledgeDocumentStatusResponse: + """Return the current lifecycle state of a batch of documents. + + Designed for the front-end's status polling loop: the page sends + every in-flight document id at once so per-document round-trips + do not multiply with concurrency. + + Args: + knowledge_base_id (`str`): + The target knowledge base id. + ids (`str`): + Comma-separated document ids. + user_id (`str`): + Injected authenticated user ID. + service (`KnowledgeBaseService`): + Injected knowledge base service. + + Returns: + `ListKnowledgeDocumentStatusResponse`: + Views for the matched documents. + """ + document_ids = [tok for tok in (s.strip() for s in ids.split(",")) if tok] + records = await service.get_document_status( + user_id, + knowledge_base_id, + document_ids, + ) + return ListKnowledgeDocumentStatusResponse( + items=[KnowledgeDocumentView.from_record(r) for r in records], + ) + + +@knowledge_base_router.post( + "/{knowledge_base_id}/documents", + response_model=UploadKnowledgeDocumentResponse, + status_code=status.HTTP_201_CREATED, + summary="Upload a document into a knowledge base", +) +async def upload_knowledge_document( + knowledge_base_id: str = Path(description="The knowledge base id."), + file: UploadFile = File( + description="The document to index (PDF, TXT, Markdown, …).", + ), + content_type: str + | None = Form( + default=None, + description=( + "Override the IANA media type used to route the upload. " + "Defaults to the type guessed from the filename." + ), + ), + user_id: str = Depends(get_current_user_id), + service: "KnowledgeBaseService" = Depends(get_knowledge_base_service), +) -> UploadKnowledgeDocumentResponse: + """Register an uploaded document and dispatch it for indexing. + + The HTTP connection covers only the upload phase: the request body + is streamed into the blob store, a ``pending`` document record is + persisted, the indexing task is dispatched, and the response is + returned. Parsing / chunking / embedding happen asynchronously in + a worker; the client tracks progress via + :func:`list_knowledge_document_status`. + + Args: + knowledge_base_id (`str`): + The knowledge base to receive the document. + file (`UploadFile`): + The uploaded file (multipart/form-data). + content_type (`str | None`, optional): + Override the IANA media type used to route the upload. + user_id (`str`): + Injected authenticated user ID. + service (`KnowledgeBaseService`): + Injected knowledge base service. + + Returns: + `UploadKnowledgeDocumentResponse`: + The server-assigned document id, filename, and the + initial lifecycle state (always ``"pending"``). + """ + record = await service.register_document( + user_id=user_id, + knowledge_base_id=knowledge_base_id, + filename=file.filename or "uploaded_file", + stream=file.file, + size=file.size or 0, + content_type=content_type or file.content_type, + ) + return UploadKnowledgeDocumentResponse( + document_id=record.id, + filename=record.data.filename, + status=record.data.status, + ) + + +@knowledge_base_router.delete( + "/{knowledge_base_id}/documents/{document_id}", + status_code=status.HTTP_204_NO_CONTENT, + summary="Delete a document from a knowledge base", +) +async def delete_knowledge_document( + knowledge_base_id: str = Path(description="The knowledge base id."), + document_id: str = Path(description="The document id."), + user_id: str = Depends(get_current_user_id), + service: "KnowledgeBaseService" = Depends(get_knowledge_base_service), +) -> None: + """Remove a document and all its chunks from a knowledge base. + + Args: + knowledge_base_id (`str`): + The knowledge base the document belongs to. + document_id (`str`): + The document to delete. + user_id (`str`): + Injected authenticated user ID. + service (`KnowledgeBaseService`): + Injected knowledge base service. + """ + await service.delete_document(user_id, knowledge_base_id, document_id) + + +# ---------------------------------------------------------------------- +# Search +# ---------------------------------------------------------------------- + + +@knowledge_base_router.post( + "/{knowledge_base_id}/search", + response_model=SearchKnowledgeBaseResponse, + summary="Search a knowledge base by natural-language query", +) +async def search_knowledge_base( + body: SearchKnowledgeBaseRequest, + knowledge_base_id: str = Path(description="The knowledge base id."), + user_id: str = Depends(get_current_user_id), + service: "KnowledgeBaseService" = Depends(get_knowledge_base_service), +) -> SearchKnowledgeBaseResponse: + """Run a similarity search over a knowledge base. + + Embeds the query with the knowledge base's configured embedding + model and returns the top-K most similar chunks. + + Args: + body (`SearchKnowledgeBaseRequest`): + The query text and ``top_k``. + knowledge_base_id (`str`): + The knowledge base to search. + user_id (`str`): + Injected authenticated user ID. + service (`KnowledgeBaseService`): + Injected knowledge base service. + + Returns: + `SearchKnowledgeBaseResponse`: + Matched chunks ordered by descending similarity. + """ + results = await service.search( + user_id=user_id, + knowledge_base_id=knowledge_base_id, + query=body.query, + top_k=body.top_k, + ) + return SearchKnowledgeBaseResponse(results=results, total=len(results)) diff --git a/src/agentscope/app/_router/_model.py b/src/agentscope/app/_router/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..06d7b56c6ef659f163aefa3bdca2b627cbaea828 --- /dev/null +++ b/src/agentscope/app/_router/_model.py @@ -0,0 +1,40 @@ +# -*- coding: utf-8 -*- +"""The model router.""" + +from fastapi import APIRouter, Depends, HTTPException, status + +from ._schema import ListModelsResponse, ListModelsRequest +from ...credential import CredentialFactory + +model_router = APIRouter( + prefix="/model", + tags=["model"], + responses={404: {"description": "Not found"}}, +) + + +@model_router.get( + "/", + response_model=ListModelsResponse, + summary="List all candidate models under the given credential type", +) +async def list_models( + body: ListModelsRequest = Depends(), +) -> ListModelsResponse: + """Return all candidate models under the given credential type. + + Args: + body (ListModelsRequest): The request body. + + Returns: + `ListModelsResponse`: The response body. + """ + credential_cls = CredentialFactory.get_credential_class(body.provider) + if credential_cls is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Provider '{body.provider}' not found.", + ) + + models = credential_cls.get_chat_model_class().list_models() + return ListModelsResponse(models=models, total=len(models)) diff --git a/src/agentscope/app/_router/_schedule.py b/src/agentscope/app/_router/_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..4eae944666fe6c5e4480b03f24335507e79c5555 --- /dev/null +++ b/src/agentscope/app/_router/_schedule.py @@ -0,0 +1,239 @@ +# -*- coding: utf-8 -*- +"""Schedule router — CRUD endpoints for scheduled agent tasks.""" +from datetime import datetime + +from fastapi import APIRouter, Depends, HTTPException, status + +from .._manager import SchedulerManager +from ..deps import ( + get_current_user_id, + get_scheduler_manager, + get_session_service, + get_storage, +) +from ._schema import ( + CreateScheduleRequest, + CreateScheduleResponse, + ListSchedulesResponse, + ScheduleSessionsResponse, + UpdateScheduleRequest, +) +from .._service import SessionService +from ..storage import ( + StorageBase, + ScheduleData, + ScheduleRecord, + ScheduleSource, +) + +schedule_router = APIRouter( + prefix="/schedule", + tags=["schedule"], + responses={404: {"description": "Not found"}}, +) + + +@schedule_router.get( + "/", + response_model=ListSchedulesResponse, + summary="List all schedules", +) +async def list_schedules( + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), +) -> ListSchedulesResponse: + """List all schedules owned by the current user. + + Args: + user_id (`str`): Authenticated user ID. + storage (`StorageBase`): Storage instance. + + Returns: + `ListSchedulesResponse`: + Paginated list of schedule records. + """ + schedules = await storage.list_schedules(user_id) + return ListSchedulesResponse(schedules=schedules, total=len(schedules)) + + +@schedule_router.post( + "/", + response_model=CreateScheduleResponse, + status_code=status.HTTP_201_CREATED, + summary="Create a new schedule", +) +async def create_schedule( + body: CreateScheduleRequest, + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), + scheduler: SchedulerManager = Depends(get_scheduler_manager), +) -> CreateScheduleResponse: + """Create a new schedule and register it with the scheduler. + + Args: + body (`CreateScheduleRequest`): Schedule configuration. + user_id (`str`): Authenticated user ID. + storage (`StorageBase`): Storage instance. + scheduler (`SchedulerManager`): Scheduler manager. + + Returns: + `CreateScheduleResponse`: + The ID of the newly created schedule. + + Raises: + `HTTPException`: 404 if the specified agent does not exist. + """ + agent = await storage.get_agent(user_id, body.agent_id) + if agent is None or agent.user_id != user_id: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Agent '{body.agent_id}' not found.", + ) + + record = ScheduleRecord( + user_id=user_id, + agent_id=body.agent_id, + data=ScheduleData( + name=body.name, + description=body.description, + cron_expression=body.cron_expression, + timezone=body.timezone, + enabled=body.enabled, + stateful=body.stateful, + permission_mode=body.permission_mode, + chat_model_config=body.chat_model_config, + source=ScheduleSource.USER, + started_at=datetime.now(), + ), + ) + await storage.upsert_schedule(user_id, record) + + if record.data.enabled: + await scheduler.register_schedule(record) + + return CreateScheduleResponse(schedule_id=record.id) + + +@schedule_router.patch( + "/{schedule_id}", + response_model=ScheduleRecord, + summary="Update a schedule", +) +async def update_schedule( + schedule_id: str, + body: UpdateScheduleRequest, + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), + scheduler: SchedulerManager = Depends(get_scheduler_manager), +) -> ScheduleRecord: + """Partially update a schedule. + + Fields omitted from the request body keep their current values. + Changing ``cron_expression`` or ``timezone`` immediately reschedules the + APScheduler job. Setting ``enable=False`` removes the job from the + scheduler without deleting the record. + + Args: + schedule_id (`str`): ID of the schedule to update. + body (`UpdateScheduleRequest`): Fields to update. + user_id (`str`): Authenticated user ID. + storage (`StorageBase`): Storage instance. + scheduler (`SchedulerManager`): Scheduler manager. + + Returns: + `ScheduleRecord`: + The updated schedule record. + + Raises: + `HTTPException`: 404 if the schedule does not exist. + """ + existing = await storage.get_schedule(user_id, schedule_id) + if existing is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Schedule '{schedule_id}' not found.", + ) + + updates = body.model_dump(exclude_none=True) + updated_data = existing.data.model_copy(update=updates) + updated_record = existing.model_copy( + update={"data": updated_data, "updated_at": datetime.now()}, + ) + await storage.upsert_schedule(user_id, updated_record) + + # Always remove the existing job first; re-register only if still enabled. + await scheduler.remove_schedule(schedule_id) + if updated_record.data.enabled: + await scheduler.register_schedule(updated_record) + + return updated_record + + +@schedule_router.delete( + "/{schedule_id}", + status_code=status.HTTP_204_NO_CONTENT, + summary="Delete a schedule", +) +async def delete_schedule( + schedule_id: str, + user_id: str = Depends(get_current_user_id), + session_service: SessionService = Depends(get_session_service), + scheduler: SchedulerManager = Depends(get_scheduler_manager), +) -> None: + """Permanently delete a schedule. + + Cancels any in-flight chat run for sessions this schedule has + triggered, removes their records via the session service, and + finally unregisters the APScheduler job. + + Args: + schedule_id (`str`): ID of the schedule to delete. + user_id (`str`): Authenticated user ID. + session_service (`SessionService`): Injected session service. + scheduler (`SchedulerManager`): Scheduler manager. + + Raises: + `HTTPException`: 404 if the schedule does not exist. + """ + deleted = await session_service.delete_schedule(user_id, schedule_id) + if not deleted: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Schedule '{schedule_id}' not found.", + ) + await scheduler.remove_schedule(schedule_id) + + +@schedule_router.get( + "/{schedule_id}/sessions", + response_model=ScheduleSessionsResponse, + summary="List execution sessions for a schedule", +) +async def list_schedule_sessions( + schedule_id: str, + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), +) -> ScheduleSessionsResponse: + """Return all sessions triggered by a given schedule. + + Args: + schedule_id (`str`): ID of the schedule. + user_id (`str`): Authenticated user ID. + storage (`StorageBase`): Storage instance. + + Returns: + `ScheduleSessionsResponse`: + List of execution sessions ordered by creation time (newest first). + + Raises: + `HTTPException`: 404 if the schedule does not exist. + """ + existing = await storage.get_schedule(user_id, schedule_id) + if existing is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Schedule '{schedule_id}' not found.", + ) + + sessions = await storage.list_sessions_by_schedule(user_id, schedule_id) + return ScheduleSessionsResponse(sessions=sessions, total=len(sessions)) diff --git a/src/agentscope/app/_router/_schema/__init__.py b/src/agentscope/app/_router/_schema/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..40428e5955babe15f9a4c5948ff89fa5542829ac --- /dev/null +++ b/src/agentscope/app/_router/_schema/__init__.py @@ -0,0 +1,110 @@ +# -*- coding: utf-8 -*- +"""Schema models for the agent service.""" + +from ._chat import ChatRequest, ChatTriggerResponse +from ._model import ListModelsResponse, ListModelsRequest +from ._tts_model import ListTTSModelsResponse, ListTTSModelsRequest +from ._schedule import ( + CreateScheduleRequest, + CreateScheduleResponse, + ListSchedulesResponse, + ScheduleSessionsResponse, + UpdateScheduleRequest, +) +from ._agent import ( + AgentSchemaResponse, + ListAgentsResponse, + CreateAgentRequest, + CreateAgentResponse, + UpdateAgentRequest, +) +from ._credential import ( + CreateCredentialRequest, + CreateCredentialResponse, + UpdateCredentialRequest, + ListCredentialsResponse, + ListCredentialSchemasResponse, +) +from ._knowledge_base import ( + CreateKnowledgeBaseRequest, + CreateKnowledgeBaseResponse, + KbEmbeddingProvider, + KbMiddlewareParametersSchemaResponse, + KnowledgeBaseView, + KnowledgeDocumentView, + ListKbEmbeddingModelsResponse, + ListKnowledgeBasesResponse, + ListKnowledgeDocumentsResponse, + ListKnowledgeDocumentStatusResponse, + ListSupportedContentTypesResponse, + SearchKnowledgeBaseRequest, + SearchKnowledgeBaseResponse, + UpdateKnowledgeBaseRequest, + UploadKnowledgeDocumentResponse, +) +from ._session import ( + CreateSessionRequest, + CreateSessionResponse, + UpdateSessionRequest, + ListSessionsResponse, + ListMessagesResponse, + SessionView, + TeamDetailResponse, + TeamMemberView, +) + +__all__ = [ + # Agent + "AgentSchemaResponse", + "ListAgentsResponse", + "CreateAgentRequest", + "CreateAgentResponse", + "UpdateAgentRequest", + "ListSchedulesResponse", + # Chat + "ChatRequest", + "ChatTriggerResponse", + # Credential + "CreateCredentialRequest", + "CreateCredentialResponse", + "UpdateCredentialRequest", + "ListCredentialsResponse", + "ListCredentialSchemasResponse", + # Knowledge base + "CreateKnowledgeBaseRequest", + "CreateKnowledgeBaseResponse", + "KbEmbeddingProvider", + "KbMiddlewareParametersSchemaResponse", + "KnowledgeBaseView", + "KnowledgeDocumentView", + "ListKbEmbeddingModelsResponse", + "ListKnowledgeBasesResponse", + "ListKnowledgeDocumentsResponse", + "ListKnowledgeDocumentStatusResponse", + "ListSupportedContentTypesResponse", + "SearchKnowledgeBaseRequest", + "SearchKnowledgeBaseResponse", + "UpdateKnowledgeBaseRequest", + "UploadKnowledgeDocumentResponse", + # Model + "ListModelsRequest", + "ListModelsResponse", + # TTS Model + "ListTTSModelsRequest", + "ListTTSModelsResponse", + # Schedule + "CreateScheduleRequest", + "CreateScheduleResponse", + "ListSchedulesResponse", + "ScheduleSessionsResponse", + "UpdateScheduleRequest", + # Session + "CreateSessionRequest", + "CreateSessionResponse", + "UpdateSessionRequest", + "ListSessionsResponse", + "ListMessagesResponse", + "SessionView", + "TeamDetailResponse", + "TeamMemberView", +] diff --git a/src/agentscope/app/_router/_schema/_agent.py b/src/agentscope/app/_router/_schema/_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..4bbf5180183790e510d7084c5017657bcc01b3b2 --- /dev/null +++ b/src/agentscope/app/_router/_schema/_agent.py @@ -0,0 +1,82 @@ +# -*- coding: utf-8 -*- +"""Request / response schemas for the agent router.""" +from pydantic import BaseModel, Field + +from ....agent import ContextConfig, ReActConfig +from ...storage import AgentRecord + + +class CreateAgentRequest(BaseModel): + """Request body for creating a new agent.""" + + name: str = Field(description="Display name of the agent.") + system_prompt: str = Field( + default="You're a helpful assistant.", + description="Base system prompt fed to the agent.", + ) + context_config: ContextConfig = Field( + default_factory=ContextConfig, + description="Context-window management configuration.", + ) + react_config: ReActConfig = Field( + default_factory=ReActConfig, + description="ReAct loop configuration.", + ) + + +class CreateAgentResponse(BaseModel): + """Response body after creating an agent.""" + + agent_id: str = Field(description="Server-assigned agent identifier.") + + +class UpdateAgentRequest(BaseModel): + """Request body for partially updating an agent. + + Omit any field to keep its current value. + """ + + name: str | None = Field(default=None, description="New display name.") + system_prompt: str | None = Field( + default=None, + description="New system prompt.", + ) + context_config: ContextConfig | None = Field( + default=None, + description="New context configuration.", + ) + react_config: ReActConfig | None = Field( + default=None, + description="New ReAct loop configuration.", + ) + + +class ListAgentsResponse(BaseModel): + """Response body for listing agents.""" + + agents: list[AgentRecord] = Field(description="Agent records.") + total: int = Field(description="Total number of agents.") + + +class AgentSchemaResponse(BaseModel): + """JSON Schema fragments used by the frontend to render the agent + create / edit forms. + + Each fragment is a self-contained JSON Schema object so the frontend + doesn't need to follow ``$ref`` links across fragments. The frontend + pairs each property with an i18n key derived from its path, so labels + and descriptions remain localizable independently of the backend. + """ + + identity: dict = Field( + description=( + "Schema for the agent's identity fields (``name``, " + "``system_prompt``)." + ), + ) + context_config: dict = Field( + description="Schema for ``ContextConfig``.", + ) + react_config: dict = Field( + description="Schema for ``ReActConfig``.", + ) diff --git a/src/agentscope/app/_router/_schema/_chat.py b/src/agentscope/app/_router/_schema/_chat.py new file mode 100644 index 0000000000000000000000000000000000000000..cd586b4b10ccb22bd931cf51a373ec0ac5cdde30 --- /dev/null +++ b/src/agentscope/app/_router/_schema/_chat.py @@ -0,0 +1,45 @@ +# -*- coding: utf-8 -*- +"""The chat endpoint schema.""" + +from pydantic import BaseModel, Field + +from ....message import Msg +from ....event import UserConfirmResultEvent, ExternalExecutionResultEvent + + +class ChatRequest(BaseModel): + """Request body for the chat endpoint.""" + + agent_id: str = Field( + description="Agent ID for the chat endpoint.", + ) + + session_id: str = Field( + description="The session to send the message to.", + ) + + input: ( + Msg + | list[Msg] + | UserConfirmResultEvent + | ExternalExecutionResultEvent + | None + ) = Field( + description="The input message(s), or agent event, or None.", + ) + + +class ChatTriggerResponse(BaseModel): + """Response body for the fire-and-forget chat trigger. + + Confirms that the chat run was scheduled. Events produced by the + run arrive separately via the session's SSE stream endpoint. + """ + + status: str = Field( + default="started", + description='Always ``"started"`` when the trigger succeeded.', + ) + session_id: str = Field( + description="Echo of the session id the run was started for.", + ) diff --git a/src/agentscope/app/_router/_schema/_credential.py b/src/agentscope/app/_router/_schema/_credential.py new file mode 100644 index 0000000000000000000000000000000000000000..82a3583761331d603f445e8248a6a27b02e0f936 --- /dev/null +++ b/src/agentscope/app/_router/_schema/_credential.py @@ -0,0 +1,42 @@ +# -*- coding: utf-8 -*- +"""Request / response schemas for the credential router.""" +from pydantic import BaseModel, Field + +from ...storage import CredentialRecord + + +class CreateCredentialRequest(BaseModel): + """Request body for creating a new credential.""" + + data: dict = Field(description="Credential payload (e.g. API keys).") + + +class CreateCredentialResponse(BaseModel): + """Response body after creating a credential.""" + + credential_id: str = Field( + description="Server-assigned credential identifier.", + ) + + +class UpdateCredentialRequest(BaseModel): + """Request body for updating an existing credential.""" + + data: dict = Field(description="New credential payload.") + + +class ListCredentialsResponse(BaseModel): + """Response body for listing credentials.""" + + credentials: list[CredentialRecord] = Field( + description="Credential records.", + ) + total: int = Field(description="Total number of credentials.") + + +class ListCredentialSchemasResponse(BaseModel): + """Response body for listing credential type schemas.""" + + schemas: list[dict] = Field( + description="JSON schemas for all registered credential types.", + ) diff --git a/src/agentscope/app/_router/_schema/_knowledge_base.py b/src/agentscope/app/_router/_schema/_knowledge_base.py new file mode 100644 index 0000000000000000000000000000000000000000..71b4665417eb33e08f99d5d0e635deb404b94974 --- /dev/null +++ b/src/agentscope/app/_router/_schema/_knowledge_base.py @@ -0,0 +1,292 @@ +# -*- coding: utf-8 -*- +"""Request / response schemas for the knowledge base router.""" +from datetime import datetime + +from pydantic import BaseModel, Field + +from ...storage import ( + CredentialRecord, + EmbeddingModelConfig, + KnowledgeDocumentRecord, + KnowledgeDocumentStatus, +) +from ....embedding import EmbeddingModelCard +from ....rag import VectorSearchResult +from ...rag.knowledge_base_manager._dimension_policy import DimensionPolicy + + +class CreateKnowledgeBaseRequest(BaseModel): + """Request body for creating a new knowledge base.""" + + name: str = Field(description="Display name of the knowledge base.") + description: str = Field( + default="", + description="Free-form description shown in the UI.", + ) + embedding_model_config: EmbeddingModelConfig = Field( + description=( + "Embedding model used both at indexing and at query time. " + "Cannot be changed after creation — switching would " + "invalidate every previously inserted vector." + ), + ) + + +class CreateKnowledgeBaseResponse(BaseModel): + """Response body after creating a knowledge base.""" + + knowledge_base_id: str = Field( + description="Server-assigned knowledge base identifier.", + ) + + +class UpdateKnowledgeBaseRequest(BaseModel): + """Request body for updating a knowledge base. + + Only mutable fields can be set here. The embedding model + configuration is pinned at creation time and cannot be changed — + switching it would invalidate every previously inserted vector. + """ + + name: str | None = Field( + default=None, + description="New display name; omit to leave unchanged.", + ) + description: str | None = Field( + default=None, + description="New free-form description; omit to leave unchanged.", + ) + + +class KnowledgeBaseView(BaseModel): + """A knowledge base record as exposed to API clients. + + Mirrors :class:`KnowledgeBaseRecord` with the internal + ``user_id`` / ``collection_name`` fields stripped — clients have + no business introspecting either. + """ + + id: str = Field(description="The knowledge base identifier.") + name: str = Field(description="Display name of the knowledge base.") + description: str = Field(description="Free-form description.") + embedding_model_config: EmbeddingModelConfig = Field( + description="Embedding model configuration pinned at creation.", + ) + created_at: datetime = Field(description="Creation timestamp.") + updated_at: datetime = Field(description="Last-update timestamp.") + + +class ListKnowledgeBasesResponse(BaseModel): + """Response body for listing the caller's knowledge bases.""" + + knowledge_bases: list[KnowledgeBaseView] = Field( + description="All knowledge bases owned by the caller.", + ) + total: int = Field(description="Total number of returned items.") + + +class KnowledgeDocumentView(BaseModel): + """A document record as exposed to API clients. + + Surfaces both the static fields the UI needs to render a row + (``filename`` / ``size``) and the live lifecycle fields the front + end polls (``status`` / ``error`` / ``chunk_count``). Internal + fields (``user_id`` / ``blob_uri`` / ``processing_node`` / lease) + are deliberately omitted — clients have no business introspecting + them. + """ + + id: str = Field(description="The document identifier.") + filename: str = Field(description="Original filename at upload time.") + size: int = Field(description="Document size in bytes.") + content_type: str | None = Field( + default=None, + description="IANA media type recorded at upload time, if any.", + ) + status: KnowledgeDocumentStatus = Field( + description="Current lifecycle state of the document.", + ) + error: str | None = Field( + default=None, + description=( + "Human-readable failure reason when ``status == 'error'``." + ), + ) + chunk_count: int = Field( + default=0, + description="Number of chunks indexed so far.", + ) + created_at: datetime = Field(description="Upload timestamp.") + updated_at: datetime = Field( + description="Last status transition timestamp.", + ) + + @classmethod + def from_record( + cls, + record: KnowledgeDocumentRecord, + ) -> "KnowledgeDocumentView": + """Project a storage record onto the API view. + + Centralised so router code stays a one-liner and the field + mapping has exactly one source of truth. + """ + return cls( + id=record.id, + filename=record.data.filename, + size=record.data.size, + content_type=record.data.content_type, + status=record.data.status, + error=record.data.error, + chunk_count=record.data.chunk_count, + created_at=record.created_at, + updated_at=record.updated_at, + ) + + +class ListKnowledgeDocumentsResponse(BaseModel): + """Response body for listing documents inside a knowledge base.""" + + documents: list[KnowledgeDocumentView] = Field( + description="One view per registered document.", + ) + total: int = Field(description="Total number of returned items.") + + +class ListKnowledgeDocumentStatusResponse(BaseModel): + """Response body for batch document-status polling.""" + + items: list[KnowledgeDocumentView] = Field( + description=( + "Subset of the requested documents that still exist. " + "Missing ids are silently omitted — clients may legitimately " + "ask about a document that was deleted between two polls." + ), + ) + + +class UploadKnowledgeDocumentResponse(BaseModel): + """Response body after uploading a document into a knowledge base.""" + + document_id: str = Field( + description="Server-assigned document identifier.", + ) + filename: str = Field( + description="The original filename of the uploaded document.", + ) + status: KnowledgeDocumentStatus = Field( + description=( + "Lifecycle state immediately after upload — always " + "``'pending'`` in the happy path; surfaced so the client " + "can seed its progress tracker without an extra round-trip." + ), + ) + + +class SearchKnowledgeBaseRequest(BaseModel): + """Request body for searching a knowledge base.""" + + query: str = Field(description="The natural-language search query.") + top_k: int = Field( + default=5, + ge=1, + le=50, + description="Maximum number of results to return.", + ) + + +class SearchKnowledgeBaseResponse(BaseModel): + """Response body for a knowledge base search.""" + + results: list[VectorSearchResult] = Field( + description="Matched chunks ordered by descending similarity score.", + ) + total: int = Field(description="Total number of returned results.") + + +class KbEmbeddingProvider(BaseModel): + """One credential and the embedding models it can serve. + + The model cards have been projected through the manager's + dimension policy: incompatible models are removed and matryoshka + cards are narrowed to the locked dimension when applicable. + """ + + credential: CredentialRecord = Field( + description="The credential record exposing these models.", + ) + models: list[EmbeddingModelCard] = Field( + description=( + "Embedding model cards available under this credential, " + "filtered to those compatible with the manager's " + "dimension policy." + ), + ) + + +class ListKbEmbeddingModelsResponse(BaseModel): + """Response body listing KB-compatible embedding models. + + The list is pre-filtered server-side against the manager's + dimension policy. The policy itself is also returned so the + front-end can render an explanatory banner and lock the dimension + selector when applicable. + """ + + providers: list[KbEmbeddingProvider] = Field( + description=( + "One entry per credential that has at least one " + "compatible embedding model." + ), + ) + policy: DimensionPolicy = Field( + description=( + "The dimension policy used to filter the cards; surfaced " + "verbatim so the UI can explain *why* models were filtered." + ), + ) + + +class KbMiddlewareParametersSchemaResponse(BaseModel): + """Response body exposing the KB middleware's parameters schema. + + The schema is derived from + :class:`agentscope.middleware.RAGMiddleware.Parameters` + via ``model_json_schema()`` so the front-end can render the + session-level KB attachment form with the same schema-driven + component used for model parameters. + """ + + parameter_schema: dict = Field( + description=( + "JSON Schema produced by `RAGMiddleware.Parameters" + "model_json_schema()`. Shaped identically to the " + "`parameter_schema` field on `ModelCard`." + ), + ) + + +class ListSupportedContentTypesResponse(BaseModel): + """Response body advertising the parser-supported upload types. + + Aggregated across every parser registered on the app — the union of + each parser's :attr:`supported_media_types` and + :meth:`supported_extensions`. The front-end uses this to populate + ```` and to reject unsupported drops on the client + before the file leaves the browser. + """ + + media_types: list[str] = Field( + description=( + "Union of IANA media types every registered parser claims " + "to handle. Deduplicated and sorted." + ), + ) + extensions: list[str] = Field( + description=( + "Filename extensions (each starting with `.`) every " + "registered parser claims to handle. Deduplicated and " + "sorted. Derived from `mimetypes` by the base parser; " + "subclasses may override the default." + ), + ) diff --git a/src/agentscope/app/_router/_schema/_mcp.py b/src/agentscope/app/_router/_schema/_mcp.py new file mode 100644 index 0000000000000000000000000000000000000000..4436462b221831e7c79fd0c96497979524f8a49a --- /dev/null +++ b/src/agentscope/app/_router/_schema/_mcp.py @@ -0,0 +1,127 @@ +# -*- coding: utf-8 -*- +"""MCP schemas for API requests and responses.""" +from enum import Enum + +from pydantic import BaseModel, Field + +from ....mcp import StdioMCPConfig, HttpMCPConfig + + +class ConnectionScope(str, Enum): + """MCP connection scope and lifecycle strategy. + + This determines how MCP connections are managed in the service layer. + """ + + SHARED = "shared" + """Shared connection across all agents/users. + - One connection per MCP config, shared globally + - Created on first use, destroyed on application shutdown + - Use case: Stateless HTTP MCP (e.g., weather API, web search) + """ + + ISOLATED = "isolated" + """Isolated connection per agent. + - One connection per (MCP config, agent) + - Created on first use per agent, destroyed on agent session end + - Use case: Stateful MCP (e.g., browser-use), STDIO MCP + """ + + EPHEMERAL = "ephemeral" + """Ephemeral connection per request. + - New connection created for each request, destroyed immediately after + - No connection pooling + - Use case: Low-frequency stateless HTTP MCP + """ + + +class MCPBase(BaseModel): + """Base MCP fields shared across request/response schemas.""" + + name: str = Field( + title="MCP Name", + description="The unique name to identify this MCP configuration.", + ) + + connection_scope: ConnectionScope = Field( + title="Connection Scope", + description="The connection scope and lifecycle strategy.", + ) + + mcp_config: StdioMCPConfig | HttpMCPConfig = Field( + discriminator="type", + title="MCP Config", + description="The base MCP server configuration.", + ) + + def validate_config(self) -> None: + """Validate the configuration. + + Raises: + ValueError: If the configuration is invalid. + """ + # STDIO MCP cannot use ephemeral mode + if ( + self.mcp_config.type == "stdio_mcp" + and self.connection_scope == ConnectionScope.EPHEMERAL + ): + raise ValueError( + "STDIO MCP does not support ephemeral mode. " + "Use 'shared' or 'isolated' instead.", + ) + + +class MCPCreateRequest(MCPBase): + """Request body for creating a new MCP configuration. + + Used in POST /mcp endpoint. Does not include server-generated fields + like creator_id, created_at, updated_at. + """ + + +class MCPUpdateRequest(BaseModel): + """Request body for partially updating an MCP configuration. + + Used in PATCH /mcp/{name} endpoint. All fields are optional. + """ + + connection_scope: ConnectionScope | None = Field( + default=None, + description="New connection scope.", + ) + mcp_config: StdioMCPConfig | HttpMCPConfig | None = Field( + default=None, + discriminator="type", + description="New MCP server configuration.", + ) + + +class MCPResponse(MCPBase): + """Response model for MCP configuration with server-generated metadata. + + Used in GET /mcp/{name}, GET /mcp (list), and POST /mcp responses. + Includes all fields from MCPBase plus server-assigned metadata. + """ + + creator_id: str = Field( + description="User ID of the creator.", + ) + + created_at: float = Field( + description="Creation timestamp (Unix epoch).", + ) + + updated_at: float = Field( + description="Last-updated timestamp (Unix epoch).", + ) + + +class ListMCPsResponse(BaseModel): + """Response model for listing MCP configurations.""" + + mcps: list[MCPResponse] = Field( + description="List of MCP configurations.", + ) + total: int = Field( + description="Total number of MCP configurations.", + ) diff --git a/src/agentscope/app/_router/_schema/_model.py b/src/agentscope/app/_router/_schema/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..2aa2174da7e7ecc4e9ba35405997253bed8bb5bb --- /dev/null +++ b/src/agentscope/app/_router/_schema/_model.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- +"""The chat model configuration, used as DTO layer.""" + +from pydantic import BaseModel, Field + +from ....model import ModelCard + + +class ListModelsResponse(BaseModel): + """List the candidate models response.""" + + models: list[ModelCard] = Field(description="The candidate models.") + total: int = Field(description="The total number of candidates.") + + +class ListModelsRequest(BaseModel): + """List the candidate models request.""" + + provider: str = Field( + description="The provider type, e.g. openai, dashscope, etc.", + ) diff --git a/src/agentscope/app/_router/_schema/_schedule.py b/src/agentscope/app/_router/_schema/_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..9261103e7b6a2ce7c7fe544bba79eebfb59d1194 --- /dev/null +++ b/src/agentscope/app/_router/_schema/_schedule.py @@ -0,0 +1,118 @@ +# -*- coding: utf-8 -*- +"""Request / response schemas for the schedule router.""" +from pydantic import BaseModel, Field + +from ...storage import ( + ScheduleRecord, + SessionRecord, + ChatModelConfig, +) +from ....permission import PermissionMode + + +class CreateScheduleRequest(BaseModel): + """Request body for creating a new schedule.""" + + name: str = Field(description="Display name of the schedule.") + + description: str = Field(default="", description="Optional description.") + + cron_expression: str = Field( + description="Standard 5-field cron expression, e.g. '0 9 * * 1-5'.", + ) + + timezone: str = Field( + default="UTC", + description="IANA timezone name, e.g. 'America/New_York' or " + "'Asia/Shanghai'.", + ) + + agent_id: str = Field(description="Agent to run when the schedule fires.") + + chat_model_config: ChatModelConfig = Field( + description="Model configuration for the auto-created session.", + ) + + enabled: bool = Field( + default=True, + description="Whether the schedule is active immediately " + "after creation.", + ) + + stateful: bool = Field( + default=False, + description="If True, consecutive executions share the same session " + "context.", + ) + + permission_mode: PermissionMode = Field( + default=PermissionMode.DONT_ASK, + description="Permission level for the agent during " + "scheduled execution.", + ) + + +class CreateScheduleResponse(BaseModel): + """Response body after creating a schedule.""" + + schedule_id: str = Field( + description="Server-assigned schedule identifier.", + ) + + +class UpdateScheduleRequest(BaseModel): + """Request body for partially updating a schedule. + + Omit any field to keep its current value. Changing ``cron_expression`` + or ``timezone`` will reschedule the APScheduler job immediately. + Changing ``enable`` to ``False`` removes the job from the scheduler + without deleting the record; setting it back to ``True`` re-registers it. + """ + + name: str | None = Field(default=None, description="New display name.") + + description: str | None = Field( + default=None, + description="New description.", + ) + + cron_expression: str | None = Field( + default=None, + description="New cron expression. Reschedules the task immediately.", + ) + + timezone: str | None = Field( + default=None, + description="New IANA timezone name.", + ) + + enabled: bool | None = Field( + default=None, + description="Set to False to pause the schedule without deleting it.", + ) + + stateful: bool | None = Field( + default=None, + description="Change whether executions share session context.", + ) + + permission_mode: PermissionMode | None = Field( + default=None, + description="New permission mode.", + ) + + +class ListSchedulesResponse(BaseModel): + """Response body for listing schedules.""" + + schedules: list[ScheduleRecord] = Field(description="Schedule records.") + total: int = Field(description="Total number of schedules.") + + +class ScheduleSessionsResponse(BaseModel): + """Response body for listing execution sessions of a schedule.""" + + sessions: list[SessionRecord] = Field( + description="Sessions triggered by this schedule.", + ) + total: int = Field(description="Total number of execution sessions.") diff --git a/src/agentscope/app/_router/_schema/_session.py b/src/agentscope/app/_router/_schema/_session.py new file mode 100644 index 0000000000000000000000000000000000000000..88a3ac9ab912fc633b73cc3ceaa27b52a0aa3d57 --- /dev/null +++ b/src/agentscope/app/_router/_schema/_session.py @@ -0,0 +1,188 @@ +# -*- coding: utf-8 -*- +"""Request / response schemas for the session router.""" +from pydantic import BaseModel, Field + +from ....permission import PermissionMode +from ...storage import ( + AgentRecord, + ChatModelConfig, + SessionKnowledgeConfig, + TTSModelConfig, + SessionRecord, + TeamRecord, +) + + +class TeamMemberView(BaseModel): + """One row in :attr:`TeamDetailResponse.members`. + + Pairs each member's :class:`AgentRecord` with its single + ``session_id`` so the UI can subscribe to the worker's chat + stream without a separate lookup. + """ + + agent: AgentRecord = Field( + description="The worker agent record.", + ) + session_id: str | None = Field( + default=None, + description=( + "The worker's session id. ``None`` if the agent is in an " + "inconsistent state (worker without a session)." + ), + ) + + +class TeamDetailResponse(BaseModel): + """Resolved team detail embedded inside :class:`SessionView.team`.""" + + team: TeamRecord = Field(description="The team record.") + leader_agent: AgentRecord | None = Field( + default=None, + description=( + "Leader's agent record (resolved from the team's " + "``session_id`` → session.agent_id)." + ), + ) + members: list[TeamMemberView] = Field( + default_factory=list, + description=( + "Worker agents listed in :attr:`TeamData.member_ids`, each " + "paired with its single session id when available." + ), + ) + + +class CreateSessionRequest(BaseModel): + """Request body for creating a new session.""" + + agent_id: str = Field(description="Agent this session belongs to.") + workspace_id: str | None = Field( + default=None, + description="Workspace this session belongs to.", + ) + name: str | None = Field( + default=None, + description="Display name. Defaults to current datetime if omitted.", + ) + chat_model_config: ChatModelConfig | None = Field( + default=None, + description="Model provider and parameters. " + "Can be set later via PATCH.", + ) + fallback_chat_model_config: ChatModelConfig | None = Field( + default=None, + description="Fallback model used when the primary model fails. " + "Can be set later via PATCH.", + ) + tts_model_config: TTSModelConfig | None = Field( + default=None, + description="TTS model configuration. Can be set later via PATCH.", + ) + knowledge_config: SessionKnowledgeConfig | None = Field( + default=None, + description=( + "Knowledge bases attached to this session plus the " + "`RAGMiddleware` parameters. Can be set later " + "via PATCH." + ), + ) + + +class CreateSessionResponse(BaseModel): + """Response body after creating a session.""" + + session_id: str = Field(description="Server-assigned session identifier.") + + +class UpdateSessionRequest(BaseModel): + """Request body for updating an existing session. + + Omit any field to keep its current value. + """ + + name: str | None = Field( + default=None, + description="New display name.", + ) + chat_model_config: ChatModelConfig | None = Field( + default=None, + description="New model configuration. " + "Replaces the existing one entirely. " + "Pass null to clear; omit to leave unchanged.", + ) + fallback_chat_model_config: ChatModelConfig | None = Field( + default=None, + description="New fallback model configuration. " + "Pass null to clear; omit to leave unchanged.", + ) + tts_model_config: TTSModelConfig | None = Field( + default=None, + description="New TTS model configuration. " + "Pass null to clear; omit to leave unchanged.", + ) + knowledge_config: SessionKnowledgeConfig | None = Field( + default=None, + description=( + "New knowledge base attachment + middleware parameters. " + "Pass null to clear; omit to leave unchanged." + ), + ) + permission_mode: PermissionMode | None = Field( + default=None, + description="New permission mode for the session.", + ) + + +class SessionView(BaseModel): + """Per-session bundle with everything the frontend needs to + render either the list view or open a session. + + Bundles three orthogonal pieces of information so opening a + session does not require a waterfall of follow-up requests: + + - the persisted :class:`SessionRecord` itself (config + state), + - whether the session has an active chat run right now, + - the team detail (resolved leader + members) when the session + participates in a team. + + Messages are intentionally **not** included here — they are + paginated separately via ``GET /sessions/{id}/messages``. + """ + + session: SessionRecord = Field( + description=( + "The persisted session record. Includes ``state`` " + "(``permission_context`` / ``tool_context`` / " + "``tasks_context``) inline." + ), + ) + is_running: bool = Field( + description="Whether a chat run is currently active on this session.", + ) + team: TeamDetailResponse | None = Field( + default=None, + description=( + "Resolved team detail when ``session.team_id`` is set " + "(leader agent + member agents with their session ids). " + "``None`` when the session does not participate in any team." + ), + ) + + +class ListSessionsResponse(BaseModel): + """Response body for listing sessions.""" + + sessions: list[SessionView] = Field( + description="Session views (record + is_running + team).", + ) + total: int = Field(description="Total number of sessions.") + + +class ListMessagesResponse(BaseModel): + """Response body for listing messages in a session.""" + + messages: list = Field(description="Messages in chronological order.") + is_running: bool = Field( + description="Whether the session is currently running.", + ) diff --git a/src/agentscope/app/_router/_schema/_tts_model.py b/src/agentscope/app/_router/_schema/_tts_model.py new file mode 100644 index 0000000000000000000000000000000000000000..8eb52b3a7c5b76cb3e49c3a5bca064bc685962a4 --- /dev/null +++ b/src/agentscope/app/_router/_schema/_tts_model.py @@ -0,0 +1,23 @@ +# -*- coding: utf-8 -*- +"""The TTS model configuration, used as DTO layer.""" + +from pydantic import BaseModel, Field + +from ....tts import TTSModelCard + + +class ListTTSModelsResponse(BaseModel): + """List the candidate TTS models response.""" + + models: list[TTSModelCard] = Field( + description="The candidate TTS models.", + ) + total: int = Field(description="The total number of candidates.") + + +class ListTTSModelsRequest(BaseModel): + """List the candidate TTS models request.""" + + provider: str = Field( + description="The provider type, e.g. dashscope_credential.", + ) diff --git a/src/agentscope/app/_router/_session.py b/src/agentscope/app/_router/_session.py new file mode 100644 index 0000000000000000000000000000000000000000..5c13610b65bca94a5ae01a7864a60e48169eb31c --- /dev/null +++ b/src/agentscope/app/_router/_session.py @@ -0,0 +1,680 @@ +# -*- coding: utf-8 -*- +"""Session router — create, list, update, delete, stream, and get messages.""" +import asyncio +import json +from typing import AsyncGenerator + +from fastapi import APIRouter, Depends, HTTPException, Query, status +from fastapi.responses import StreamingResponse + +from ..._utils._common import _generate_id +from ..deps import ( + get_current_user_id, + get_message_bus, + get_session_service, + get_storage, +) +from ._schema import ( + CreateSessionRequest, + CreateSessionResponse, + ListMessagesResponse, + ListSessionsResponse, + SessionView, + TeamDetailResponse, + TeamMemberView, + UpdateSessionRequest, +) +from ..message_bus import MessageBus, MessageBusKeys +from .._service import SessionService, SessionProjection, SubagentHitlProjector +from ..storage import ( + AgentRecord, + ChatModelConfig, + SessionKnowledgeConfig, + TTSModelConfig, + SessionConfig, + SessionRecord, + StorageBase, + TeamRecord, +) +from ...message import ToolCallState +from ...event import CustomEvent + + +async def _build_team_detail( + storage: StorageBase, + user_id: str, + team: TeamRecord, +) -> TeamDetailResponse: + """Resolve a team's leader agent + member agents into a + :class:`TeamDetailResponse` for the session list endpoint. + + Args: + storage (`StorageBase`): + Application storage. Used to look up the leader session, + each member agent, and each member's session. + user_id (`str`): + The owner user id. + team (`TeamRecord`): + The team to resolve. Caller has already loaded it. + + Returns: + `TeamDetailResponse`: + The team plus its resolved leader and member agents (each + member paired with its session id when available). + """ + leader_agent: AgentRecord | None = None + leader_session = await storage.get_session(user_id, "", team.session_id) + if leader_session is not None: + leader_agent = await storage.get_agent( + user_id, + leader_session.agent_id, + ) + + members: list[TeamMemberView] = [] + for member_id in team.data.member_ids: + agent = await storage.get_agent(user_id, member_id) + if agent is None: + continue + sessions = await storage.list_sessions(user_id, member_id) + session_id = sessions[0].id if sessions else None + members.append(TeamMemberView(agent=agent, session_id=session_id)) + + return TeamDetailResponse( + team=team, + leader_agent=leader_agent, + members=members, + ) + + +session_router = APIRouter( + prefix="/sessions", + tags=["sessions"], + responses={404: {"description": "Not found"}}, +) + + +async def _ensure_credential_exists( + storage: StorageBase, + user_id: str, + config: ChatModelConfig | TTSModelConfig | None, +) -> None: + """Validate that the credential referenced by ``config`` belongs to the + given user. No-op when ``config`` is ``None``. + + Args: + storage (`StorageBase`): Injected storage backend. + user_id (`str`): The authenticated user ID. + config (`ChatModelConfig | TTSModelConfig | None`): Model config to + validate. Pass ``None`` to skip the check. + + Raises: + `HTTPException`: 404 if the credential does not exist or does not + belong to the user. + """ + if config is None: + return + credentials = await storage.list_credentials(user_id) + if not any(c.id == config.credential_id for c in credentials): + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Credential '{config.credential_id}' not found.", + ) + + +async def _ensure_knowledge_bases_exist( + storage: StorageBase, + user_id: str, + config: SessionKnowledgeConfig | None, +) -> None: + """Validate every KB id in ``config`` belongs to the given user. + + No-op when ``config`` is ``None`` or its ``knowledge_base_ids`` + list is empty. + + Args: + storage (`StorageBase`): Injected storage backend. + user_id (`str`): The authenticated user ID. + config (`SessionKnowledgeConfig | None`): + Knowledge config to validate. Pass ``None`` to skip. + + Raises: + `HTTPException`: 404 if any KB id does not exist or is not + owned by the user. + """ + if config is None or not config.knowledge_base_ids: + return + for kb_id in config.knowledge_base_ids: + kb = await storage.get_knowledge_base(user_id, kb_id) + if kb is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Knowledge base '{kb_id}' not found.", + ) + + +@session_router.get( + "/", + response_model=ListSessionsResponse, + summary="List sessions for an agent", +) +async def list_sessions( + agent_id: str = Query(description="Filter sessions by agent ID."), + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), + message_bus: MessageBus = Depends(get_message_bus), +) -> ListSessionsResponse: + """Return all sessions for an agent as enriched + :class:`SessionView` entries. + + Each entry bundles three things the chat UI needs to render + without follow-up requests: the session record (incl. + ``state``), whether a chat run is currently active, and — when + the session participates in a team — the resolved team detail + (leader agent + member agents with their session ids). + + Args: + agent_id (`str`): + Agent whose sessions to list. + user_id (`str`): + Injected authenticated user ID. + storage (`StorageBase`): + Injected storage backend. + message_bus (`MessageBus`): + Injected message bus (used for ``session_is_running``). + + Returns: + `ListSessionsResponse`: + Enriched session views and their count. + + Raises: + `HTTPException`: 404 if the agent does not exist or does not + belong to the authenticated user. + """ + # Direct ownership check via get_agent — handles both source=user + # and source=team agents (the latter aren't returned by + # storage.list_agents but are still owned by the user; reachable + # via team navigation). + agent = await storage.get_agent(user_id, agent_id) + if agent is None or agent.user_id != user_id: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Agent '{agent_id}' not found.", + ) + + sessions = await storage.list_sessions(user_id, agent_id) + views: list[SessionView] = [] + for session in sessions: + team_detail = None + if session.team_id: + team_record = await storage.get_team(user_id, session.team_id) + if team_record is not None: + team_detail = await _build_team_detail( + storage, + user_id, + team_record, + ) + views.append( + SessionView( + session=session, + is_running=await message_bus.is_locked( + MessageBusKeys.session_lock(session.id), + ), + team=team_detail, + ), + ) + return ListSessionsResponse(sessions=views, total=len(views)) + + +@session_router.post( + "/", + response_model=CreateSessionResponse, + status_code=status.HTTP_201_CREATED, + summary="Create a new session", +) +async def create_session( + body: CreateSessionRequest, + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), +) -> CreateSessionResponse: + """Create (or resume) a session for a given agent and workspace. + + At most one session exists per ``(user_id, agent_id, workspace_id)`` + triple — a second call with the same triple updates the existing session + rather than creating a duplicate. + + Args: + body (`CreateSessionRequest`): Agent, workspace, and model config. + user_id (`str`): Injected authenticated user ID. + storage (`StorageBase`): Injected storage backend. + + Returns: + `CreateSessionResponse`: The session identifier. + + Raises: + `HTTPException`: 404 if the agent or credential does not exist or + does not belong to the authenticated user. + """ + agent = await storage.get_agent(user_id, body.agent_id) + if agent is None or agent.user_id != user_id: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Agent '{body.agent_id}' not found.", + ) + + await _ensure_credential_exists(storage, user_id, body.chat_model_config) + await _ensure_credential_exists( + storage, + user_id, + body.fallback_chat_model_config, + ) + await _ensure_credential_exists(storage, user_id, body.tts_model_config) + await _ensure_knowledge_bases_exist( + storage, + user_id, + body.knowledge_config, + ) + + session_record = await storage.upsert_session( + user_id=user_id, + agent_id=body.agent_id, + config=SessionConfig( + workspace_id=body.workspace_id or _generate_id(), + chat_model_config=body.chat_model_config, + fallback_chat_model_config=body.fallback_chat_model_config, + tts_model_config=body.tts_model_config, + knowledge_config=body.knowledge_config, + **({"name": body.name} if body.name is not None else {}), + ), + ) + return CreateSessionResponse(session_id=session_record.id) + + +@session_router.delete( + "/{session_id}", + status_code=status.HTTP_204_NO_CONTENT, + summary="Delete a session", +) +async def delete_session( + session_id: str, + agent_id: str = Query(description="Agent the session belongs to."), + user_id: str = Depends(get_current_user_id), + session_service: SessionService = Depends(get_session_service), +) -> None: + """Permanently delete a session and all its associated state. + + Cancels any in-flight chat run for this session (and for every + worker session if this one is a team leader) before dropping + storage records and bus state. The cancel path is cross-process: + whichever worker is actually running the session will receive the + cancel broadcast and abort. + + Args: + session_id (`str`): The session to delete. + agent_id (`str`): The agent the session belongs to. + user_id (`str`): Injected authenticated user ID. + session_service (`SessionService`): Injected session service. + + Raises: + `HTTPException`: 404 if the session does not exist or does not belong + to the authenticated user. + """ + deleted = await session_service.delete_session( + user_id, + agent_id, + session_id, + ) + if not deleted: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Session '{session_id}' not found.", + ) + + +@session_router.patch( + "/{session_id}", + response_model=SessionRecord, + summary="Update a session", +) +async def update_session( + session_id: str, + body: UpdateSessionRequest, + agent_id: str = Query(description="Agent the session belongs to."), + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), +) -> SessionRecord: + """Update the model configuration of an existing session. + + Args: + session_id (`str`): The session to update. + body (`UpdateSessionRequest`): Fields to update. + user_id (`str`): Injected authenticated user ID. + storage (`StorageBase`): Injected storage backend. + + Returns: + `SessionRecord`: The full session record after the update. + + Raises: + `HTTPException`: 404 if the session, agent, or credential does not + exist or does not belong to the authenticated user. + """ + existing = await storage.get_session(user_id, agent_id, session_id) + if existing is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Session '{session_id}' not found.", + ) + + await _ensure_credential_exists(storage, user_id, body.chat_model_config) + await _ensure_credential_exists( + storage, + user_id, + body.fallback_chat_model_config, + ) + await _ensure_credential_exists(storage, user_id, body.tts_model_config) + await _ensure_knowledge_bases_exist( + storage, + user_id, + body.knowledge_config, + ) + + updated_state = existing.state + if body.permission_mode is not None: + updated_ctx = existing.state.permission_context.model_copy( + update={"mode": body.permission_mode}, + ) + + updated_state = existing.state.model_copy( + update={ + "permission_context": updated_ctx, + }, + ) + + # PATCH semantics: only fields explicitly present in the request body are + # applied. ``exclude_unset=True`` lets clients distinguish "leave + # unchanged" (omit) from "clear" (send ``null``) — required for clearing + # ``fallback_chat_model_config``. + config_updates = body.model_dump( + exclude_unset=True, + exclude={"permission_mode"}, + ) + + return await storage.upsert_session( + user_id=user_id, + agent_id=agent_id, + config=SessionConfig.model_validate( + {**existing.config.model_dump(mode="json"), **config_updates}, + ), + state=updated_state, + session_id=session_id, + ) + + +# ---------------------------------------------------------------------- +# Messages: fetch persisted messages for a session +# ---------------------------------------------------------------------- + + +@session_router.get( + "/{session_id}/messages", + response_model=ListMessagesResponse, + summary="List messages for a session", +) +async def list_messages( + session_id: str, + agent_id: str = Query(description="Agent the session belongs to."), + offset: int = Query(0, ge=0, description="Pagination offset."), + limit: int = Query(50, ge=1, le=200, description="Max messages."), + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), + message_bus: MessageBus = Depends(get_message_bus), +) -> ListMessagesResponse: + """Return persisted messages for a session. + + Args: + session_id: The session to query. + agent_id: Agent the session belongs to. + offset: Pagination offset. + limit: Maximum number of messages to return. + user_id: Injected authenticated user ID. + storage: Injected storage backend. + message_bus: Injected message bus. + + Returns: + Messages and running status. + """ + existing = await storage.get_session(user_id, agent_id, session_id) + if existing is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Session '{session_id}' not found.", + ) + + messages = await storage.list_messages( + user_id, + session_id, + offset=offset, + limit=limit, + ) + return ListMessagesResponse( + messages=messages, + is_running=await message_bus.is_locked( + MessageBusKeys.session_lock(session_id), + ), + ) + + +# ---------------------------------------------------------------------- +# Stream: live SSE connection for session events +# ---------------------------------------------------------------------- + +_HEARTBEAT_INTERVAL_SECS = 30 +# Interval between SSE heartbeat comment frames (``:\\n\\n``). + + +async def _worker_still_asking( + storage: StorageBase, + user_id: str, + worker_agent_id: str, + worker_session_id: str, + reply_id: str, +) -> bool: + """Return whether a worker session is still parked on the ASKING + tool call identified by ``reply_id``. + + This is the reconcile-on-read check (design §3.5): the worker + session's own ``state.context`` is the single source of truth for + "does this confirmation still need answering". A leader-side + pending projection whose worker has already resolved / cancelled + the call is a ghost and must not be replayed. + + Mirrors the wakeup guard in + :meth:`ChatService._run_impl` — a request is "still asking" when + the tail ``AssistantMsg`` of the worker carries a tool call in + ``ASKING`` or ``SUBMITTED`` state for the matching ``reply_id``. + + Args: + storage (`StorageBase`): + Application storage. + user_id (`str`): + The owner user id. + worker_agent_id (`str`): + The worker agent that owns the session. + worker_session_id (`str`): + The worker session to inspect. + reply_id (`str`): + The reply id the pending request belongs to. + + Returns: + `bool`: + ``True`` if the worker is still awaiting confirmation for + ``reply_id``; ``False`` otherwise (resolved, cancelled, or + the session/record is gone). + """ + session = await storage.get_session( + user_id, + worker_agent_id, + worker_session_id, + ) + if session is None or not session.state.context: + return False + last_msg = session.state.context[-1] + if last_msg.role != "assistant" or last_msg.id != reply_id: + return False + return any( + tc.state in (ToolCallState.ASKING, ToolCallState.SUBMITTED) + for tc in last_msg.get_content_blocks("tool_call") + ) + + +@session_router.get( + "/{session_id}/stream", + summary="Subscribe to a session's event stream (SSE)", + response_description="Server-Sent Events stream of AgentEvent objects", +) +async def stream_session_events( + session_id: str, + agent_id: str = Query(description="Agent the session belongs to."), + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), + message_bus: MessageBus = Depends(get_message_bus), +) -> StreamingResponse: + """Subscribe to a session's live event stream. + + Returns a ``text/event-stream`` that first replays any buffered + events from the current run's replay log (if a run is in progress + or just finished), then streams live events as they are produced + by :meth:`ChatService.run`. The connection stays open + until the client disconnects — subsequent runs on the same session + are delivered over the same connection. + + A heartbeat comment frame (``:\\n\\n``) is sent every 30 seconds to + keep the connection alive through reverse proxies. + + Args: + session_id (`str`): + The session to subscribe to. + agent_id (`str`): + The agent that owns the session (used for ownership + validation). + user_id (`str`): + Injected authenticated user id. + storage (`StorageBase`): + Injected storage backend (ownership check only). + message_bus (`MessageBus`): + Injected message bus (replay + live subscription). + + Returns: + `StreamingResponse`: + SSE stream of AgentEvent frames + periodic heartbeats. + """ + existing = await storage.get_session(user_id, agent_id, session_id) + if existing is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Session '{session_id}' not found.", + ) + + async def _sse_generator() -> AsyncGenerator[str, None]: + # 1. Replay buffered events from the current run (if any). + for _entry_id, event in await message_bus.log_read( + MessageBusKeys.session_events(session_id), + max_count=MessageBusKeys.SESSION_REPLAY_MAX_LEN, + ): + yield f"data: {json.dumps(event)}\n\n" + + # 1b. Inject pending subagent HITL cards projected onto this + # session as a team leader (design §3.5). These live in a + # durable Redis hash — NOT in the replay log (trimmed per + # run) nor in the leader's own Msg history — so a fresh + # reconnect after the worker parked still surfaces them. + # + # Reconcile-on-read: the worker session's own context is the + # SSOT. Inject only when the worker is still ASKING; drop and + # delete ghosts (worker resolved/cancelled without clearing). + projection = SessionProjection(message_bus) + for payload in await projection.list( + session_id, + SubagentHitlProjector.KIND, + ): + if not await _worker_still_asking( + storage, + user_id, + payload["worker_agent_id"], + payload["worker_session_id"], + payload["reply_id"], + ): + await projection.delete( + session_id, + SubagentHitlProjector.KIND, + SubagentHitlProjector.entry_id( + payload["worker_session_id"], + payload["reply_id"], + ), + ) + continue + custom = CustomEvent( + name=SubagentHitlProjector.EVT_REQUIRE, + value=payload, + ) + yield f"data: {json.dumps(custom.model_dump(mode='json'))}\n\n" + + # 2. Live subscribe via a background feeder task that pushes + # events into a queue. The main loop reads from the queue + # with a timeout so we can interleave heartbeat frames. + # + # We avoid calling ``wait_for(__anext__())`` on the async + # generator directly because cancelling a suspended + # ``__anext__`` leaves the generator in a "running" state + # that prevents ``aclose()`` from working. + queue: asyncio.Queue[dict | None] = asyncio.Queue() + + async def _feeder() -> None: + """Read from the bus subscription and forward to the queue. + + Pushes ``None`` as a sentinel when the subscription ends + (which in practice only happens if the bus shuts down). + """ + try: + async for evt in message_bus.subscribe( + MessageBusKeys.session_events(session_id), + ): + await queue.put( + {k: v for k, v in evt.items() if k != "_entry_id"}, + ) + except asyncio.CancelledError: + pass + finally: + await queue.put(None) + + feeder_task = asyncio.create_task( + _feeder(), + name=f"sse-feeder:{session_id}", + ) + + try: + while True: + try: + item = await asyncio.wait_for( + queue.get(), + timeout=_HEARTBEAT_INTERVAL_SECS, + ) + if item is None: + break + yield f"data: {json.dumps(item)}\n\n" + except asyncio.TimeoutError: + yield ":\n\n" + finally: + feeder_task.cancel() + try: + await feeder_task + except asyncio.CancelledError: + pass + + return StreamingResponse( + _sse_generator(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", + }, + ) diff --git a/src/agentscope/app/_router/_tts_model.py b/src/agentscope/app/_router/_tts_model.py new file mode 100644 index 0000000000000000000000000000000000000000..758f75e0e35310e6e1806a52c41c56b2812d2abd --- /dev/null +++ b/src/agentscope/app/_router/_tts_model.py @@ -0,0 +1,40 @@ +# -*- coding: utf-8 -*- +"""The TTS model router.""" + +from fastapi import APIRouter, Depends, HTTPException, status + +from ._schema import ListTTSModelsResponse, ListTTSModelsRequest +from ...credential import CredentialFactory + +tts_model_router = APIRouter( + prefix="/tts-model", + tags=["tts-model"], + responses={404: {"description": "Not found"}}, +) + + +@tts_model_router.get( + "/", + response_model=ListTTSModelsResponse, + summary="List all candidate TTS models under the given credential type", +) +async def list_tts_models( + body: ListTTSModelsRequest = Depends(), +) -> ListTTSModelsResponse: + """Return all candidate TTS models under the given credential type. + + Args: + body (ListTTSModelsRequest): The request body. + + Returns: + `ListTTSModelsResponse`: The response body. + """ + credential_cls = CredentialFactory.get_credential_class(body.provider) + if credential_cls is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Provider '{body.provider}' not found.", + ) + + models = credential_cls.list_tts_models() + return ListTTSModelsResponse(models=models, total=len(models)) diff --git a/src/agentscope/app/_router/_workspace.py b/src/agentscope/app/_router/_workspace.py new file mode 100644 index 0000000000000000000000000000000000000000..eb7ea66cd8270a0d1b972af48f6186be9034128b --- /dev/null +++ b/src/agentscope/app/_router/_workspace.py @@ -0,0 +1,220 @@ +# -*- coding: utf-8 -*- +"""Workspace router — manage MCP clients and skills on a workspace.""" +from fastapi import APIRouter, Depends, HTTPException, Query, status +from pydantic import BaseModel, Field + +from ..deps import ( + get_current_user_id, + get_storage, + get_workspace_manager, +) +from ..workspace_manager import WorkspaceManagerBase +from ..storage import StorageBase +from ...mcp import MCPClient +from ...skill import Skill +from ...workspace import WorkspaceBase + +workspace_router = APIRouter(prefix="/workspace", tags=["workspace"]) + + +class AddSkillRequest(BaseModel): + """The request to add skill.""" + + skill_path: str + + +class ToolInfo(BaseModel): + """The tool info.""" + + name: str + description: str | None = None + + +class MCPClientStatus(MCPClient): + """MCPClient enriched with live tool list and health status.""" + + is_healthy: bool = False + tools: list[ToolInfo] = Field(default_factory=list) + + +async def _resolve_workspace( + user_id: str, + agent_id: str, + session_id: str, + storage: StorageBase, + workspace_manager: WorkspaceManagerBase, +) -> WorkspaceBase: + """Resolve the workspace for the given session, raising 404 if not + found.""" + session_record = await storage.get_session(user_id, agent_id, session_id) + if session_record is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Session {session_id!r} not found.", + ) + return await workspace_manager.get_workspace( + user_id, + agent_id, + session_id, + session_record.config.workspace_id, + ) + + +# --------------------------------------------------------------------------- +# MCP endpoints +# --------------------------------------------------------------------------- + + +@workspace_router.get("/mcp") +async def list_mcps( + agent_id: str = Query(...), + session_id: str = Query(...), + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), + workspace_manager: WorkspaceManagerBase = Depends(get_workspace_manager), +) -> list[MCPClientStatus]: + """Return all MCP clients with live tool list and health status.""" + workspace = await _resolve_workspace( + user_id, + agent_id, + session_id, + storage, + workspace_manager, + ) + clients = await workspace.list_mcps() + + results = [] + for client in clients: + base = client.model_dump() + try: + mcp_tools = await client.list_tools() + tools = [ + ToolInfo(name=t.name, description=t.description) + for t in mcp_tools + ] + results.append( + MCPClientStatus( + **base, + is_healthy=True, + tools=tools, + ), + ) + except Exception: + results.append( + MCPClientStatus( + **base, + is_healthy=False, + ), + ) + + return results + + +@workspace_router.post("/mcp", status_code=status.HTTP_201_CREATED) +async def add_mcp( + mcp: MCPClient, + agent_id: str = Query(...), + session_id: str = Query(...), + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), + workspace_manager: WorkspaceManagerBase = Depends(get_workspace_manager), +) -> None: + """Add an MCP client to the session's workspace.""" + workspace = await _resolve_workspace( + user_id, + agent_id, + session_id, + storage, + workspace_manager, + ) + await workspace.add_mcp(mcp) + + +@workspace_router.delete( + "/mcp/{mcp_name}", + status_code=status.HTTP_204_NO_CONTENT, +) +async def remove_mcp( + mcp_name: str, + agent_id: str = Query(...), + session_id: str = Query(...), + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), + workspace_manager: WorkspaceManagerBase = Depends(get_workspace_manager), +) -> None: + """Remove an MCP client from the session's workspace by name.""" + workspace = await _resolve_workspace( + user_id, + agent_id, + session_id, + storage, + workspace_manager, + ) + await workspace.remove_mcp(mcp_name) + + +# --------------------------------------------------------------------------- +# Skill endpoints +# --------------------------------------------------------------------------- + + +@workspace_router.get("/skill") +async def list_skills( + agent_id: str = Query(...), + session_id: str = Query(...), + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), + workspace_manager: WorkspaceManagerBase = Depends(get_workspace_manager), +) -> list[Skill]: + """Return all skills available in the session's workspace.""" + workspace = await _resolve_workspace( + user_id, + agent_id, + session_id, + storage, + workspace_manager, + ) + return await workspace.list_skills() + + +@workspace_router.post("/skill", status_code=status.HTTP_201_CREATED) +async def add_skill( + body: AddSkillRequest, + agent_id: str = Query(...), + session_id: str = Query(...), + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), + workspace_manager: WorkspaceManagerBase = Depends(get_workspace_manager), +) -> None: + """Add a skill to the session's workspace from the given path.""" + workspace = await _resolve_workspace( + user_id, + agent_id, + session_id, + storage, + workspace_manager, + ) + await workspace.add_skill(body.skill_path) + + +@workspace_router.delete( + "/skill/{skill_name}", + status_code=status.HTTP_204_NO_CONTENT, +) +async def remove_skill( + skill_name: str, + agent_id: str = Query(...), + session_id: str = Query(...), + user_id: str = Depends(get_current_user_id), + storage: StorageBase = Depends(get_storage), + workspace_manager: WorkspaceManagerBase = Depends(get_workspace_manager), +) -> None: + """Remove a skill from the session's workspace by name.""" + workspace = await _resolve_workspace( + user_id, + agent_id, + session_id, + storage, + workspace_manager, + ) + await workspace.remove_skill(skill_name) diff --git a/src/agentscope/app/_service/__init__.py b/src/agentscope/app/_service/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7676c2973aab89c13c2f6be9bddc3f540c6d3e76 --- /dev/null +++ b/src/agentscope/app/_service/__init__.py @@ -0,0 +1,29 @@ +# -*- coding: utf-8 -*- +"""Service layer for the AgentScope app.""" +from ._chat import ChatService +from ._embedding import get_embedding_model +from ._index_sweeper import IndexSweeper +from ._index_task_consumer import IndexTaskConsumer +from ._index_worker import IndexWorker +from ._knowledge_base import KnowledgeBaseService +from ._model import get_model +from ._tts_model import get_tts_model +from ._session import SessionService +from ._session_projection import SessionProjection +from ._projectors import SubagentHitlProjector +from ._toolkit import get_toolkit + +__all__ = [ + "ChatService", + "IndexSweeper", + "IndexTaskConsumer", + "IndexWorker", + "KnowledgeBaseService", + "SessionService", + "SessionProjection", + "SubagentHitlProjector", + "get_embedding_model", + "get_model", + "get_tts_model", + "get_toolkit", +] diff --git a/src/agentscope/app/_service/_chat.py b/src/agentscope/app/_service/_chat.py new file mode 100644 index 0000000000000000000000000000000000000000..f5abd761fdca4801282d2ddc4bc391aa0894b14d --- /dev/null +++ b/src/agentscope/app/_service/_chat.py @@ -0,0 +1,588 @@ +# -*- coding: utf-8 -*- +"""Chat service encapsulating agent execution + persistence logic. + +This is the single source of truth for running an agent against a +session. Both the HTTP chat endpoint and the wakeup dispatcher call +:meth:`ChatService.run`, guaranteeing identical message persistence, +middleware wiring, and state handling. + +Events produced by the agent are not exposed back through this method +— they are published to the message bus inside the run, and any client +that wants them subscribes through the +``GET /sessions/{sid}/stream`` SSE endpoint. +""" +from fastapi import HTTPException + +from ..message_bus import MessageBus, MessageBusKeys +from .._bus_ops import publish_session_event +from ..rag.knowledge_base_manager import KnowledgeBaseManagerBase +from ..storage import StorageBase, AgentRecord, SessionRecord +from .._manager import BackgroundTaskManager, SchedulerManager +from ..workspace_manager import WorkspaceManagerBase +from ..middleware import ( + InboxMiddleware, + StateChangeMiddleware, + ToolOffloadMiddleware, +) +from ...middleware import TTSMiddleware, RAGMiddleware +from ...rag import KnowledgeBase +from .._types import ( + AgentMiddlewareFactory, + AgentToolFactory, + EventProjector, + SubAgentTemplate, +) +from ._model import get_model +from ._tts_model import get_tts_model +from ._toolkit import get_toolkit +from ._session_projection import SessionProjection +from ._projectors import SubagentHitlProjector + +from ..._logging import logger +from ...agent import Agent, ModelConfig +from ...event import ( + AgentEvent, + ReplyStartEvent, + UserConfirmResultEvent, + ExternalExecutionResultEvent, +) +from ...message import AssistantMsg, Msg, ToolCallState +from ...permission import AdditionalWorkingDirectory + + +class ChatService: + """Run an agent against a session, persisting input/reply messages + and updated agent state. + + Shared by the HTTP chat endpoint and the wakeup dispatcher so both + paths go through identical validation, assembly, and persistence. + + Session serialisation and event fan-out are both handled by the + :class:`MessageBus`: :meth:`bus.session_run` acquires a distributed + lock (guaranteeing at most one chat run per session across all + processes), and :meth:`bus.session_publish_event` writes each event + to both a replay log (for late-joining subscribers) and a live + Pub/Sub channel. + """ + + def __init__( + self, + storage: StorageBase, + workspace_manager: WorkspaceManagerBase, + scheduler_manager: SchedulerManager, + background_task_manager: BackgroundTaskManager, + message_bus: MessageBus, + knowledge_base_manager: KnowledgeBaseManagerBase | None = None, + extra_agent_middlewares: AgentMiddlewareFactory | None = None, + extra_agent_tools: AgentToolFactory | None = None, + custom_subagent_templates: dict[str, SubAgentTemplate] | None = None, + custom_agent_cls: type[Agent] | None = None, + extra_projectors: list[EventProjector] | None = None, + ) -> None: + """Initialize chat service. + + Args: + storage (`StorageBase`): + Application storage backend. + workspace_manager (`WorkspaceManagerBase`): + Provides per-session workspace (tools, MCPs, skills) used + during agent assembly. + scheduler_manager (`SchedulerManager`): + Application scheduler — passed through to + :func:`get_toolkit` so the agent toolkit gets the four + ``Schedule*`` tools. + background_task_manager (`BackgroundTaskManager`): + Tracks offloaded long-running tool tasks. Also provides + the :class:`ToolStop` tool through + :func:`get_toolkit`. + message_bus (`MessageBus`): + Application-wide message bus. Provides session-level + distributed locking (via :meth:`session_run`), event + replay + live fan-out (via :meth:`session_publish_event`), + and inbox delivery (via :class:`InboxMiddleware`). + knowledge_base_manager (`KnowledgeBaseManagerBase | None`, \ + optional): + The application's knowledge base manager. When + provided and the session config carries a + ``knowledge_config``, a + :class:`~agentscope.middleware.RAGMiddleware` + is attached to the agent at run time. ``None`` + disables knowledge-base wiring even for sessions that + have one configured. + extra_agent_middlewares (`AgentMiddlewareFactory | None`, \ + optional): + Async factory invoked at every chat turn to produce + user/session-specific middlewares to attach to the agent. + extra_agent_tools (`AgentToolFactory | None`, optional): + Async factory invoked at every chat turn to produce + user/session-specific tools to register in the toolkit. + custom_subagent_templates (`dict[str, SubAgentTemplate] | None`,\ + optional): + Sub-agent template registry, keyed by template type. + Passed through to :func:`get_toolkit` so that + ``AgentCreate`` can route to the appropriate template + when a ``subagent_type`` is specified. + custom_agent_cls (`type[Agent] | None`, optional): + Custom :class:`Agent` subclass for assembling agents. + Falls back to :class:`Agent` when ``None``. + extra_projectors (`list[EventProjector] | None`, optional): + Additional cross-session event projectors to run after + the built-in ones (mirrors the ``extra_agent_*`` + injection style). Each is invoked once per produced + event to mirror a UI feed onto another session; see + :class:`~agentscope.app._types.EventProjector`. + """ + self._storage = storage + self._workspace_manager = workspace_manager + self._scheduler_manager = scheduler_manager + self._background_task_manager = background_task_manager + self._message_bus = message_bus + self._knowledge_base_manager = knowledge_base_manager + self._extra_agent_middlewares = extra_agent_middlewares + self._extra_agent_tools = extra_agent_tools + self._sub_agent_templates = custom_subagent_templates + self._agent_cls = custom_agent_cls or Agent + self._projection = SessionProjection(message_bus) + self._projectors: list[EventProjector] = [ + SubagentHitlProjector(storage), + *(extra_projectors or []), + ] + + async def run( + self, + user_id: str, + session_id: str, + agent_id: str, + input_msg: Msg + | list[Msg] + | UserConfirmResultEvent + | ExternalExecutionResultEvent + | None = None, + ) -> None: + """Drive a chat run to completion. + + Persists input messages (Case A) or the incoming continuation + event applied to the existing reply (Case B), runs the agent + while publishing every produced event to the message bus, and + persists the rebuilt reply ``Msg`` + updated agent state when + finished. + + Session serialisation is handled by the bus's distributed lock + (:meth:`MessageBus.session_run`); events are simultaneously + persisted to the replay log and fanned out on the live channel + via :meth:`MessageBus.session_publish_event`. Exceptions are + logged and swallowed so a single failed fire does not tear + down its trigger (HTTP request task, wakeup dispatcher, …). + + Args: + user_id (`str`): + Authenticated caller's user ID. + session_id (`str`): + Target session ID. + agent_id (`str`): + Agent to run. + input_msg: + One of: + + - ``Msg`` / ``list[Msg]``: new user message(s) (Case A). + - ``None``: continue from current state — used by the + wakeup dispatcher when there is no fresh user input + but pending inbox content needs draining (Case A + with no input). + - ``UserConfirmResultEvent`` / + ``ExternalExecutionResultEvent``: resume an awaiting + tool call (Case B). + """ + try: + await self._run_impl(user_id, session_id, agent_id, input_msg) + except Exception as e: + logger.exception( + "ChatService.run failed for user_id=%s session_id=%s " + "agent_id=%s, error=%s", + user_id, + session_id, + agent_id, + str(e), + ) + + async def _run_impl( + self, + user_id: str, + session_id: str, + agent_id: str, + input_msg: Msg + | list[Msg] + | UserConfirmResultEvent + | ExternalExecutionResultEvent + | None, + ) -> None: + """The actual chat-run body; wrapped by :meth:`run` for error + swallowing. Separated so the try/except doesn't bury the + per-step logic at one extra indentation level.""" + + # ---------------------------------------------------------------- + # 1. Load records + resolve workspace ONCE here, reused below. + # Reject missing records up front with a clear error so the + # downstream assembly code can rely on non-None values. + # ---------------------------------------------------------------- + agent_record = await self._storage.get_agent(user_id, agent_id) + if agent_record is None: + raise HTTPException( + status_code=404, + detail=f"Agent {agent_id!r} not found.", + ) + session_record = await self._storage.get_session( + user_id, + agent_id, + session_id, + ) + if session_record is None: + raise HTTPException( + status_code=404, + detail=( + f"Session {session_id!r} not found for " + f"agent {agent_id!r}." + ), + ) + workspace = await self._workspace_manager.get_workspace( + user_id, + agent_id, + session_id, + session_record.config.workspace_id, + ) + + # Add workspace working directory to the permission context + if ( + workspace.workdir + not in session_record.state.permission_context.working_directories + ): + session_record.state.permission_context.working_directories[ + workspace.workdir + ] = AdditionalWorkingDirectory( + path=workspace.workdir, + source="session", + ) + + # ---------------------------------------------------------------- + # 2. Middlewares — framework-supplied first, then caller extras. + # Background-tool completions deliver their results via + # ``message_bus.inbox_push + enqueue_wakeup``, so the dispatcher + # (any process) wakes an idle session — no in-process retrigger + # plumbing is needed here. + # ---------------------------------------------------------------- + middlewares: list = [ + InboxMiddleware(self._message_bus), + StateChangeMiddleware( + message_bus=self._message_bus, + session_id=session_id, + ), + ToolOffloadMiddleware( + bg_manager=self._background_task_manager, + message_bus=self._message_bus, + user_id=user_id, + agent_id=agent_id, + ), + ] + if self._extra_agent_middlewares is not None: + middlewares.extend( + await self._extra_agent_middlewares( + user_id, + agent_id, + session_id, + ), + ) + + # ---------------------------------------------------------------- + # 2b. TTS middleware — inject when the session has a TTS config. + # ---------------------------------------------------------------- + tts_cfg = session_record.config.tts_model_config + if tts_cfg is not None: + tts_model = await get_tts_model( + user_id, + tts_cfg, + self._storage, + ) + middlewares.append(TTSMiddleware(tts_model)) + + # ---------------------------------------------------------------- + # 2c. Knowledge-base middleware — inject when the session has KBs + # attached. Each KB resolves to its own :class:`KnowledgeBase` handle + # (own embedding model + vector store), so the middleware can + # retrieve across heterogeneous KBs in one fan-out. + # ---------------------------------------------------------------- + kb_cfg = session_record.config.knowledge_config + if ( + kb_cfg is not None + and kb_cfg.knowledge_base_ids + and self._knowledge_base_manager is not None + ): + knowledges: list[KnowledgeBase] = [] + for kb_id in kb_cfg.knowledge_base_ids: + try: + knowledge = ( + await self._knowledge_base_manager.get_knowledge( + user_id, + kb_id, + ) + ) + except Exception: # pylint: disable=broad-except + # A KB the session referenced was deleted (or its + # credential revoked) — log and skip so the chat + # turn can still run with the remaining KBs. + logger.exception( + "Skipping knowledge base %r for session %r: " + "failed to resolve runtime handle.", + kb_id, + session_id, + ) + continue + knowledges.append(knowledge) + if knowledges: + middlewares.append( + RAGMiddleware( + knowledge_bases=knowledges, + parameters=RAGMiddleware.Parameters( + **(kb_cfg.parameters or {}), + ), + ), + ) + + # ---------------------------------------------------------------- + # 3. Toolkit (workspace tools + planning + ToolStop + schedule + + # team + extras + skills + mcps). + # ---------------------------------------------------------------- + toolkit = await get_toolkit( + storage=self._storage, + workspace=workspace, + scheduler_manager=self._scheduler_manager, + background_task_manager=self._background_task_manager, + message_bus=self._message_bus, + middlewares=middlewares, + user_id=user_id, + agent_record=agent_record, + session_record=session_record, + extra_factory=self._extra_agent_tools, + sub_agent_templates=self._sub_agent_templates, + ) + + # ---------------------------------------------------------------- + # 4. Model + fallback (resolved from session's config). + # ---------------------------------------------------------------- + model_cfg = session_record.config.chat_model_config + if not model_cfg: + raise HTTPException( + status_code=404, + detail=f"No model configuration found for agent {agent_id}", + ) + model = await get_model(user_id, model_cfg, self._storage) + + fallback_cfg = session_record.config.fallback_chat_model_config + fallback_model = ( + await get_model(user_id, fallback_cfg, self._storage) + if fallback_cfg is not None + else None + ) + + # ---------------------------------------------------------------- + # 5. Assemble the Agent. + # ---------------------------------------------------------------- + agent_state = session_record.state + agent_state.session_id = session_id + agent = self._agent_cls( + name=agent_record.data.name, + system_prompt=agent_record.data.system_prompt, + model=model, + toolkit=toolkit, + model_config=ModelConfig(fallback_model=fallback_model), + context_config=agent_record.data.context_config, + react_config=agent_record.data.react_config, + state=agent_state, + middlewares=middlewares, + offloader=workspace, + ) + + # ---------------------------------------------------------------- + # 6. Guard: skip wake-up driven runs when the agent is parked on + # an awaiting tool call. + # + # Wake-ups deliver pending inbox content (team messages, etc.) by + # poking the dispatcher to run the session with ``input_msg=None``. + # If the agent is currently parked on an ``ASKING`` or + # ``SUBMITTED`` tool call (waiting for user confirmation or + # external-execution results), kicking off another ``None`` run + # would hit :meth:`Agent._check_incoming_event`, which rightly + # rejects ``None`` when there is something to confirm — and fail + # the run noisily. The inbox content is safe to leave queued: + # whenever the user does confirm (or the external result lands), + # the resuming run's next reasoning step lets + # :class:`InboxMiddleware` drain the queue naturally. + # ---------------------------------------------------------------- + if input_msg is None and agent.state.context: + last_msg = agent.state.context[-1] + if last_msg.role == "assistant" and last_msg.name == agent.name: + awaiting = [ + tc + for tc in last_msg.get_content_blocks("tool_call") + if tc.state + in (ToolCallState.ASKING, ToolCallState.SUBMITTED) + ] + if awaiting: + logger.info( + "Skipping wake-up for session %s: agent is parked " + "on %d awaiting tool call(s); inbox messages will " + "be drained when the agent resumes.", + session_id, + len(awaiting), + ) + return + + # ---------------------------------------------------------------- + # 7. Run the agent inside the distributed session lock + # ---------------------------------------------------------------- + lock_key = MessageBusKeys.session_lock(session_id) + events_key = MessageBusKeys.session_events(session_id) + async with self._message_bus.acquire_lock( + lock_key, + ttl_secs=MessageBusKeys.SESSION_RUN_TTL_SECS, + ): + try: + reply_msg: Msg | None = None + + if input_msg is None or isinstance(input_msg, (Msg, list)): + # Case A: new reply (user message(s), or retrigger with + # empty input) + if isinstance(input_msg, (Msg, list)): + input_msgs = ( + [input_msg] + if isinstance(input_msg, Msg) + else input_msg + ) + for msg in input_msgs: + await self._storage.upsert_message( + user_id, + session_id, + msg, + ) + + async for event in agent.reply_stream(inputs=input_msg): + await publish_session_event( + self._message_bus, + session_id, + event.model_dump(mode="json"), + ) + await self._project_event( + user_id, + session_record, + agent_record, + event, + ) + if isinstance(event, ReplyStartEvent): + reply_msg = AssistantMsg( + id=event.reply_id, + name=event.name, + content=[], + ) + elif reply_msg is not None: + reply_msg.append_event(event) + + else: + # Case B: continuation (UserConfirmResult + # / ExternalExecResult) + reply_msg = await self._storage.get_message( + user_id, + session_id, + agent.state.reply_id, + ) + + if reply_msg is None: + logger.warning( + "Reply message %r not found in storage for " + "session %r; tool-call state changes from the " + "incoming event will not be persisted.", + agent.state.reply_id, + session_id, + ) + elif input_msg: + reply_msg.append_event(input_msg) + + async for event in agent.reply_stream(inputs=input_msg): + await publish_session_event( + self._message_bus, + session_id, + event.model_dump(mode="json"), + ) + await self._project_event( + user_id, + session_record, + agent_record, + event, + ) + if reply_msg is not None: + reply_msg.append_event(event) + + # Persist the reply Msg (upsert: overwrite if same id, + # append if new). + if reply_msg is not None: + await self._storage.upsert_message( + user_id, + session_id, + reply_msg, + ) + + # Persist the updated agent state. MUST happen inside + # the session lock: if we released the lock first, + # another process could acquire it and load a stale + # state from storage before this write lands. + await self._storage.update_session_state( + user_id=user_id, + agent_id=agent_id, + session_id=session_id, + state=agent.state, + ) + finally: + await self._message_bus.log_trim(events_key) + + async def _project_event( + self, + user_id: str, + session_record: SessionRecord, + agent_record: AgentRecord, + event: AgentEvent, + ) -> None: + """Run every registered projector against one produced event. + + Each :class:`~agentscope.app._types.EventProjector` decides + whether the event is relevant to its cross-session UI feed and, + if so, mirrors it onto the owning session via the shared + :class:`SessionProjection`. Projectors are independent: one + failing must neither tear down the producing run nor block the + others, so each call is guarded individually and its error + logged. Adding a feed means adding a projector — no change here. + + Args: + user_id (`str`): + The owner user id. + session_record (`SessionRecord`): + The currently-running session's record. + agent_record (`AgentRecord`): + The currently-running agent's record. + event (`AgentEvent`): + The event just published to this session's channel. + """ + for projector in self._projectors: + try: + await projector.maybe_project( + user_id, + session_record, + agent_record, + event, + self._projection, + ) + except Exception as e: # pylint: disable=broad-except + logger.warning( + "Projector %s failed on event %s from session %s: %s", + type(projector).__name__, + type(event).__name__, + session_record.id, + str(e), + ) diff --git a/src/agentscope/app/_service/_embedding.py b/src/agentscope/app/_service/_embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..b32c697bac6e08925813d93b3f1a821c8e7821df --- /dev/null +++ b/src/agentscope/app/_service/_embedding.py @@ -0,0 +1,101 @@ +# -*- coding: utf-8 -*- +"""Embedding model service: builds an EmbeddingModelBase from stored +credential + config. + +Mirrors :mod:`._model` (which does the same for chat models). +""" +from fastapi import HTTPException, status + +from ..storage import StorageBase, EmbeddingModelConfig +from ...credential import CredentialFactory +from ...embedding import EmbeddingModelBase + + +async def get_embedding_model( + user_id: str, + config: EmbeddingModelConfig, + storage: StorageBase, +) -> EmbeddingModelBase: + """Construct an embedding model from a stored credential and config. + + This is the embedding counterpart of + :func:`~agentscope.app._service._model.get_model`. It loads the + user's credential from storage, resolves the matching embedding + model class, looks up the model card for ``context_size``, and + constructs a ready-to-use instance. + + Args: + user_id (`str`): + The authenticated user id (credential owner). + config (`EmbeddingModelConfig`): + The embedding model configuration containing + ``type``, ``credential_id``, ``model``, and + ``parameters``. + storage (`StorageBase`): + The storage backend for loading credentials. + + Returns: + `EmbeddingModelBase`: + A configured embedding model instance. + + Raises: + `HTTPException`: + 404 if the credential is not found. + 400 if the provider does not support embedding. + """ + # 1. Load credential from storage. + credential_record = await storage.get_credential( + user_id, + config.credential_id, + ) + if credential_record is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Credential {config.credential_id!r} not found.", + ) + + credential = CredentialFactory.from_dict(credential_record.data) + + # 2. Resolve the embedding model class from the credential type. + credential_cls = CredentialFactory.get_credential_class(config.type) + if credential_cls is None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Provider {config.type!r} not found.", + ) + + embedding_cls = credential_cls.get_embedding_model_class() + if embedding_cls is None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=( + f"Provider {config.type!r} does not support " + f"embedding models." + ), + ) + + # 3. Look up the model card for context_size. + context_size: int | None = None + for card in embedding_cls.list_models(): + if card.name == config.model: + context_size = card.context_size + break + + # 4. Build parameters (provider-specific, no dimensions). + parameters = ( + embedding_cls.Parameters(**config.parameters) + if config.parameters + else None + ) + + # 5. Construct the model — dimensions is first-class, not in parameters. + kwargs: dict = { + "credential": credential, + "model": config.model, + "dimensions": config.dimensions, + "parameters": parameters, + } + if context_size is not None: + kwargs["context_size"] = context_size + + return embedding_cls(**kwargs) diff --git a/src/agentscope/app/_service/_index_sweeper.py b/src/agentscope/app/_service/_index_sweeper.py new file mode 100644 index 0000000000000000000000000000000000000000..cff65dc55f4bd3e4ceb5fc7f3024f19df228170e --- /dev/null +++ b/src/agentscope/app/_service/_index_sweeper.py @@ -0,0 +1,150 @@ +# -*- coding: utf-8 -*- +"""Background sweep for stuck knowledge-document indexing jobs. + +The indexing pipeline relies on two storage-level signals to keep +moving when something goes wrong: + +- a *lease* per in-flight document — its ``lease_expires_at`` is the + upper bound on how long a worker may sit on the document before + another worker is allowed to take over; +- a *creation timestamp* on every ``pending`` record — used to catch + documents that were never picked up by a worker (e.g. process died + right after the upload endpoint persisted the record). + +The sweeper periodically scans storage for both classes of stuck +records and re-enqueues them on the index-task channel. Re-enqueue is +safe because the worker's CAS lease acquisition rejects duplicates, +so multiple nodes running their own sweeper does not produce double +processing. +""" +import asyncio +from datetime import datetime, timedelta +from typing import TYPE_CHECKING + +from ..._logging import logger +from .._bus_ops import enqueue_index_task + +if TYPE_CHECKING: + from ..message_bus import MessageBus + from ..storage import StorageBase + + +class IndexSweeper: + """Periodically re-enqueues documents stuck in indexing. + + Lifecycle is wired into the app's lifespan: :meth:`start` schedules + the background task and runs an immediate sweep so that documents + left stuck by the previous process generation get picked up at + once; :meth:`stop` cancels the loop on shutdown. + """ + + def __init__( + self, + storage: "StorageBase", + message_bus: "MessageBus", + interval: timedelta = timedelta(seconds=60), + pending_grace: timedelta = timedelta(minutes=5), + ) -> None: + """Initialize the sweeper. + + Args: + storage (`StorageBase`): + Used to find stuck records and as the contract holder + for the lease semantics. + message_bus (`MessageBus`): + The same bus the upload endpoint uses. Re-enqueuing + a document re-enters the worker pipeline, where the + CAS lease acquisition decides whether to actually + process or bail. + interval (`timedelta`, defaults to ``60s``): + How often the loop wakes up. Roughly one order of + magnitude shorter than the typical lease TTL — fast + enough to recover from crashes within a few minutes, + slow enough not to thrash storage. + pending_grace (`timedelta`, defaults to ``5min``): + A record may legitimately sit in ``pending`` while the + bus push is still queued; only after this grace period + do we treat the record as orphaned. + """ + self._storage = storage + self._bus = message_bus + self._interval = interval + self._pending_grace = pending_grace + self._task: asyncio.Task[None] | None = None + + async def start(self) -> None: + """Start the background sweep loop and run one immediate sweep.""" + if self._task is not None: + return + # Catch up from any state the previous generation left behind. + await self._sweep_once() + self._task = asyncio.create_task( + self._loop(), + name="kb-index-sweeper", + ) + + async def stop(self) -> None: + """Cancel the sweep loop and wait for it to exit.""" + if self._task is None: + return + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + self._task = None + + async def _loop(self) -> None: + """Run sweeps forever until cancelled.""" + interval_seconds = self._interval.total_seconds() + while True: + try: + await asyncio.sleep(interval_seconds) + except asyncio.CancelledError: + return + try: + await self._sweep_once() + except asyncio.CancelledError: + return + except Exception: # noqa: BLE001 — keep the loop alive + logger.exception("Sweep iteration failed") + + async def _sweep_once(self) -> None: + """Find and re-enqueue every stuck document. + + De-duplication: a document showing up in both the + expired-lease and orphan-pending queries (a record that was + never picked up and whose lease pre-dates the grace period) + is enqueued only once per sweep, by record id. + """ + now = datetime.now() + pending_threshold = now - self._pending_grace + + seen: set[str] = set() + stuck = ( + await self._storage.list_knowledge_documents_with_expired_lease( + now=now, + ) + ) + orphans = await self._storage.list_knowledge_documents_pending_since( + threshold=pending_threshold, + ) + for record in (*stuck, *orphans): + if record.id in seen: + continue + seen.add(record.id) + try: + await enqueue_index_task( + self._bus, + user_id=record.user_id, + knowledge_base_id=record.knowledge_base_id, + document_id=record.id, + ) + except Exception: # noqa: BLE001 — keep iterating + logger.exception( + "Failed to re-enqueue document %s", + record.id, + ) + + if seen: + logger.info("Re-enqueued %d stuck document(s)", len(seen)) diff --git a/src/agentscope/app/_service/_index_task_consumer.py b/src/agentscope/app/_service/_index_task_consumer.py new file mode 100644 index 0000000000000000000000000000000000000000..f05090a5e1d40f6a7bc1965e4a80cc4dd5fc85d6 --- /dev/null +++ b/src/agentscope/app/_service/_index_task_consumer.py @@ -0,0 +1,216 @@ +# -*- coding: utf-8 -*- +"""Single per-worker-process consumer of the shared index-task channel. + +One asyncio task per worker process. Subscribes to the shared +:meth:`~agentscope.app.message_bus.MessageBusKeys.index_tasks_signal` +channel and drains the durable +:meth:`~agentscope.app.message_bus.MessageBusKeys.index_tasks_queue` +on each signal. For each queued entry it invokes +:meth:`IndexWorker.process` directly — the worker holds its own +semaphore so we can fire-and-forget multiple ``process`` calls without +overrunning resources. + +Mirrors :class:`~agentscope.app._manager.WakeupDispatcher`. The two +patterns are deliberately identical: both subscribe to a signal, +drain a queue, dispatch each entry, and run forever inside an +``async with`` block. Keeping them shaped the same makes it cheap to +reason about either one once you've read the other. + +The bus exposes only transport-level primitives — there is no +``enqueue_index_task`` or ``dequeue_index_task`` method on it. The +key constants live on :class:`~agentscope.app.message_bus. +MessageBusKeys` (next to every other application-layer key) and the +composition is inline here because the consumer is the only sink for +the channel; introducing a separate ``IndexTaskBroker`` would be +ceremony without gain. +""" +import asyncio +from typing import TYPE_CHECKING, Any, Self + +from ..message_bus import MessageBusKeys +from ..._logging import logger + +if TYPE_CHECKING: + from ..message_bus import MessageBus + from ._index_worker import IndexWorker + + +class IndexTaskConsumer: + """Subscribe-then-drain consumer that feeds :class:`IndexWorker`. + + Args: + message_bus (`MessageBus`): + Application message bus. The consumer only uses the two + transport-level primitives — ``subscribe`` (for the + signal channel) and ``queue_drain`` (for the durable + task queue). + worker (`IndexWorker`): + The worker that owns the parse → chunk → index pipeline. + ``process`` is invoked once per queue entry; the worker's + internal semaphore + lease CAS handle concurrency and + deduplication. + max_batch (`int`, defaults to ``32``): + Maximum entries drained per signal. Keeps a single + signal from monopolising the loop when the queue is + backed up; remaining entries are picked up on the next + signal or the next sweeper-driven eager drain. + """ + + def __init__( + self, + message_bus: "MessageBus", + worker: "IndexWorker", + max_batch: int = 32, + ) -> None: + self._bus = message_bus + self._worker = worker + self._max_batch = max_batch + self._task: asyncio.Task | None = None + # In-flight ``worker.process`` calls. Tracked so ``__aexit__`` + # can cancel + drain them; otherwise the event-loop teardown + # would swallow exceptions raised inside the worker. + self._inflight: set[asyncio.Task[Any]] = set() + + async def __aenter__(self) -> Self: + """Start the consumer loop and wait until its subscription + is live. + + After the subscription is established, an initial drain runs + synchronously so tasks queued while every worker was down + get picked up immediately on startup, without waiting for + a fresh signal. + """ + ready = asyncio.Event() + self._task = asyncio.create_task( + self._loop(ready), + name="index-task-consumer", + ) + await ready.wait() + await self._drain_and_dispatch() + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: Any, + ) -> None: + """Cancel the consumer loop and drain any in-flight work. + + Cancellation of the worker's ``process`` calls is a clean + shutdown signal — the worker holds the storage lease and + will let it expire so the sweeper re-dispatches the document + on the next loop tick. + """ + if self._task is not None: + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + self._task = None + + for task in list(self._inflight): + task.cancel() + if self._inflight: + await asyncio.gather(*self._inflight, return_exceptions=True) + self._inflight.clear() + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + async def _loop(self, ready: asyncio.Event) -> None: + """Long-lived loop: subscribe to the signal and drain on each + received signal. + + Args: + ready (`asyncio.Event`): + Signalled after the underlying SUBSCRIBE completes. + :meth:`__aenter__` blocks on this so the producer + can publish a signal immediately after start-up + without racing. + """ + try: + async for _signal in self._bus.subscribe( + MessageBusKeys.index_tasks_signal(), + on_ready=ready.set, + ): + await self._drain_and_dispatch() + except Exception: # pylint: disable=broad-except + logger.exception( + "IndexTaskConsumer loop crashed; subscription ended.", + ) + finally: + # If ``subscribe`` raises before ``on_ready`` fires, the + # ``__aenter__`` coroutine would deadlock on ``ready.wait()``. + # Set the event unconditionally on the way out so startup + # cannot stall on a transient bus failure. + ready.set() + + async def _drain_and_dispatch(self) -> None: + """Read up to a batch of task entries and dispatch each one.""" + try: + entries = await self._bus.queue_drain( + MessageBusKeys.index_tasks_queue(), + max_count=self._max_batch, + ) + except Exception: # pylint: disable=broad-except + logger.exception("IndexTaskConsumer: drain failed.") + return + + for _entry_id, payload in entries: + try: + user_id = payload["user_id"] + knowledge_base_id = payload["knowledge_base_id"] + document_id = payload["document_id"] + except (KeyError, TypeError): + logger.warning( + "IndexTaskConsumer: skipping malformed entry %r", + payload, + ) + continue + + self._spawn( + user_id=user_id, + knowledge_base_id=knowledge_base_id, + document_id=document_id, + ) + + def _spawn( + self, + *, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> None: + """Run :meth:`IndexWorker.process` as a tracked background task. + + We do not ``await`` ``worker.process`` inline — a slow parse + would block draining the next signal. The worker holds its + own concurrency semaphore, so spawning many tasks at once is + safe; they will queue at the semaphore. + """ + task = asyncio.create_task( + self._worker.process( + user_id=user_id, + knowledge_base_id=knowledge_base_id, + document_id=document_id, + ), + name=f"index-task:{knowledge_base_id}:{document_id}", + ) + self._inflight.add(task) + task.add_done_callback(self._on_done) + + def _on_done(self, task: asyncio.Task[Any]) -> None: + """Drop the task reference and log any uncaught exception.""" + self._inflight.discard(task) + if task.cancelled(): + return + exc = task.exception() + if exc is not None: + logger.exception( + "IndexTaskConsumer: worker.process(%s) raised", + task.get_name(), + exc_info=exc, + ) diff --git a/src/agentscope/app/_service/_index_worker.py b/src/agentscope/app/_service/_index_worker.py new file mode 100644 index 0000000000000000000000000000000000000000..387676641f7f7877c1a7056d83303bc37af87650 --- /dev/null +++ b/src/agentscope/app/_service/_index_worker.py @@ -0,0 +1,534 @@ +# -*- coding: utf-8 -*- +"""Background indexing pipeline for one knowledge document. + +The :class:`IndexWorker` owns the post-upload half of the document +lifecycle. Given a ``document_id`` it: + +1. acquires the processing lease via storage CAS (so only one worker + in the cluster handles the document at a time); +2. reads the bytes back from the blob store (streamed); +3. routes to a parser by IANA media type; +4. chunks the resulting sections; +5. embeds + writes to the vector store through + :class:`~agentscope.rag.KnowledgeBase`; +6. transitions the status through ``parsing → chunking → indexing → + ready`` (or ``error``) on the way. + +The worker is intentionally embeddable: a single instance can live +inside the API process (embedded deployment) or inside a dedicated +worker process (dedicated deployment). Coordination across workers +is done entirely through the storage lease — workers do not need to +know about each other. +""" +import asyncio +import contextlib +import mimetypes +from concurrent.futures import ProcessPoolExecutor +from datetime import timedelta +from typing import TYPE_CHECKING + +from ..._logging import logger + +if TYPE_CHECKING: + from ..rag.blob_store import BlobStoreBase + from ..rag.knowledge_base_manager import KnowledgeBaseManagerBase + from ..storage import StorageBase + from ...rag import ChunkerBase, ParserBase, Section + +# Read blob bytes in chunks bounded so the worker never holds the whole +# file in memory at once even when the parser is byte-oriented. +_READ_CHUNK = 1 << 20 # 1 MiB + + +def _build_parser_registry( + parsers: "list[ParserBase] | dict[str, ParserBase]", +) -> "dict[str, ParserBase]": + """Normalise the user-supplied parser registry. + + Two input shapes are accepted: + + - **List** — each parser's ``supported_media_types`` is expanded; + duplicate media types resolve to the **last** parser in the list + (so callers can layer custom parsers over the defaults), with a + warning logged for each override so silent shadowing is not + possible. + - **Dict** — the caller's mapping is used verbatim. This is the + escape hatch for callers who want full control over routing (one + parser bound to multiple types, type aliases, etc.); a warning is + logged when a parser is registered against a media type it does + not declare in ``supported_media_types``, since that almost + always indicates a typo. + + Args: + parsers (`list[ParserBase] | dict[str, ParserBase]`): + The user-supplied parser registry. + + Returns: + `dict[str, ParserBase]`: + The resolved ``media_type → parser`` routing table. + """ + if isinstance(parsers, dict): + for media_type, parser in parsers.items(): + declared = getattr(parser, "supported_media_types", ()) + if declared and media_type not in declared: + logger.warning( + "Parser %s registered for media type %r but it only " + "declares %s — proceeding with the caller-supplied " + "mapping.", + type(parser).__name__, + media_type, + list(declared), + ) + return dict(parsers) + + registry: dict[str, "ParserBase"] = {} + for parser in parsers: + for media_type in parser.supported_media_types: + previous = registry.get(media_type) + if previous is not None and previous is not parser: + logger.warning( + "Parser %s overrides %s for media type %r — later " + "entries in `parsers` win. Pass a " + "`dict[str, ParserBase]` if you want explicit " + "routing.", + type(parser).__name__, + type(previous).__name__, + media_type, + ) + registry[media_type] = parser + return registry + + +class IndexWorker: + """Drive one document through parse → chunk → index. + + Multiple invocations of :meth:`process` are run concurrently up to + a per-worker semaphore. The semaphore protects shared resources + that scale with the number of in-flight parses (memory for big + PDFs, embedding API rate budget), while the lease CAS in storage + protects against the *cross-worker* version of the same race. + """ + + def __init__( + self, + storage: "StorageBase", + blob_store: "BlobStoreBase", + knowledge_base_manager: "KnowledgeBaseManagerBase", + parsers: "list[ParserBase] | dict[str, ParserBase]", + chunker: "ChunkerBase", + node_id: str, + max_concurrency: int = 4, + lease_ttl: timedelta = timedelta(seconds=90), + parser_executor: ProcessPoolExecutor | None = None, + ) -> None: + """Initialize the worker. + + Args: + storage (`StorageBase`): + Document records, lease, status. + blob_store (`BlobStoreBase`): + Source of the document bytes. + knowledge_base_manager (`KnowledgeBaseManagerBase`): + Resolves the :class:`KnowledgeBase` runtime for embedding + and vector store writes. + parsers (`list[ParserBase] | dict[str, ParserBase]`): + Parsers used to dispatch uploads by IANA media type. + Two input shapes are accepted: + + - **List** — each parser's ``supported_media_types`` is + expanded into a routing table; later entries override + earlier ones for overlapping types, with a warning + logged at construction time. + - **Dict** — caller-supplied ``media_type → parser`` + routing table used verbatim. ``supported_media_types`` + is **not** consulted, but a warning is logged if a + parser is registered against a media type it does not + declare. + + Same registry the upload service uses, passed in by DI. + chunker (`ChunkerBase`): + The shared chunker. + node_id (`str`): + Stable identifier for this worker process. Used as + ``processing_node`` on the lease so the sweeper can + tell whose work expired. Typically + ``f"{hostname}:{pid}:{uuid}"``. + max_concurrency (`int`, defaults to ``4``): + Maximum number of documents processed concurrently by + this worker. Higher values trade memory for + throughput; tune per embedding-API rate limits and + per-document parse cost. + lease_ttl (`timedelta`, defaults to ``90s``): + How long a single processing lease lives. The worker + renews periodically so long-running parses do not + trip the sweeper. + parser_executor (`ProcessPoolExecutor | None`, optional): + Process pool used to off-load CPU-intensive parses + (PDF, Office). ``None`` runs parses in the event-loop + thread, which is fine for plain text but unsafe for + third-party byte-oriented parsers. Injected so a + single pool can be shared across the app (built in + lifespan). + """ + self._storage = storage + self._blob_store = blob_store + self._manager = knowledge_base_manager + self._parsers_by_media_type = _build_parser_registry(parsers) + self._chunker = chunker + self._node_id = node_id + self._lease_ttl = lease_ttl + self._sem = asyncio.Semaphore(max_concurrency) + self._parser_executor = parser_executor + # Renewal cadence: refresh while there is still half the lease + # left so a one-cycle missed renewal doesn't drop the lease. + self._renew_interval = max(lease_ttl / 2, timedelta(seconds=5)) + + async def process( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> None: + """Run the full indexing pipeline for one document. + + Steps: + + 1. **Lease** — CAS-acquire the processing lease; bail if some + other worker already holds it (duplicate dispatch / sweep). + 2. **Throttle** — wait on the per-worker semaphore so the + number of in-flight parses stays bounded. + 3. **Pipeline** — parse → chunk → embed + write vector store, + updating status before each phase. A background heartbeat + keeps the lease alive while parsing runs. + 4. **Finalise** — on success mark ``ready`` with the final + chunk count; on failure mark ``error`` with a sanitised + message. The lease is released regardless. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The parent knowledge base id. + document_id (`str`): + The document to process. + """ + acquired = await self._storage.acquire_knowledge_document_lease( + user_id=user_id, + knowledge_base_id=knowledge_base_id, + document_id=document_id, + processing_node=self._node_id, + lease_ttl=self._lease_ttl, + ) + if not acquired: + logger.debug( + "Skipping %s — another worker holds the lease.", + document_id, + ) + return + + pipeline_task = asyncio.create_task( + self._guarded_pipeline( + user_id, + knowledge_base_id, + document_id, + ), + name=f"pipeline:{document_id}", + ) + heartbeat_task = asyncio.create_task( + self._heartbeat(user_id, knowledge_base_id, document_id), + name=f"lease-renew:{document_id}", + ) + try: + # Race the pipeline against the heartbeat: if the heartbeat + # returns first, the lease was stolen mid-flight (e.g. the + # sweeper reaped this worker after a renewal gap) — we MUST + # stop the pipeline before it writes the vector store again, + # otherwise the worker that just took over and this one will + # both insert the same chunks. + await asyncio.wait( + {pipeline_task, heartbeat_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + + if not pipeline_task.done(): + # Heartbeat reached the end first; only `_heartbeat`'s + # lost-lease branch returns, so cancel the pipeline and + # surface it as a terminal error for this document. + pipeline_task.cancel() + with contextlib.suppress( + asyncio.CancelledError, + Exception, + ): + await pipeline_task + raise RuntimeError( + f"Lost lease on {document_id} during processing; " + "another worker has taken over.", + ) + + # Pipeline finished first; stop the heartbeat and re-raise + # whatever the pipeline raised (if anything). + heartbeat_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await heartbeat_task + await pipeline_task + except Exception as exc: # noqa: BLE001 — terminal error sink + await self._mark_error( + user_id, + knowledge_base_id, + document_id, + exc, + ) + finally: + # Release is CAS-guarded server-side on ``processing_node`` + # (storage._base.release_knowledge_document_lease) — calling + # it after a stolen lease is a safe no-op. + await self._storage.release_knowledge_document_lease( + user_id=user_id, + knowledge_base_id=knowledge_base_id, + document_id=document_id, + processing_node=self._node_id, + ) + + async def _guarded_pipeline( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> None: + """Run the throttled pipeline inside the per-worker semaphore.""" + async with self._sem: + await self._run_pipeline( + user_id, + knowledge_base_id, + document_id, + ) + + # ------------------------------------------------------------------ + # Pipeline + # ------------------------------------------------------------------ + + async def _run_pipeline( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> None: + """Walk the document through parse → chunk → index.""" + record = await self._storage.get_knowledge_document( + user_id, + knowledge_base_id, + document_id, + ) + if record is None: + logger.warning( + "Document %s vanished before processing.", + document_id, + ) + return + + data = record.data + media_type = ( + data.content_type or mimetypes.guess_type(data.filename)[0] + ) + if not media_type: + raise ValueError( + f"Cannot determine media type for {data.filename!r}.", + ) + parser = self._parsers_by_media_type.get(media_type) + if parser is None: + raise ValueError( + f"No parser registered for media type {media_type!r}.", + ) + + # ---- parsing ---- + await self._storage.update_knowledge_document_status( + user_id, + knowledge_base_id, + document_id, + "parsing", + ) + file_bytes = await self._read_blob(data.blob_uri) + sections = await self._parse(parser, file_bytes, data.filename) + + # ---- chunking ---- + await self._storage.update_knowledge_document_status( + user_id, + knowledge_base_id, + document_id, + "chunking", + ) + chunks = await self._chunker.chunk(sections) + + # ---- indexing ---- + await self._storage.update_knowledge_document_status( + user_id, + knowledge_base_id, + document_id, + "indexing", + ) + knowledge = await self._manager.get_knowledge( + user_id, + knowledge_base_id, + ) + await knowledge.insert_document( + chunks=chunks, + document_id=document_id, + document_metadata={ + "filename": data.filename, + "media_type": media_type, + "size_bytes": data.size, + }, + ) + + # ---- ready ---- + await self._storage.update_knowledge_document_status( + user_id, + knowledge_base_id, + document_id, + "ready", + chunk_count=len(chunks), + ) + + async def _parse( + self, + parser: "ParserBase", + file_bytes: bytes, + filename: str, + ) -> "list[Section]": + """Run the parser, optionally on the process pool.""" + if self._parser_executor is None: + return await parser.parse(file_bytes, filename) + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + self._parser_executor, + _run_parser_sync, + parser, + file_bytes, + filename, + ) + + async def _read_blob(self, blob_uri: str) -> bytes: + """Stream the blob into memory in bounded chunks. + + We buffer the whole file before handing it to the parser + because today's parser API is byte-oriented (``parse(file: + bytes, filename: str)``). The read loop still avoids large + single allocations and gives us a single place to upgrade to a + true streaming parser API later — only this method needs to + change. + """ + buffer = bytearray() + async with self._blob_store.open(blob_uri) as fp: + while True: + chunk = await fp.read(_READ_CHUNK) + if not chunk: + break + buffer.extend(chunk) + return bytes(buffer) + + # ------------------------------------------------------------------ + # Lease heartbeat + # ------------------------------------------------------------------ + + async def _heartbeat( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> None: + """Renew the lease in the background while processing runs. + + Two exit paths: + + - The surrounding pipeline finishes first and cancels this + task — silent return via :class:`asyncio.CancelledError`. + - The renewal fails (the sweeper reaped this worker and + another worker now holds the lease). The task **returns + normally** in this case; :meth:`process` is racing this task + against the pipeline and treats a normal return as the + "lost-lease" signal, cancelling the pipeline before it + double-writes the vector store. + + Anything other than ``ok=False`` keeps the loop alive. + """ + interval_seconds = self._renew_interval.total_seconds() + while True: + try: + await asyncio.sleep(interval_seconds) + except asyncio.CancelledError: + return + ok = await self._storage.renew_knowledge_document_lease( + user_id=user_id, + knowledge_base_id=knowledge_base_id, + document_id=document_id, + processing_node=self._node_id, + lease_ttl=self._lease_ttl, + ) + if not ok: + logger.warning( + "Lost lease on %s while processing.", + document_id, + ) + return + + # ------------------------------------------------------------------ + # Error sink + # ------------------------------------------------------------------ + + async def _mark_error( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + exc: BaseException, + ) -> None: + """Persist a sanitised error and mark the document failed.""" + logger.exception( + "Indexing failed for %s/%s", + knowledge_base_id, + document_id, + exc_info=exc, + ) + message = _sanitise_error(exc) + try: + await self._storage.update_knowledge_document_status( + user_id, + knowledge_base_id, + document_id, + "error", + error=message, + ) + except Exception: # noqa: BLE001 — last-resort log + logger.exception( + "Failed to persist error status for %s", + document_id, + ) + + +# ---------------------------------------------------------------------- +# Module-level helpers (picklable for ProcessPoolExecutor) +# ---------------------------------------------------------------------- + + +def _run_parser_sync( + parser: "ParserBase", + file_bytes: bytes, + filename: str, +) -> "list[Section]": + """Run an async parser to completion inside a sync executor.""" + return asyncio.run(parser.parse(file_bytes, filename)) + + +def _sanitise_error(exc: BaseException) -> str: + """Reduce an exception to a single user-facing line. + + Only the exception class name + first line of its message are + kept — stack traces and filesystem paths stay inside the worker + log and out of the user-visible record. + """ + raw = str(exc) or exc.__class__.__name__ + first_line = raw.splitlines()[0].strip() + cls = exc.__class__.__name__ + if not first_line: + return cls + return f"{cls}: {first_line[:240]}" diff --git a/src/agentscope/app/_service/_knowledge_base.py b/src/agentscope/app/_service/_knowledge_base.py new file mode 100644 index 0000000000000000000000000000000000000000..516d85bf407046446646829cdc6fa7e470066bc8 --- /dev/null +++ b/src/agentscope/app/_service/_knowledge_base.py @@ -0,0 +1,515 @@ +# -*- coding: utf-8 -*- +"""Knowledge base service: HTTP-side orchestration. + +The router stays thin and DTO-shaped; everything HTTP-side that needs +to coordinate persistence, the blob store, the indexing pipeline, +and the vector store goes through this service. + +The split with :class:`~agentscope.rag.KnowledgeBase` +is deliberate. ``KnowledgeBase`` is a **library-mode** handle that only +depends on the vector store; embedded users instantiate one and drive +the parse → chunk → embed pipeline themselves. ``KnowledgeBaseService`` +is **service-mode** orchestration: it owns the document records +(status / blob / lease) and is the single source of truth for "what +documents exist in this KB" when the app is running over HTTP. The +two views are intentionally not blended — mixing library-mode inserts +with service-mode listing would leave records out of sync, and the +project's stance is that a knowledge base is managed end-to-end in one +mode. +""" +import uuid +from typing import IO, TYPE_CHECKING + +from fastapi import HTTPException, status + +from ..rag.knowledge_base_manager import ( + DimensionPolicyError, + KnowledgeBaseNotFoundError, +) +from ..storage import ( + KnowledgeDocumentData, + KnowledgeDocumentRecord, +) +from ..._logging import logger +from .._bus_ops import enqueue_index_task + +if TYPE_CHECKING: + from ..rag.blob_store import BlobStoreBase + from ..rag.knowledge_base_manager import KnowledgeBaseManagerBase + from ..message_bus import MessageBus + from ..storage import ( + EmbeddingModelConfig, + KnowledgeBaseRecord, + StorageBase, + ) + from ...rag import VectorSearchResult + + +class KnowledgeBaseService: + """HTTP service for knowledge bases. + + Owns the document lifecycle in service mode: register on upload, + enqueue an index task, query status during indexing, and clean up + record + blob + vector store on delete. All parsing / chunking / + embedding work happens inside the + :class:`~agentscope.app._service.IndexWorker`; the service only + hands off (via the message bus) and observes. + """ + + def __init__( + self, + storage: "StorageBase", + knowledge_base_manager: "KnowledgeBaseManagerBase", + blob_store: "BlobStoreBase", + message_bus: "MessageBus", + ) -> None: + """Initialize the service. + + Args: + storage (`StorageBase`): + The application storage backend; documents are + persisted here, not inside the vector store. + knowledge_base_manager (`KnowledgeBaseManagerBase`): + Resolves the :class:`KnowledgeBase` runtime used to clear + vector store records on document deletion. + blob_store (`BlobStoreBase`): + Owns the bytes from upload until the worker is done. + The service writes on upload and deletes on document + removal. + message_bus (`MessageBus`): + Application message bus. The service publishes one + index-task entry per uploaded document via + :func:`~agentscope.app._bus_ops.enqueue_index_task`; + a co-located or out-of-process + :class:`IndexTaskConsumer` drains and processes them. + """ + self._storage = storage + self._manager = knowledge_base_manager + self._blob_store = blob_store + self._bus = message_bus + + # ------------------------------------------------------------------ + # Knowledge base CRUD + # ------------------------------------------------------------------ + + async def create_knowledge_base( + self, + user_id: str, + name: str, + description: str, + embedding_model_config: "EmbeddingModelConfig", + ) -> "KnowledgeBaseRecord": + """Delegate creation to the manager, mapping policy errors. + + Args: + user_id (`str`): + The owner user id. + name (`str`): + Display name. + description (`str`): + Free-form description. + embedding_model_config (`EmbeddingModelConfig`): + Embedding model configuration; pinned to the record. + + Returns: + `KnowledgeBaseRecord`: + The newly persisted record. + + Raises: + `HTTPException`: + ``409`` when the requested embedding dimension + violates the manager's dimension policy. + """ + try: + return await self._manager.create_knowledge_base( + user_id=user_id, + name=name, + description=description, + embedding_model_config=embedding_model_config, + ) + except DimensionPolicyError as exc: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=str(exc), + ) from exc + + async def list_knowledge_bases( + self, + user_id: str, + ) -> "list[KnowledgeBaseRecord]": + """List all knowledge base records owned by the given user. + + Args: + user_id (`str`): + The owner user id. + + Returns: + `list[KnowledgeBaseRecord]`: + All knowledge base records belonging to the user. + """ + return await self._manager.list_knowledge_bases(user_id) + + async def update_knowledge_base( + self, + user_id: str, + knowledge_base_id: str, + name: str | None = None, + description: str | None = None, + ) -> "KnowledgeBaseRecord": + """Update mutable fields on a knowledge base, raising 404 if absent. + + Only ``name`` and ``description`` are mutable. The embedding + model configuration is pinned at creation time. + """ + record = await self._manager.update_knowledge_base( + user_id=user_id, + knowledge_base_id=knowledge_base_id, + name=name, + description=description, + ) + if record is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Knowledge base {knowledge_base_id!r} not found.", + ) + return record + + async def delete_knowledge_base( + self, + user_id: str, + knowledge_base_id: str, + ) -> None: + """Delete a knowledge base, raising 404 if absent. + + Documents under the KB are cascade-deleted at the storage + layer; blob files referenced by those records are released + best-effort here so disk space is reclaimed even though the + manager + storage cascade would otherwise orphan them. + """ + documents = await self._storage.list_knowledge_documents( + user_id, + knowledge_base_id, + ) + for document in documents: + await self._delete_blob_quietly(document.data.blob_uri) + + deleted = await self._manager.delete_knowledge_base( + user_id, + knowledge_base_id, + ) + if not deleted: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Knowledge base {knowledge_base_id!r} not found.", + ) + + # ------------------------------------------------------------------ + # Document management + # ------------------------------------------------------------------ + + async def register_document( + self, + user_id: str, + knowledge_base_id: str, + filename: str, + stream: IO[bytes], + size: int, + content_type: str | None = None, + ) -> KnowledgeDocumentRecord: + """Persist an uploaded document and enqueue it for indexing. + + Streams ``stream`` into the blob store (so the bytes never + live fully in memory), records a ``pending`` document, and + pushes an index-task entry onto the message bus. Returns + immediately — a worker (in-process or dedicated) takes over + from here and the client tracks progress via + :meth:`get_document_status`. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The target knowledge base id. + filename (`str`): + The original filename. + stream (`IO[bytes]`): + A synchronous binary stream — typically + ``UploadFile.file`` from FastAPI. + size (`int`): + Byte length declared by the uploader. Persisted on + the record for the UI; not authoritative. + content_type (`str | None`, optional): + IANA media type; ``None`` lets the worker fall back + to a filename guess at processing time. + + Returns: + `KnowledgeDocumentRecord`: + The persisted record (``status='pending'``) with the + final ``blob_uri`` filled in. + + Raises: + `HTTPException`: + ``404`` if the knowledge base does not exist. + """ + # Authorise before touching the blob store: raising after a + # write would leave the blob orphaned. + await self._authorise_kb(user_id, knowledge_base_id) + + document_id = uuid.uuid4().hex + blob_uri = await self._blob_store.write_stream( + key=f"kb/{knowledge_base_id}/{document_id}", + stream=stream, + ) + + record = KnowledgeDocumentRecord( + id=document_id, + user_id=user_id, + knowledge_base_id=knowledge_base_id, + data=KnowledgeDocumentData( + filename=filename, + size=size, + content_type=content_type, + blob_uri=blob_uri, + ), + ) + try: + stored = await self._storage.upsert_knowledge_document( + user_id, + record, + ) + except Exception: + # Storage write failed — drop the blob so the orphan + # sweeper doesn't later see a referenced-by-nobody file. + await self._delete_blob_quietly(blob_uri) + raise + + await enqueue_index_task( + self._bus, + user_id=user_id, + knowledge_base_id=knowledge_base_id, + document_id=document_id, + ) + return stored + + async def list_documents( + self, + user_id: str, + knowledge_base_id: str, + ) -> list[KnowledgeDocumentRecord]: + """List every document registered against a knowledge base. + + Service-mode source of truth: reads from storage, NOT the + vector store. Documents in ``pending`` / ``parsing`` / + ``chunking`` / ``indexing`` / ``error`` show up here even + though they have no chunks in the vector store yet. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The target knowledge base id. + + Returns: + `list[KnowledgeDocumentRecord]`: + Every document registered against the knowledge base, + in unspecified order. + + Raises: + `HTTPException`: + ``404`` if the knowledge base does not exist. + """ + await self._authorise_kb(user_id, knowledge_base_id) + return await self._storage.list_knowledge_documents( + user_id, + knowledge_base_id, + ) + + async def get_document_status( + self, + user_id: str, + knowledge_base_id: str, + document_ids: list[str], + ) -> list[KnowledgeDocumentRecord]: + """Batch-fetch documents for status polling. + + The endpoint backing this method accepts a comma-separated list + of ids so the front-end can ask "what's the state of these N + in-flight uploads" in a single round-trip. Records that do + not exist or do not belong to the user are silently skipped — + the front-end may legitimately ask about a document that was + deleted between two polls. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The target knowledge base id. + document_ids (`list[str]`): + Document ids to look up. + + Returns: + `list[KnowledgeDocumentRecord]`: + One record per matched id; missing ids omitted. + + Raises: + `HTTPException`: + ``404`` if the knowledge base does not exist. + """ + await self._authorise_kb(user_id, knowledge_base_id) + records: list[KnowledgeDocumentRecord] = [] + for document_id in document_ids: + record = await self._storage.get_knowledge_document( + user_id, + knowledge_base_id, + document_id, + ) + if record is not None: + records.append(record) + return records + + async def delete_document( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> None: + """Remove a document end-to-end: vector store, record, blob. + + Order is chosen so that a crash mid-way always leaves a + recoverable state: + + 1. Vector store delete (idempotent — re-deleting an already + empty document_id is harmless). + 2. Storage record delete. + 3. Blob delete (idempotent). + + A failure at step 1 surfaces as an exception to the caller and + the record + blob are left untouched, so a retry sees the same + state. Failures at steps 2/3 leave a small amount of orphan + data but the user-visible deletion has already succeeded from + the vector store's point of view. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The target knowledge base id. + document_id (`str`): + The document to delete. + + Raises: + `HTTPException`: + ``404`` if the knowledge base does not exist. + """ + record = await self._storage.get_knowledge_document( + user_id, + knowledge_base_id, + document_id, + ) + if record is None: + # 404 if the KB does not exist, otherwise treat the + # missing document as already-deleted (idempotent). + await self._authorise_kb(user_id, knowledge_base_id) + return + + knowledge = await self._resolve_knowledge(user_id, knowledge_base_id) + await knowledge.delete_document(document_id) + await self._storage.delete_knowledge_document( + user_id, + knowledge_base_id, + document_id, + ) + await self._delete_blob_quietly(record.data.blob_uri) + + # ------------------------------------------------------------------ + # Search + # ------------------------------------------------------------------ + + async def search( + self, + user_id: str, + knowledge_base_id: str, + query: str, + top_k: int = 5, + ) -> "list[VectorSearchResult]": + """Search a knowledge base by text query. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The knowledge base to search. + query (`str`): + The natural-language query. + top_k (`int`, defaults to ``5``): + Maximum number of results. + + Returns: + `list[VectorSearchResult]`: + The top hits ordered by descending similarity score. + + Raises: + `HTTPException`: + ``404`` if the knowledge base does not exist. + """ + knowledge = await self._resolve_knowledge(user_id, knowledge_base_id) + return await knowledge.search(queries=[query], top_k=top_k) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + async def _authorise_kb( + self, + user_id: str, + knowledge_base_id: str, + ) -> "KnowledgeBaseRecord": + """Look the KB record up so we can 404 cleanly. + + The check is intentionally separate from :meth:`_resolve_knowledge` + because document-level endpoints (list / delete) need to refuse + unknown KBs without paying the embedding-model construction cost + that :meth:`_resolve_knowledge` triggers. + """ + record = await self._storage.get_knowledge_base( + user_id, + knowledge_base_id, + ) + if record is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Knowledge base {knowledge_base_id!r} not found.", + ) + return record + + async def _resolve_knowledge( + self, + user_id: str, + knowledge_base_id: str, + ) -> "object": + """Resolve a :class:`KnowledgeBase` and translate not-found to 404.""" + try: + return await self._manager.get_knowledge( + user_id, + knowledge_base_id, + ) + except KnowledgeBaseNotFoundError as exc: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=str(exc), + ) from exc + + async def _delete_blob_quietly(self, blob_uri: str) -> None: + """Best-effort blob delete — swallow backend errors. + + Treated as cleanup: if the blob store is unavailable the + record/vector-store state is still consistent and a future + sweep can reclaim the disk space. Surface only via logs. + """ + try: + await self._blob_store.delete(blob_uri) + except Exception: # noqa: BLE001 — cleanup only + logger.exception( + "Failed to delete blob %s", + blob_uri, + ) diff --git a/src/agentscope/app/_service/_model.py b/src/agentscope/app/_service/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..4a3b7b8f7ee64a74dfbe1ed527ed328327f7c3fb --- /dev/null +++ b/src/agentscope/app/_service/_model.py @@ -0,0 +1,50 @@ +# -*- coding: utf-8 -*- +"""Model service: builds a ChatModelBase from stored credential + config.""" +from fastapi import HTTPException, status + +from ..storage import StorageBase, ChatModelConfig +from ...credential import CredentialFactory +from ...model import ChatModelBase + + +async def get_model( + user_id: str, + config: ChatModelConfig, + storage: StorageBase, +) -> ChatModelBase: + """Get the model instance from the configuration and storage. + + Args: + user_id (`str`): + The user id. + config (`ChatModelConfig`): + The chat model configuration. + storage (`StorageBase`): + The storage instance. + + Returns: + `ChatModelBase`: + The model instance. + """ + credential_record = await storage.get_credential( + user_id, + config.credential_id, + ) + if credential_record is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Credential {config.credential_id!r} not found.", + ) + + credential = CredentialFactory.from_dict(credential_record.data) + model_cls = credential.get_chat_model_class() + parameters = ( + model_cls.Parameters(**config.parameters) + if config.parameters + else None + ) + return model_cls( + credential=credential, + model=config.model, + parameters=parameters, + ) diff --git a/src/agentscope/app/_service/_projectors/__init__.py b/src/agentscope/app/_service/_projectors/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..344784917e020ddf5d002fa69c7b4209aca081dd --- /dev/null +++ b/src/agentscope/app/_service/_projectors/__init__.py @@ -0,0 +1,13 @@ +# -*- coding: utf-8 -*- +"""Built-in event projectors. + +Each projector mirrors one cross-session UI feed onto the owning +session via the shared +:class:`~agentscope.app._service._session_projection.SessionProjection` +primitive. See :class:`~agentscope.app._types.EventProjector`. +""" +from ._subagent_hitl import SubagentHitlProjector + +__all__ = [ + "SubagentHitlProjector", +] diff --git a/src/agentscope/app/_service/_projectors/_subagent_hitl.py b/src/agentscope/app/_service/_projectors/_subagent_hitl.py new file mode 100644 index 0000000000000000000000000000000000000000..df86d880d41e6b26fc4fc00cf67d9130b6221d31 --- /dev/null +++ b/src/agentscope/app/_service/_projectors/_subagent_hitl.py @@ -0,0 +1,270 @@ +# -*- coding: utf-8 -*- +"""Projector that bridges team-member HITL events to leader sessions. + +When a team *member* (worker) session hits a tool call that needs human +confirmation, the worker run parks on an ``ASKING`` tool call in its +**own** session — invisible to a client subscribed only to the *leader* +session's event stream. This projector mirrors each such pending +request onto the leader session so the leader UI can render and resolve +it, and clears it when the request is answered. + +It is a thin strategy over the generic +:class:`~agentscope.app._service._session_projection.SessionProjection` +primitive: this file holds only the HITL-specific policy (which events +matter, how to resolve the leader, what the card payload looks like). +The durable hash, live notification, and key conventions all live in +the shared primitive. + +Persistence model (see the design doc, §2.4): + +- **Authoritative state** is the worker session's own ``state.context`` + (the ``ASKING`` tool call). The projection is only a mirror. +- The projected hash entry carries **no TTL**: a legitimate + confirmation can stay pending indefinitely. Stale entries (worker + cancelled/crashed without clearing) are healed by reconcile-on-read + at SSE replay time. +""" +from datetime import datetime +from typing import TYPE_CHECKING + +from ....event import ( + RequireUserConfirmEvent, + RequireExternalExecutionEvent, + UserConfirmResultEvent, + ExternalExecutionResultEvent, + ReplyEndEvent, +) + +if TYPE_CHECKING: + from ...storage import AgentRecord, SessionRecord, StorageBase + from ....event import AgentEvent + from .._session_projection import SessionProjection + + +class SubagentHitlProjector: + """Project pending team-member HITL requests onto leader sessions. + + Holds the storage handle it needs to resolve a worker's team (and + thus its leader). The :class:`SessionProjection` it writes through + is passed in per call by + :class:`~agentscope.app._service.ChatService`, which also resolves + confirm-routing and SSE replay through this projector's + :meth:`resolve` / :meth:`entry_id` helpers. + """ + + KIND = "subagent_hitl" + """Projection feed key (namespaces the entry within a session's + shared projection hash).""" + + EVT_REQUIRE = "subagent_require_user_confirm" + """``CustomEvent.name`` used to push/replay a pending request to the + leader's event stream.""" + + EVT_RESULT = "subagent_user_confirm_result" + """``CustomEvent.name`` used to tell the leader UI a pending request + has been resolved and its card should be cleared.""" + + def __init__(self, storage: "StorageBase") -> None: + """Bind the storage backend. + + Args: + storage (`StorageBase`): + Application storage, used to resolve the team (and hence + the leader session) a worker session belongs to. + """ + self._storage = storage + + @staticmethod + def entry_id(worker_session_id: str, reply_id: str) -> str: + """Return the projection entry id for one pending request. + + Args: + worker_session_id (`str`): + The worker session that emitted the HITL request. + reply_id (`str`): + The worker-side reply id the request belongs to. + + Returns: + `str`: + The entry id, ``"{worker_session_id}:{reply_id}"``. + """ + return f"{worker_session_id}:{reply_id}" + + async def maybe_project( + self, + user_id: str, + session_record: "SessionRecord", + agent_record: "AgentRecord", + event: "AgentEvent", + projection: "SessionProjection", + ) -> None: + """Mirror a worker HITL event onto its team's leader session. + + When a *worker* session emits an HITL request, write the pending + card (durable entry + live notification) onto the leader; when + it resolves one or its reply ends, clear the card. No-op for + non-team sessions and for the leader session itself (a leader's + own HITL reaches its client directly). + + Args: + user_id (`str`): + The owner of the running session. + session_record (`SessionRecord`): + The currently-running session's record. + agent_record (`AgentRecord`): + The currently-running agent's record. Only + ``source == "team"`` agents forward. + event (`AgentEvent`): + The event just published to this session's channel. + projection (`SessionProjection`): + Shared primitive used to write the durable entry and the + live notification. + """ + # Fast path: only team-member sessions forward anything, and + # only for the event kinds we care about. + if agent_record.source != "team" or not session_record.team_id: + return + if not isinstance( + event, + ( + RequireUserConfirmEvent, + RequireExternalExecutionEvent, + UserConfirmResultEvent, + ExternalExecutionResultEvent, + ReplyEndEvent, + ), + ): + return + + team = await self._storage.get_team( + user_id, + session_record.team_id, + ) + if team is None or team.session_id == session_record.id: + # No team, or this IS the leader session — nothing to mirror. + return + leader_sid = team.session_id + + if isinstance( + event, + (RequireUserConfirmEvent, RequireExternalExecutionEvent), + ): + payload = { + "worker_session_id": session_record.id, + "worker_agent_id": agent_record.id, + "worker_agent_name": agent_record.data.name, + "reply_id": event.reply_id, + "event_type": ( + "require_user_confirm" + if isinstance(event, RequireUserConfirmEvent) + else "require_external_execution" + ), + "event": event.model_dump(mode="json"), + "created_at": datetime.now().isoformat(), + } + await projection.upsert( + leader_sid, + self.KIND, + self.entry_id(session_record.id, event.reply_id), + payload, + ) + await projection.publish(leader_sid, self.EVT_REQUIRE, payload) + else: + # Clear the pending card. ``ReplyEndEvent`` is the primary + # clear signal (the resume's continuation event is NOT + # republished through the stream); the explicit result + # events clear early when they do flow through. All are + # idempotent — deleting an already-gone entry is a no-op. + await projection.delete( + leader_sid, + self.KIND, + self.entry_id(session_record.id, event.reply_id), + ) + await projection.publish( + leader_sid, + self.EVT_RESULT, + { + "worker_session_id": session_record.id, + "reply_id": event.reply_id, + }, + ) + + @classmethod + async def resolve( + cls, + projection: "SessionProjection", + leader_sid: str, + reply_id: str, + ) -> dict | None: + """Find the pending entry for ``reply_id`` under a leader. + + Used by the confirm-routing entry point (the chat router): given + a confirm result POSTed to the leader session, locate which + worker session it actually belongs to so the result can be + forwarded there. + + Args: + projection (`SessionProjection`): + The shared projection store to scan. + leader_sid (`str`): + The leader session the confirm result was POSTed to. + reply_id (`str`): + The worker-side reply id carried by the confirm result. + + Returns: + `dict | None`: + The stored payload (with ``worker_session_id`` / + ``worker_agent_id``), or ``None`` when no pending entry + matches — meaning the confirm is the leader's own. + """ + for entry in await projection.list(leader_sid, cls.KIND): + if entry.get("reply_id") == reply_id: + return entry + return None + + @classmethod + async def purge( + cls, + projection: "SessionProjection", + leader_sid: str, + ) -> None: + """Drop every pending HITL entry for a leader session. + + Used when the leader session (or its team) is deleted. Scoped to + this feed so other projections on the same session survive. + + Args: + projection (`SessionProjection`): + The shared projection store. + leader_sid (`str`): + The leader session to purge. + """ + await projection.purge(leader_sid, cls.KIND) + + @classmethod + async def drop_worker( + cls, + projection: "SessionProjection", + leader_sid: str, + worker_sid: str, + ) -> None: + """Drop every pending entry that originated from one worker. + + Used when a single worker session is deleted while the leader + survives. + + Args: + projection (`SessionProjection`): + The shared projection store. + leader_sid (`str`): + The leader session the entries are projected onto. + worker_sid (`str`): + The worker session whose entries should be dropped. + """ + for entry in await projection.list(leader_sid, cls.KIND): + if entry.get("worker_session_id") == worker_sid: + await projection.delete( + leader_sid, + cls.KIND, + cls.entry_id(worker_sid, entry["reply_id"]), + ) diff --git a/src/agentscope/app/_service/_session.py b/src/agentscope/app/_service/_session.py new file mode 100644 index 0000000000000000000000000000000000000000..570a402f13e256d28369e9e960bf1c2516d44e9e --- /dev/null +++ b/src/agentscope/app/_service/_session.py @@ -0,0 +1,473 @@ +# -*- coding: utf-8 -*- +"""Cross-resource session lifecycle service. + +Owns the "stop in-flight runs + delete records + drop bus state" +cascades that ``DELETE /sessions/{sid}``, ``DELETE /agents/{aid}``, +``DELETE /schedules/{sid}`` and the agent-facing +:class:`~agentscope.app._tools.TeamDelete` / +:class:`~agentscope.app._manager._scheduler._tools.ScheduleDelete` +tools all share. + +Layering +======== + +Methods deliberately delegate down the cascade so the bus-touching +logic lives in exactly one place — :meth:`delete_session`. Higher-level +methods only orchestrate which sessions to delete, then ask storage +to clean its own non-session scope (records, indexes, back-refs). + +:: + + delete_session ← atomic: cancel run, storage.delete_session, + bus.session_purge + │ + delete_team → service.delete_agent per worker + → storage.delete_team (record + leader detach) + │ + delete_agent → service.delete_session per session + → service.delete_schedule per owned schedule + → storage.delete_agent (agent record + team back-refs) + │ + delete_schedule → service.delete_session per spawned session + → storage.delete_schedule (schedule record + indexes) + +Storage's own internal cascades (e.g. +``storage.delete_agent`` re-iterating sessions) become idempotent +no-ops because the records are already gone — they still execute, but +do no work and never touch the bus, so the storage layer stays +unaware of the message bus. + +Separation of concerns +====================== + +Storage and message bus are treated as distinct backends — they may +live in different databases in the future. The service is the **only** +component that touches both in the same call. Storage code never +imports the bus; bus code never imports storage. +""" +import asyncio + +from ..message_bus import MessageBus, MessageBusKeys +from ..storage import StorageBase +from ._session_projection import SessionProjection +from ._projectors import SubagentHitlProjector +from ..._logging import logger + + +class SessionService: + """Cancel in-flight chat runs and cascade-delete related records. + + The cancel side broadcasts via + :meth:`MessageBus.session_publish_cancel`, then polls + :meth:`MessageBus.session_is_running` until the run-lock clears or + a timeout expires — so the implementation is multi-process and + multi-node by construction. + + Args: + storage (`StorageBase`): + Persistent storage backend. Owns durable records and their + cascades among themselves. + message_bus (`MessageBus`): + Live message bus. Owns transient per-session state (events + log, inbox, run-lock, cancel channel). + """ + + _CANCEL_POLL_INTERVAL_SECS: float = 0.1 + """Interval between :meth:`MessageBus.session_is_running` polls + while waiting for a cancelled run to release its distributed + run-lock.""" + + def __init__( + self, + storage: StorageBase, + message_bus: MessageBus, + ) -> None: + """Bind dependencies. + + Args: + storage (`StorageBase`): Persistent storage backend. + message_bus (`MessageBus`): Live message bus. + """ + self._storage = storage + self._bus = message_bus + self._projection = SessionProjection(message_bus) + + # ------------------------------------------------------------------ + # Cancel + # ------------------------------------------------------------------ + + async def cancel_session_run( + self, + session_id: str, + *, + timeout: float = 10.0, + ) -> bool: + """Broadcast a session cancel and wait for the chat-run lock to + clear. + + Publishes one cancel payload on the bus's shared cancel channel, + unconditionally. Every process's + :class:`~agentscope.app._manager.CancelDispatcher` reacts to the + broadcast by cancelling whatever it locally holds for the + session — the chat-run asyncio task **and** any background + tasks owned by that session. The publisher does not need to + know which worker holds which piece. + + After publishing, polls + :meth:`MessageBus.session_is_running` until the distributed + chat-run lock clears. Only the chat run holds a distributed + lock; BG tasks do not, so this poll only waits for the chat + run. Returns immediately when no chat run was active. + + Idempotent: calling on an idle session just sends a no-op + broadcast and observes a clear lock. + + Args: + session_id (`str`): + The session whose chat run + BG tasks should be + cancelled. + timeout (`float`, defaults to ``10.0``): + Maximum seconds to wait for the chat-run lock to + release. On timeout the method returns ``False`` so + callers can proceed (e.g. with cascade delete) instead + of hanging on a process that may have died. + + Returns: + `bool`: + ``True`` if the chat-run lock was confirmed released + within ``timeout`` seconds (or was never held). + ``False`` if the lock was still held when the timeout + expired. + """ + await self._bus.publish( + MessageBusKeys.session_cancel_channel(), + {"session_id": session_id}, + ) + + deadline = asyncio.get_event_loop().time() + timeout + while True: + if not await self._bus.is_locked( + MessageBusKeys.session_lock(session_id), + ): + return True + if asyncio.get_event_loop().time() >= deadline: + logger.warning( + "Session %s did not release its run-lock within " + "%.1fs after cancel; proceeding anyway.", + session_id, + timeout, + ) + return False + await asyncio.sleep(self._CANCEL_POLL_INTERVAL_SECS) + + # ------------------------------------------------------------------ + # Delete cascades — every higher-level method delegates to + # ``delete_session`` so the cancel + bus-purge logic exists in + # exactly one place. + # ------------------------------------------------------------------ + + async def delete_session( + self, + user_id: str, + agent_id: str, + session_id: str, + ) -> bool: + """Cancel, delete and bus-purge a single session. + + This is the atomic primitive — every other cascade delegates + here for per-session work. + + Steps: + + 1. Cancel any in-flight run for ``session_id`` (cross-process + via the bus cancel channel). + 2. Delete the session record (and its storage-side cascade: + message log, schedule-session index, team dissolution when + this session leads one — recursive into worker agents). + 3. Purge transient bus state for ``session_id`` (events log, + inbox). + + Worker sessions that storage cascades through are picked up + here too: when this session is a team leader, + ``storage.delete_session`` calls ``storage.delete_team`` → + ``storage.delete_agent`` → ``storage.delete_session`` for each + worker, and we mirror that on the bus side by purging worker + sessions identified up front via + :meth:`_team_worker_session_ids`. + + Args: + user_id (`str`): The owner user id. + agent_id (`str`): The agent that owns the session. + session_id (`str`): The session to delete. + + Returns: + `bool`: + ``True`` if the session record existed and was deleted, + ``False`` otherwise. Mirrors + :meth:`StorageBase.delete_session`. + """ + # Identify all bus-purge targets before storage mutates anything. + worker_sids = await self._team_worker_session_ids( + user_id, + agent_id, + session_id, + ) + all_sids = [session_id, *worker_sids] + + # Clean leader-side subagent HITL projections before storage + # cascades remove the records we need to resolve roles from. + await self._purge_subagent_hitl(user_id, agent_id, session_id) + + await self._cancel_runs(all_sids) + deleted = await self._storage.delete_session( + user_id, + agent_id, + session_id, + ) + await self._purge_bus(all_sids) + return deleted + + async def delete_team(self, user_id: str, team_id: str) -> bool: + """Cancel, delete and bus-purge a team. + + Delegates worker dissolution to :meth:`delete_agent` (one call + per ``member_id``) so the per-session cancel + bus purge runs + for each worker. The leader's own session is **not** deleted — + teams dissolve, leaders survive (and have their ``team_id`` + cleared by ``storage.delete_team``). + + Args: + user_id (`str`): The owner user id. + team_id (`str`): The team to dissolve. + + Returns: + `bool`: + ``True`` if the team record existed and was deleted. + """ + team = await self._storage.get_team(user_id, team_id) + if team is None: + # Still call storage.delete_team so it can clean any index + # residue, but the return value will be False. + return await self._storage.delete_team(user_id, team_id) + + for member_id in team.data.member_ids: + await self.delete_agent(user_id, member_id) + + # storage.delete_team will iterate member_ids again to delete + # each worker agent — those calls are now no-ops because the + # agents are already gone, leaving only the leader-detach and + # team-record cleanup work. + return await self._storage.delete_team(user_id, team_id) + + async def delete_agent(self, user_id: str, agent_id: str) -> bool: + """Cancel, delete and bus-purge every session and schedule + owned by an agent, then drop the agent record. + + Delegates per-session work to :meth:`delete_session` and + per-schedule work to :meth:`delete_schedule`, then asks + storage to clean the remaining agent-scoped state (the agent + record, the agent index entry, and any team back-references). + + Args: + user_id (`str`): The owner user id. + agent_id (`str`): The agent to delete. + + Returns: + `bool`: + ``True`` if the agent record existed and was deleted. + """ + for session in await self._storage.list_sessions(user_id, agent_id): + await self.delete_session(user_id, agent_id, session.id) + + for schedule in await self._storage.list_schedules(user_id): + if schedule.agent_id == agent_id: + await self.delete_schedule(user_id, schedule.id) + + # storage.delete_agent re-iterates sessions and schedules — + # those re-runs are idempotent no-ops because the records were + # already removed above. What remains is the agent record, + # the agent index entry, and team back-reference scrubbing. + return await self._storage.delete_agent(user_id, agent_id) + + async def delete_schedule( + self, + user_id: str, + schedule_id: str, + ) -> bool: + """Cancel, delete and bus-purge every session spawned by a + schedule, then drop the schedule record. + + Args: + user_id (`str`): The owner user id. + schedule_id (`str`): The schedule to delete. + + Returns: + `bool`: + ``True`` if the schedule record existed and was deleted. + """ + for session in await self._storage.list_sessions_by_schedule( + user_id, + schedule_id, + ): + await self.delete_session( + user_id, + session.agent_id, + session.id, + ) + + # storage.delete_schedule re-iterates the same sessions — + # idempotent no-ops; only schedule record + indexes remain. + return await self._storage.delete_schedule(user_id, schedule_id) + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + async def _team_worker_session_ids( + self, + user_id: str, + agent_id: str, + session_id: str, + ) -> list[str]: + """Return the session ids of every worker in the team that + ``session_id`` leads, or ``[]`` when the session does not lead + a team. + + Mirrors :meth:`StorageBase.delete_session`'s own team-leader + cascade so the bus side can purge the same sessions. + + Args: + user_id (`str`): The owner user id. + agent_id (`str`): + The agent that owns ``session_id``. May be empty when + unknown; team-leader lookup does not depend on it. + session_id (`str`): + The candidate leader session. + + Returns: + `list[str]`: + Worker session ids, empty when this session is not a + team leader. + """ + session = await self._storage.get_session( + user_id, + agent_id, + session_id, + ) + if session is None or not session.team_id: + return [] + team = await self._storage.get_team(user_id, session.team_id) + if team is None or team.session_id != session_id: + return [] + sids: list[str] = [] + for member_id in team.data.member_ids: + worker_sessions = await self._storage.list_sessions( + user_id, + member_id, + ) + sids.extend(s.id for s in worker_sessions) + return sids + + async def _cancel_runs(self, session_ids: list[str]) -> None: + """Cancel every in-flight run in ``session_ids`` concurrently. + + Args: + session_ids (`list[str]`): + Sessions whose runs should be cancelled. + """ + if not session_ids: + return + await asyncio.gather( + *(self.cancel_session_run(sid) for sid in session_ids), + ) + + async def _purge_bus(self, session_ids: list[str]) -> None: + """Drop bus state (events log + inbox) for each id concurrently. + + Args: + session_ids (`list[str]`): + Sessions whose bus state should be purged. + """ + if not session_ids: + return + await asyncio.gather( + *(self._purge_session_bus(sid) for sid in session_ids), + ) + + async def _purge_session_bus(self, session_id: str) -> None: + """Drop all per-session bus state for one session.""" + await self._bus.log_trim( + MessageBusKeys.session_events(session_id), + ) + await self._bus.queue_delete( + MessageBusKeys.inbox(session_id), + ) + await self._bus.registry_drop( + MessageBusKeys.bg_tasks(session_id), + ) + + async def _purge_subagent_hitl( + self, + user_id: str, + agent_id: str, + session_id: str, + ) -> None: + """Clean leader-side subagent HITL projections for a session + about to be deleted (design §3.7). + + Two cases, resolved from the session's role: + + - **Leader session** (it leads a team): purge the entire hash + keyed by this session — every projected member card goes. + - **Worker session** (it has a ``team_id`` but is not the + leader): drop just this worker's entries from the *leader's* + hash, leaving sibling members' cards intact. + + Must run before storage cascades remove the team / session + records this resolution depends on. Failures are swallowed — a + stale projection is self-healed by reconcile-on-read and must + not block the delete cascade. + + Args: + user_id (`str`): + The owner user id. + agent_id (`str`): + The agent that owns ``session_id``. + session_id (`str`): + The session being deleted. + """ + try: + session = await self._storage.get_session( + user_id, + agent_id, + session_id, + ) + if session is None or not session.team_id: + # Not in a team — also clear any hash that may have been + # created with this session as a (future) leader key. + await SubagentHitlProjector.purge(self._projection, session_id) + return + + team = await self._storage.get_team(user_id, session.team_id) + if team is None: + await SubagentHitlProjector.purge(self._projection, session_id) + return + + if team.session_id == session_id: + # Leader session — drop the whole projection store. + await SubagentHitlProjector.purge(self._projection, session_id) + else: + # Worker session — drop only its entries from the + # leader's store. + await SubagentHitlProjector.drop_worker( + self._projection, + team.session_id, + session_id, + ) + except Exception as e: # pylint: disable=broad-except + logger.warning( + "Failed to purge subagent HITL projection for session " + "%s: %s", + session_id, + str(e), + ) diff --git a/src/agentscope/app/_service/_session_projection.py b/src/agentscope/app/_service/_session_projection.py new file mode 100644 index 0000000000000000000000000000000000000000..3da50d2c088ff0a5b227010d653c7d1fba59a54f --- /dev/null +++ b/src/agentscope/app/_service/_session_projection.py @@ -0,0 +1,186 @@ +# -*- coding: utf-8 -*- +"""Generic cross-session UI projection primitive. + +A *projection* mirrors a UI card owned by one session onto another +session's event stream, so a client subscribed only to the target +session can render and resolve it. The canonical use is team HITL: a +worker (member) session parks on a tool call awaiting confirmation in +its own session — invisible to a client watching only the *leader* — +so the pending request is projected onto the leader. + +This class is **pure mechanism**: it knows nothing about teams, +workers, leaders, or HITL. It is the reusable substrate every such +feature shares — a durable per-session hash plus a live notification — +so a new projection feature is a small strategy object (an +``EventProjector``) over this primitive, not a new bus wrapper class. + +Backed entirely by the message-bus generic registry primitives +(``registry_*``) and :meth:`MessageBus.session_publish_event`; no +business methods are added to ``MessageBus``. Key conventions live in +:class:`~agentscope.app.message_bus.MessageBusKeys`. + +Persistence model: + +- The Redis hash is the **only durable** record of a projected card — + the target session's event channel (replay log + pub/sub) is not + durable across runs. It carries **no TTL**: a legitimate card can + stay pending indefinitely. Authoritative truth for whether a card is + still live lives in the *owning* session's own state; stale hash + entries are healed by reconcile-on-read at SSE replay time, not by + expiry. +- ``kind`` partitions the hash so one target session can host several + independent feeds (HITL, progress, errors, …) without collision. +""" +import json +from typing import TYPE_CHECKING + +from ..message_bus import MessageBusKeys +from .._bus_ops import publish_session_event +from ...event import CustomEvent + +if TYPE_CHECKING: + from ..message_bus import MessageBus + + +class SessionProjection: + """Durable per-session store of UI cards projected from elsewhere. + + A thin stateless wrapper around the message bus — construct one + wherever needed (it holds only a bus reference). Entries are grouped + per ``(target_session_id, kind)``; within a feed each entry is keyed + by a caller-chosen ``entry_id``. + + Live notification piggybacks on the target session's existing event + channel via :meth:`publish`, so front-ends receive updates over the + same ``GET /sessions/{sid}/stream`` SSE connection they already use. + """ + + def __init__(self, message_bus: "MessageBus") -> None: + """Bind the message bus. + + Args: + message_bus (`MessageBus`): + Application message bus; only its generic ``registry_*`` + primitives and :meth:`session_publish_event` are used. + """ + self._bus = message_bus + + async def upsert( + self, + target_sid: str, + kind: str, + entry_id: str, + payload: dict, + ) -> None: + """Persist (or overwrite) one projected entry. + + Args: + target_sid (`str`): + The session the entry is projected onto. + kind (`str`): + The projection feed (e.g. ``"subagent_hitl"``). + entry_id (`str`): + Identity of the entry within the feed. + payload (`dict`): + The entry to store (JSON-serializable). + """ + await self._bus.registry_set( + MessageBusKeys.projection_namespace(target_sid), + MessageBusKeys.projection_field(kind, entry_id), + json.dumps(payload), + ) + + async def delete( + self, + target_sid: str, + kind: str, + entry_id: str, + ) -> None: + """Remove one projected entry. + + Idempotent: a no-op when the entry is already gone. + + Args: + target_sid (`str`): + The session the entry was projected onto. + kind (`str`): + The projection feed. + entry_id (`str`): + Identity of the entry within the feed. + """ + await self._bus.registry_del( + MessageBusKeys.projection_namespace(target_sid), + MessageBusKeys.projection_field(kind, entry_id), + ) + + async def list(self, target_sid: str, kind: str) -> list[dict]: + """Return every entry in one feed for a target session. + + Args: + target_sid (`str`): + The session whose projections to read. + kind (`str`): + The projection feed to filter by. + + Returns: + `list[dict]`: + All stored payloads in the feed; empty when none. + """ + raw = await self._bus.registry_getall( + MessageBusKeys.projection_namespace(target_sid), + ) + prefix = MessageBusKeys.projection_field_prefix(kind) + return [ + json.loads(value) + for field, value in raw.items() + if field.startswith(prefix) + ] + + async def purge(self, target_sid: str, kind: str | None = None) -> None: + """Drop projected entries for a target session. + + Args: + target_sid (`str`): + The session to purge. + kind (`str | None`, optional): + When given, drop only that feed's entries (preserving + other feeds on the same session). When ``None``, drop + the session's entire projection store in one shot. + """ + if kind is None: + await self._bus.registry_drop( + MessageBusKeys.projection_namespace(target_sid), + ) + return + ns = MessageBusKeys.projection_namespace(target_sid) + prefix = MessageBusKeys.projection_field_prefix(kind) + raw = await self._bus.registry_getall(ns) + for field in raw: + if field.startswith(prefix): + await self._bus.registry_del(ns, field) + + async def publish( + self, + target_sid: str, + event_name: str, + value: dict, + ) -> None: + """Send a live ``CustomEvent`` to a target session's channel. + + Notifies front-ends subscribed to the target session that a + projected card should be rendered or cleared. + + Args: + target_sid (`str`): + The session to notify. + event_name (`str`): + The ``CustomEvent.name`` carried to the front-end. + value (`dict`): + The event payload. + """ + custom = CustomEvent(name=event_name, value=value) + await publish_session_event( + self._bus, + target_sid, + custom.model_dump(mode="json"), + ) diff --git a/src/agentscope/app/_service/_toolkit.py b/src/agentscope/app/_service/_toolkit.py new file mode 100644 index 0000000000000000000000000000000000000000..f3eff63653d4ab34c261ed996f069765c170f00d --- /dev/null +++ b/src/agentscope/app/_service/_toolkit.py @@ -0,0 +1,192 @@ +# -*- coding: utf-8 -*- +"""Toolkit assembly for an (agent, session) pair. + +The single entry point :func:`get_toolkit` gathers every tool source — +workspace builtins, MCPs, skills, planning tools (Task*), background-task +control (ToolStop), schedule control (Schedule*), team participation +tools, and caller-supplied extras — into one :class:`Toolkit`. +""" +from typing import Any + +from .._manager import BackgroundTaskManager, SchedulerManager +from ..message_bus import MessageBus +from .._tool import AgentCreate, TeamCreate, TeamDelete, TeamSay +from .._types import AgentToolFactory, SubAgentTemplate +from ..storage import AgentRecord, SessionRecord, StorageBase +from ...middleware import MiddlewareBase +from ...tool import ( + TaskCreate, + TaskGet, + TaskList, + TaskUpdate, + Toolkit, + ToolGroup, +) +from ...workspace import WorkspaceBase + + +async def get_toolkit( + *, + storage: StorageBase, + workspace: WorkspaceBase, + scheduler_manager: SchedulerManager, + background_task_manager: BackgroundTaskManager, + message_bus: MessageBus, + middlewares: list[MiddlewareBase], + user_id: str, + agent_record: AgentRecord, + session_record: SessionRecord, + extra_factory: AgentToolFactory | None = None, + sub_agent_templates: dict[str, SubAgentTemplate] | None = None, +) -> Toolkit: + """Assemble the complete :class:`Toolkit` for one chat turn. + + Tool sources (in attachment order): + + 1. Workspace builtins (Bash / Read / Write / Grep / …) + 2. Planning tools (:class:`TaskCreate` / :class:`TaskList` / + :class:`TaskGet` / :class:`TaskUpdate`) + 3. Background-task control (:class:`ToolStop`, from + :meth:`BackgroundTaskManager.list_tools`) + 4. Schedule control (:class:`ScheduleCreate` / :class:`ScheduleView` + / :class:`ScheduleDelete` / :class:`ScheduleList`, from + :meth:`SchedulerManager.list_tools`). Only attached when the + session has a model configured (Schedule tools need a model to + fire new chats with). + 5. Team tools — selected inline by ``agent_record.source``: + worker (``"team"``) gets only ``TeamSay``; everyone else gets + the full leader-side toolset + (``TeamCreate / AgentCreate / TeamSay / TeamDelete``) + 6. Caller-supplied extras (``extra_factory``) + + Plus the workspace's skills and MCPs, which become the toolkit's + ``skills_or_loaders`` and ``mcps`` parameters. + + Args: + storage (`StorageBase`): + Application storage backend; needed by team tools to read + fresh team / session state at call time, and by schedule + tools. + workspace (`WorkspaceBase`): + Pre-resolved per-session workspace (caller resolves it + via :meth:`WorkspaceManagerBase.get_workspace`). Used here + for tool / skill / MCP discovery. + scheduler_manager (`SchedulerManager`): + Application scheduler. Provides the four schedule tools and + persists schedules through it. + background_task_manager (`BackgroundTaskManager`): + Application background-task registry. Provides the + :class:`ToolStop` tool bound to its live task dict. + message_bus (`MessageBus`): + Application message bus; passed to team tools so they can + push HintBlocks + wakeups when delivering inter-session + messages. + middlewares (`list[MiddlewareBase]`): + The agent middlewares that may provide tools to the agent via the + `list_tools` interface. + user_id (`str`): + Caller user id. + agent_record (`AgentRecord`): + Pre-loaded agent record (loaded once by the caller). Its + ``source`` field determines which team tools are attached. + session_record (`SessionRecord`): + Pre-loaded session record (loaded once by the caller). + Used for the schedule-tool model configuration. + extra_factory (`AgentToolFactory | None`, optional): + Async factory invoked once per assembly to produce + user/session-specific extra tools. + sub_agent_templates (`dict[str, SubAgentTemplate] | None`, \ +optional): + Sub-agent template registry, keyed by template type. + Passed to the ``AgentCreate`` tool so it can route to + the appropriate template when a ``subagent_type`` is + specified by the leader agent. + + Returns: + `Toolkit`: Fully populated toolkit (tools + skills + MCPs). + """ + + tool_groups = [] + + # The general tools running in the workspace + tools = await workspace.list_tools() + + # Planning tools — always on. + tools += [TaskCreate(), TaskList(), TaskGet(), TaskUpdate()] + + # Background-task control. + tools += await background_task_manager.list_tools( + session_id=session_record.id, + ) + + # Schedule control. Requires a model config on this session because + # ``ScheduleCreate`` records it into new ``ScheduleRecord`` instances. + if session_record.config.chat_model_config is not None: + # Add schedule tools as a tool group + tool_groups.append( + ToolGroup( + name="schedule_tools", + description=( + """Tools for managing cron schedules. A cron schedule is \ +a recurring task that fires at a specified time — at that point, a new \ +session is created and an agent will be invoked to complete the given task \ +autonomously. + +## When to Use This Tool Group +- When you need to create a new cron schedule that triggers at a specific \ +time or interval" +- When you're asked to list, inspect, stop, or delete existing cron schedules +""" + ), + tools=await scheduler_manager.list_tools( + user_id=user_id, + agent_id=agent_record.id, + chat_model_config=session_record.config.chat_model_config, + ), + ), + ) + + # Team tools — variant based on ``agent_record.source``. A worker + # only gets TeamSay (to report back); a user-owned agent always + # gets the full leader-side toolset. Each tool checks its own + # preconditions (am I in a team? am I the leader?) at call time + # against fresh storage, which is why the full set can be attached + # unconditionally without needing a stale snapshot of team_id. + team_tool_kwargs: dict[str, Any] = { + "storage": storage, + "message_bus": message_bus, + "user_id": user_id, + "session_id": session_record.id, + "agent_id": agent_record.id, + } + if agent_record.source == "team": + tools.append(TeamSay(**team_tool_kwargs, role="worker")) + else: + tools += [ + TeamCreate(**team_tool_kwargs), + AgentCreate( + **team_tool_kwargs, + sub_agent_templates=sub_agent_templates or {}, + ), + TeamSay(**team_tool_kwargs, role="leader"), + TeamDelete(**team_tool_kwargs), + ] + + # Caller-supplied extras. + if extra_factory is not None: + tools += await extra_factory( + user_id, + agent_record.id, + session_record.id, + ) + + # Tools from middleware + for mw in middlewares: + tools.extend(await mw.list_tools()) + + return Toolkit( + tools=tools, + skills_or_loaders=await workspace.list_skills(), + mcps=await workspace.list_mcps(), + tool_groups=tool_groups, + ) diff --git a/src/agentscope/app/_service/_tts_model.py b/src/agentscope/app/_service/_tts_model.py new file mode 100644 index 0000000000000000000000000000000000000000..e60a845c2d46a2073c1b0dbc1fa747a28e618a2e --- /dev/null +++ b/src/agentscope/app/_service/_tts_model.py @@ -0,0 +1,68 @@ +# -*- coding: utf-8 -*- +"""TTS model service: builds a TTSModelBase from stored credential + config.""" +from typing import Type + +from fastapi import HTTPException, status + +from ..storage import StorageBase, TTSModelConfig +from ...credential import CredentialFactory +from ...tts import TTSModelBase + + +async def get_tts_model( + user_id: str, + config: TTSModelConfig, + storage: StorageBase, +) -> TTSModelBase: + """Get the TTS model instance from the configuration and storage. + + Args: + user_id (`str`): + The user id. + config (`TTSModelConfig`): + The TTS model configuration. + storage (`StorageBase`): + The storage instance. + + Returns: + `TTSModelBase`: + The TTS model instance. + """ + credential_record = await storage.get_credential( + user_id, + config.credential_id, + ) + if credential_record is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Credential {config.credential_id!r} not found.", + ) + + credential = CredentialFactory.from_dict(credential_record.data) + tts_classes = credential.get_tts_model_classes() + if not tts_classes: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Provider {config.type!r} does not support TTS models.", + ) + + tts_cls = _resolve_tts_class(tts_classes, config.model) + parameters = ( + tts_cls.Parameters(**config.parameters) if config.parameters else None + ) + return tts_cls( + credential=credential, + model=config.model, + parameters=parameters, + ) + + +def _resolve_tts_class( + classes: list[Type[TTSModelBase]], + model: str, +) -> Type[TTSModelBase]: + """Pick the TTS class that lists the given model name.""" + for cls in classes: + if any(card.name == model for card in cls.list_models()): + return cls + return classes[0] diff --git a/src/agentscope/app/_tool/__init__.py b/src/agentscope/app/_tool/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9ae334f11dca81ad57d22526625af626b8e79245 --- /dev/null +++ b/src/agentscope/app/_tool/__init__.py @@ -0,0 +1,39 @@ +# -*- coding: utf-8 -*- +"""Framework-builtin tools wired into team-participating agents. + +These tools differ from the workspace-provided builtins (Bash, Read, +Task series, …) in two ways: + +1. **Construction depends on app-level resources** — they bind a + :class:`StorageBase` + :class:`MessageBus` reference plus the + request-scoped ``user_id`` / ``session_id`` / ``agent_id`` at agent + assembly time, and call storage / bus directly in their + ``__call__`` — except ``TeamDelete``, which delegates to + :class:`SessionService` for cascade deletion. +2. **Visibility depends only on the agent's source field** — + ``source='user'`` agents always see the full leader-side toolset + (``TeamCreate / AgentCreate / TeamSay / TeamDelete``) regardless of + whether they currently lead a team. Each tool checks the storage + state at ``__call__`` time and fails clearly if its precondition + is not met. ``source='team'`` workers only see ``TeamSay``. + + The benefit is that a single chat run can ``TeamCreate`` and then + ``AgentCreate`` immediately afterward — there is no toolkit + refresh point because the toolkit never changed; only the + underlying storage state did, which the next tool reads fresh. + +Selection of the right subset by ``agent.source`` happens inline in +:func:`get_toolkit`; there is no separate "team tool factory" helper. +""" +from ._agent_create import AgentCreate, DEFAULT_SUB_AGENT_TEMPLATE +from ._team_create import TeamCreate +from ._team_delete import TeamDelete +from ._team_say import TeamSay + +__all__ = [ + "AgentCreate", + "DEFAULT_SUB_AGENT_TEMPLATE", + "TeamCreate", + "TeamDelete", + "TeamSay", +] diff --git a/src/agentscope/app/_tool/_agent_create.py b/src/agentscope/app/_tool/_agent_create.py new file mode 100644 index 0000000000000000000000000000000000000000..9d55fd15521a379043e59ae2da6c29346d5c8f8f --- /dev/null +++ b/src/agentscope/app/_tool/_agent_create.py @@ -0,0 +1,530 @@ +# -*- coding: utf-8 -*- +"""The AgentCreate tool — spawns a worker into the current team.""" +from __future__ import annotations + +import copy +import json +from typing import TYPE_CHECKING + +from pydantic import Field + +from ._team_tool_base import _TeamToolBase +from .._types import SubAgentTemplate +from ..message_bus import MessageBusKeys +from .._bus_ops import enqueue_run_trigger +from ..storage import AgentData, AgentRecord, SessionConfig +from ...message import HintBlock, TextBlock, ToolResultState +from ...permission import PermissionContext +from ...state import AgentState +from ...tool import ToolChunk, ParamsBase + +if TYPE_CHECKING: + from ..message_bus import MessageBus + from ..storage import StorageBase + + +_DEFAULT_SYSTEM_PROMPT_TEMPLATE = ( + "You are {member_name}, a member of team '{team_name}' led by " + "{leader_name}.\n\n" + "Team purpose: {team_description}\n\n" + "Your role: {member_description}\n\n" + "You communicate with the team leader and other members " + "through the TeamSay tool. " + "Speak on the team only when you have something " + "external to share — your private reasoning stays private." +) + +DEFAULT_SUB_AGENT_TEMPLATE = SubAgentTemplate( + type="default", + description="Default worker agent with standard configuration.", + system_prompt_template=_DEFAULT_SYSTEM_PROMPT_TEMPLATE, + override_leader_mode=False, + extend_leader_permission_rules=True, + extend_leader_working_directories=True, +) +# The built-in default sub-agent template. +# +# Used when no custom templates are registered, or when the leader +# agent creates a member without specifying a ``subagent_type`` (or +# explicitly specifies ``subagent_type="default"``). Developers can +# override this by registering their own template with +# ``type="default"`` via :func:`~agentscope.app.create_app`. +# +# The default template fully follows the leader: the worker inherits +# the leader's permission mode, working directories, and rules — +# matching the intuition that a generic worker should behave like the +# leader unless the developer registers a more opinionated template. + + +def _merge_leader_permissions( + template: SubAgentTemplate, + leader_context: PermissionContext, +) -> PermissionContext: + """Build the worker's permission context from the template, layered + with the leader's runtime state according to the template's three + inherit-from-leader flags. + + - ``override_leader_mode``: if True, the template's + :attr:`PermissionContext.mode` wins; otherwise the worker + inherits the leader's mode. + - ``extend_leader_permission_rules``: if True, the leader's + allow/deny/ask rules are appended after the template's rules for + each tool, so the worker doesn't re-prompt for permissions the + user has already granted in the leader session. The template's + rules appear first in each list, so the engine — which returns + on the first matching rule per stage — evaluates the template's + intent before the leader's. + - ``extend_leader_working_directories``: if True, the leader's + working directories are merged in; on key (path) collisions the + template's entry wins. + + The template fields are deep-copied so the returned context is + independent of both the template and the leader state. + """ + merged = template.permission_context.model_copy(deep=True) + + if not template.override_leader_mode: + merged.mode = leader_context.mode + + if template.extend_leader_working_directories: + for path, wd in leader_context.working_directories.items(): + merged.working_directories.setdefault( + path, + wd.model_copy(deep=True), + ) + + if template.extend_leader_permission_rules: + for attr in ("allow_rules", "deny_rules", "ask_rules"): + merged_rules: dict = getattr(merged, attr) + for tool_name, rules in getattr(leader_context, attr).items(): + merged_rules.setdefault(tool_name, []).extend( + r.model_copy(deep=True) for r in rules + ) + + return merged + + +class _AgentCreateParams(ParamsBase): + """Parameters for :class:`AgentCreate`.""" + + name: str = Field( + description=( + 'Short identifier for the new member, e.g. ``"researcher"`` ' + 'or ``"coder-1"``. Other members address it via ' + "``TeamSay(to=)``, so it MUST be unique within " + "the team." + ), + ) + description: str = Field( + description=( + "One-sentence summary of the member's role — e.g. " + '``"Researches background information on the target topic"``. ' + "Becomes part of the member's system prompt so it understands " + "its place in the team." + ), + ) + prompt: str = Field( + description=( + "The first task delivered to the member as a user message. " + "The member begins executing immediately upon creation, so " + "make this concrete and self-contained — do not just say " + '``"wait for instructions"`` (use TeamSay later instead). ' + "Include any context, constraints, deliverables, and " + "deadlines the member needs." + ), + ) + + +class AgentCreate(_TeamToolBase): + """Spawn a new worker member into the team you lead.""" + + name: str = "AgentCreate" + is_state_injected: bool = True + + description: str = """Add a new member to the team you lead. + +## When to Use This Tool +After ``TeamCreate``, call this for each member you want on the team. \ +Each call: +- Creates a worker agent dedicated to this team. +- Delivers ``prompt`` as the worker's first user message — **the worker \ +starts executing it immediately**. (So DONT use ``TeamSay`` right after \ +creating one agent). + +## When NOT to Use This Tool +- You're not currently leading a team. Call ``TeamCreate`` first. +- The new member would duplicate an existing member's role; reuse the \ +existing member via ``TeamSay`` instead. + +## Effects +- Use the ``name`` you chose as ``to=`` in ``TeamSay`` to direct \ +messages to this member specifically. Names must be unique within the \ +team (including against the leader's name); duplicates are rejected. +- Members spawned this way live only as long as the team — they are \ +deleted when ``TeamDelete`` is called. + +## Important +- You are responsible for organising the team, assigning tasks, collecting \ +every member's report, and producing the final answer — all members report \ +directly to you. Therefore, **DO NOT** encourage members to communicate with \ +each other, and **AVOID** creating "integrator"-style members; both make the \ +overall communication topology unnecessarily complex. +""" + + input_schema: dict = _AgentCreateParams.model_json_schema() + + def __init__( + self, + storage: "StorageBase", + message_bus: "MessageBus", + user_id: str, + session_id: str, + agent_id: str, + sub_agent_templates: dict[str, SubAgentTemplate] | None = None, + ) -> None: + """Bind request-scoped identifiers plus sub-agent templates. + + Extends :meth:`_TeamToolBase.__init__` with an optional + template registry. The built-in ``"default"`` template is + always present as a fallback; developers can override it by + registering their own template with ``type="default"``. + + When more than one template type is available (i.e. custom + templates were registered), the tool's ``input_schema`` is + dynamically extended with a ``subagent_type`` enum field so + the leader agent can choose which type to create. + + Args: + storage (`StorageBase`): + Application storage backend. + message_bus (`MessageBus`): + Application message bus for inter-session delivery. + user_id (`str`): + The owner user id of the calling agent. + session_id (`str`): + The current session id of the calling agent. + agent_id (`str`): + The id of the agent invoking the tool. + sub_agent_templates (`dict[str, SubAgentTemplate] | None`, \ +optional): + Template registry keyed by template type. The + built-in ``"default"`` template is injected + automatically if not already present. + """ + super().__init__(storage, message_bus, user_id, session_id, agent_id) + + self._sub_agent_templates: dict[str, SubAgentTemplate] = dict( + sub_agent_templates or {}, + ) + if "default" not in self._sub_agent_templates: + self._sub_agent_templates["default"] = DEFAULT_SUB_AGENT_TEMPLATE + + # Only expose subagent_type when the developer registered + # custom templates — a single "default" type is redundant in + # the schema and would confuse the LLM. + has_custom_templates = set(self._sub_agent_templates) != {"default"} + if has_custom_templates: + schema = copy.deepcopy( + _AgentCreateParams.model_json_schema(), + ) + type_descriptions = "\n".join( + f"- ``{t.type!r}`` — {t.description}" + for t in self._sub_agent_templates.values() + ) + schema["properties"]["subagent_type"] = { + "type": "string", + "enum": list(self._sub_agent_templates), + "description": ( + "The type of sub-agent template to use. " + "Available types:\n\n" + f"{type_descriptions}\n\n" + "Each type has pre-configured system prompt, " + "permissions, and task context." + ), + } + self.input_schema = schema + + async def __call__( + self, + name: str, + description: str, + prompt: str, + subagent_type: str = "default", + _agent_state: AgentState | None = None, + ) -> ToolChunk: + """Spawn the worker agent + session directly via storage. + + Reads the current session + team records from storage to + enforce two preconditions: the calling session must be in a + team, and it must be that team's leader. + + The worker's configuration (system prompt, context/react + config, permission context, task context) is determined by + the :class:`SubAgentTemplate` matching ``subagent_type``. The + leader's user-confirmed permission rules and working + directories are merged into the template's permission context + so the worker does not re-prompt for permissions the user has + already granted. + + Args: + name (`str`): + Short identifier for the worker, unique within the + team. Used as the ``to`` target in ``TeamSay``. + description (`str`): + One-sentence summary of the worker's role. + prompt (`str`): + First task delivered as a user message to the worker. + subagent_type (`str`, defaults to ``"default"``): + Template type to use. Must match a registered + :class:`SubAgentTemplate.type`. + _agent_state (`AgentState | None`, optional): + Live leader state injected by the toolkit. + + Returns: + `ToolChunk`: + A success message containing the new member id, or an + error chunk on failure. + """ + try: + session = await self._storage.get_session( + self._user_id, + self._agent_id, + self._session_id, + ) + if session is None or session.team_id is None: + return ToolChunk( + content=[ + TextBlock( + text=( + "AgentCreate: this session is not in " + "any team — call TeamCreate first." + ), + ), + ], + state=ToolResultState.ERROR, + ) + team = await self._storage.get_team( + self._user_id, + session.team_id, + ) + if team is None: + return ToolChunk( + content=[ + TextBlock( + text=( + "AgentCreate: team " + f"{session.team_id} no longer exists." + ), + ), + ], + state=ToolResultState.ERROR, + ) + if team.session_id != self._session_id: + return ToolChunk( + content=[ + TextBlock( + text=( + "AgentCreate: only the team leader " + "can add members; this session is a " + "worker." + ), + ), + ], + state=ToolResultState.ERROR, + ) + + # Look up leader session for chat-model inheritance + name. + leader_session = await self._storage.get_session( + self._user_id, + "", # agent_id unused at storage level + team.session_id, + ) + if leader_session is None: + return ToolChunk( + content=[ + TextBlock( + text=( + f"AgentCreate: leader session " + f"{team.session_id} for team {team.id} is " + f"missing — team is in an inconsistent " + f"state." + ), + ), + ], + state=ToolResultState.ERROR, + ) + + # Resolve the template. + template = self._sub_agent_templates.get(subagent_type) + if template is None: + available = list(self._sub_agent_templates) + return ToolChunk( + content=[ + TextBlock( + text=( + f"AgentCreate: unknown subagent_type " + f"{subagent_type!r}; expected one of " + f"{available}." + ), + ), + ], + state=ToolResultState.ERROR, + ) + + # Enforce team-scoped name uniqueness. TeamSay routes by + # ``name`` (not agent_id), so duplicates would be ambiguous + # and unaddressable. The leader's name participates too — + # workers must not collide with it. + leader_agent_record = await self._storage.get_agent( + self._user_id, + leader_session.agent_id, + ) + existing_names: set[str] = set() + if leader_agent_record is not None: + existing_names.add(leader_agent_record.data.name) + for member_id in team.data.member_ids: + member_record = await self._storage.get_agent( + self._user_id, + member_id, + ) + if member_record is not None: + existing_names.add(member_record.data.name) + if name in existing_names: + return ToolChunk( + content=[ + TextBlock( + text=( + f"AgentCreate: a team member named " + f"{name!r} already exists. Member names " + f"must be unique within the team " + f"(including the leader's name); pick " + f"another." + ), + ), + ], + state=ToolResultState.ERROR, + ) + + # Resolve leader name early — needed both for the system + # prompt template and for the initial team-message hint. + leader_name = ( + leader_agent_record.data.name + if leader_agent_record is not None + else leader_session.agent_id + ) + + # 1. Build worker AgentRecord (source="team" so it's hidden + # from the global agent list). + system_prompt = template.system_prompt_template.format( + team_name=team.data.name, + team_description=team.data.description, + member_name=name, + member_description=description, + leader_name=leader_name, + ) + worker_agent = AgentRecord( + user_id=self._user_id, + source="team", + data=AgentData( + name=name, + system_prompt=system_prompt, + context_config=template.context_config.model_copy( + deep=True, + ), + react_config=template.react_config.model_copy( + deep=True, + ), + ), + ) + await self._storage.upsert_agent(self._user_id, worker_agent) + + # 2. Build worker SessionRecord, inheriting leader's model + # config. The template's permission context is the base; + # on top of it we merge the leader's mode and/or rules + # and/or working directories according to the template's + # inherit-from-leader flags. See + # :func:`_merge_leader_permissions` for the policy. + leader_permission_context = ( + _agent_state.permission_context + if _agent_state is not None + else leader_session.state.permission_context + ) + worker_permission_context = _merge_leader_permissions( + template, + leader_permission_context, + ) + worker_state = AgentState( + permission_context=worker_permission_context, + tasks_context=template.tasks_context.model_copy( + deep=True, + ), + ) + worker_session = await self._storage.upsert_session( + user_id=self._user_id, + agent_id=worker_agent.id, + config=SessionConfig( + workspace_id=leader_session.config.workspace_id, + name=f"team:{team.id}/{name}", + chat_model_config=( + leader_session.config.chat_model_config + ), + fallback_chat_model_config=( + leader_session.config.fallback_chat_model_config + ), + ), + state=worker_state, + ) + await self._storage.set_session_team_id( + self._user_id, + worker_session.id, + team.id, + ) + + # 3. Append worker to team.member_ids. + team.data.member_ids = [ + *team.data.member_ids, + worker_agent.id, + ] + await self._storage.upsert_team(self._user_id, team) + + # 4. Deliver the initial task to the worker's inbox + wakeup. + hint = HintBlock( + hint=( + f'\n' + f"{prompt}\n" + f"" + ), + source=json.dumps( + { + "label": "team_message", + "sublabel": leader_name, + }, + ensure_ascii=False, + ), + ) + await self._message_bus.queue_push( + MessageBusKeys.inbox(worker_session.id), + hint.model_dump(mode="json"), + ) + await enqueue_run_trigger( + self._message_bus, + user_id=self._user_id, + session_id=worker_session.id, + agent_id=worker_agent.id, + ) + + return ToolChunk( + content=[ + TextBlock( + text=( + f"Member {name!r} added to team " + f"{team.data.name!r}." + ), + ), + ], + ) + except Exception as e: # pylint: disable=broad-except + return ToolChunk( + content=[TextBlock(text=f"AgentCreate failed: {e}")], + state=ToolResultState.ERROR, + ) diff --git a/src/agentscope/app/_tool/_team_create.py b/src/agentscope/app/_tool/_team_create.py new file mode 100644 index 0000000000000000000000000000000000000000..cda39816778375440d5a0199f3dc75042e84a663 --- /dev/null +++ b/src/agentscope/app/_tool/_team_create.py @@ -0,0 +1,145 @@ +# -*- coding: utf-8 -*- +"""The TeamCreate tool — establishes a new team led by the current session.""" +from pydantic import Field + +from ._team_tool_base import _TeamToolBase +from ..storage import TeamData, TeamRecord +from ...message import TextBlock, ToolResultState +from ...tool import ToolChunk, ParamsBase + + +class _TeamCreateParams(ParamsBase): + """Parameters for :class:`TeamCreate`.""" + + name: str = Field( + description=( + "Display name of the team. Used by the user to identify the " + "team and shown in the team UI." + ), + ) + description: str = Field( + description=( + "What the team is for — its overall goal or shared context. " + "This becomes the team's charter and is wired into every " + "member's system prompt so all members share the same " + "high-level understanding of why the team exists." + ), + ) + + +class TeamCreate(_TeamToolBase): + """Create a new team and become its leader.""" + + name: str = "TeamCreate" + + description: str = """Create a new team led by your current session and \ +return its team id. + +## When to Use This Tool +Use this tool when the task you've been given is best decomposed into \ +parallel sub-tasks executed by multiple specialised agents (members) \ +under your coordination. After creating the team, use ``AgentCreate`` to \ +spawn each member with its own role, prompt, and permission mode. NOTE: \ +the ``prompt`` you pass to ``AgentCreate`` is delivered to that member \ +automatically, so do **NOT** call ``TeamSay`` right after ``AgentCreate`` — \ +just wait for the members to report back. + +## When NOT to Use This Tool +- The task is small enough to handle yourself. +- You already lead a team in this session — a session can only lead \ +one team at a time. +""" + + input_schema: dict = _TeamCreateParams.model_json_schema() + + async def __call__( + self, + name: str, + description: str, + ) -> ToolChunk: + """Create the team directly via storage. + + Reads the current session record from storage to enforce the + precondition: a session can only lead one team at a time. + This makes the tool safe to attach unconditionally to + ``source='user'`` agents — calling it when the session + already leads a team returns a clear error rather than + silently corrupting state. + + Args: + name (`str`): + Display name of the team. + description (`str`): + Description / charter of the team. + + Returns: + `ToolChunk`: + A success message containing the team id, or an error + chunk if a precondition fails or creation failed. + """ + try: + session = await self._storage.get_session( + self._user_id, + self._agent_id, + self._session_id, + ) + if session is None: + return ToolChunk( + content=[ + TextBlock( + text=( + "TeamCreate: session " + f"{self._session_id} not found." + ), + ), + ], + state=ToolResultState.ERROR, + ) + if session.team_id is not None: + return ToolChunk( + content=[ + TextBlock( + text=( + "TeamCreate: this session is already " + f"part of team {session.team_id}. A " + "session can only lead one team at a " + "time — dissolve the current one with " + "TeamDelete first." + ), + ), + ], + state=ToolResultState.ERROR, + ) + + team = TeamRecord( + user_id=self._user_id, + session_id=self._session_id, + data=TeamData( + name=name, + description=description, + member_ids=[], + ), + ) + await self._storage.upsert_team(self._user_id, team) + await self._storage.set_session_team_id( + self._user_id, + self._session_id, + team.id, + ) + + return ToolChunk( + content=[ + TextBlock( + text=( + f"Team {team.id} ({team.data.name}) created. " + f"You are the leader. Use AgentCreate to add " + f"members, then TeamSay to coordinate them." + ), + ), + ], + ) + except Exception as e: # pylint: disable=broad-except + return ToolChunk( + content=[TextBlock(text=f"TeamCreate failed: {e}")], + state=ToolResultState.ERROR, + ) diff --git a/src/agentscope/app/_tool/_team_delete.py b/src/agentscope/app/_tool/_team_delete.py new file mode 100644 index 0000000000000000000000000000000000000000..d8f7efa428b70ef818ea5024bad99a71da62f74e --- /dev/null +++ b/src/agentscope/app/_tool/_team_delete.py @@ -0,0 +1,130 @@ +# -*- coding: utf-8 -*- +"""The TeamDelete tool — dissolves the team led by the current session.""" +from ._team_tool_base import _TeamToolBase +from ...message import TextBlock, ToolResultState +from ...tool import ToolChunk, ParamsBase + + +class _TeamDeleteParams(ParamsBase): + """Parameters for :class:`TeamDelete` — none.""" + + +class TeamDelete(_TeamToolBase): + """Dissolve the team you currently lead and clean up all members.""" + + name: str = "TeamDelete" + + description: str = """Dissolve the team you currently lead. + +## When to Use This Tool +- The team has finished its work and you want to clean up. +- The team is unrecoverably stuck and you want to start over. +- You have collected the deliverables you need from each member. + +## When NOT to Use This Tool +- Members are still producing useful output and you may want their \ +follow-up; dissolving deletes them and they cannot be revived. +- You want to remove only one specific member — there is no "remove \ +single member" tool in v1, only whole-team dissolution. + +## Effects +- Every member agent + its session is deleted. +- The team record is deleted. +- Your own session continues to exist but is no longer associated with \ +any team — the team-related tools become unavailable on subsequent \ +reasoning steps. + +This is irreversible. +""" + + input_schema: dict = _TeamDeleteParams.model_json_schema() + + async def __call__(self) -> ToolChunk: + """Dissolve the bound session's team via :class:`SessionService`. + + Reads the current session + team records from storage to + enforce: caller must be in a team AND must be its leader. + Then delegates the actual cancel + delete + bus-purge cascade + to :meth:`SessionService.delete_team`, which routes every + member through the shared session-level primitive — so worker + chat runs are cancelled cross-process and their bus state is + cleaned up the same way ``DELETE /agents`` would. + + Returns: + `ToolChunk`: + A confirmation message, or an error chunk if a + precondition fails. + """ + try: + session = await self._storage.get_session( + self._user_id, + self._agent_id, + self._session_id, + ) + if session is None or session.team_id is None: + return ToolChunk( + content=[ + TextBlock( + text=( + "TeamDelete: this session is not in " + "any team." + ), + ), + ], + state=ToolResultState.ERROR, + ) + team = await self._storage.get_team( + self._user_id, + session.team_id, + ) + if team is None: + return ToolChunk( + content=[ + TextBlock( + text=( + "TeamDelete: team " + f"{session.team_id} no longer exists." + ), + ), + ], + state=ToolResultState.ERROR, + ) + if team.session_id != self._session_id: + return ToolChunk( + content=[ + TextBlock( + text=( + "TeamDelete: only the team leader " + "can dissolve the team; this session " + "is a worker." + ), + ), + ], + state=ToolResultState.ERROR, + ) + + # Local import to avoid a circular dependency between + # ``_tools`` and ``_service`` at module load. + from .._service import SessionService # noqa: PLC0415 + + session_service = SessionService( + storage=self._storage, + message_bus=self._message_bus, + ) + await session_service.delete_team(self._user_id, team.id) + return ToolChunk( + content=[ + TextBlock( + text=( + f"Team {team.id} dissolved. All members " + f"deleted; your session is no longer " + f"leading any team." + ), + ), + ], + ) + except Exception as e: # pylint: disable=broad-except + return ToolChunk( + content=[TextBlock(text=f"TeamDelete failed: {e}")], + state=ToolResultState.ERROR, + ) diff --git a/src/agentscope/app/_tool/_team_say.py b/src/agentscope/app/_tool/_team_say.py new file mode 100644 index 0000000000000000000000000000000000000000..d14795602eef6748871d09293edb7be538dcec8f --- /dev/null +++ b/src/agentscope/app/_tool/_team_say.py @@ -0,0 +1,339 @@ +# -*- coding: utf-8 -*- +"""The TeamSay tool — sends a message to one or all team members.""" +from typing import Any + +from pydantic import Field + +from ._team_tool_base import _TeamToolBase +from ..message_bus import MessageBusKeys +from .._bus_ops import enqueue_run_trigger +from ...message import HintBlock, TextBlock, ToolResultState +from ...tool import ToolChunk, ParamsBase + + +class _TeamSayParams(ParamsBase): + """Parameters for :class:`TeamSay`.""" + + content: str = Field( + description=( + "The message text. Plain natural-language; the recipient " + "sees it as a user message in its context." + ), + ) + to: str | None = Field( + default=None, + description=( + "Recipient member name. Pass ``null`` (the default) to " + "broadcast to every other member of the team. To address " + "a specific peer use that member's name." + ), + ) + + +_LEADER_DESCRIPTION = """Send a message to a specific team member or \ +broadcast to all members. + +## When to Use This Tool +- Pass **new** requirements or context from the user to a specific member. +- Broadcast an update or coordination message to all members. +- Ask a member a follow-up question when you need clarification. + +## When NOT to Use This Tool +- DO NOT repeatedly call this to check on a member's progress — members \ +will automatically notify you via ``TeamSay`` when they finish their task. \ +Wait for their message instead of polling. +- DO NOT call this right after creating a member by ``AgentCreate``, the \ +member will receive its initial task from the ``prompt`` of the \ +``AgentCreate`` call and report back when done — just wait for their message. \ +- The session is not in a team yet (call ``TeamCreate`` first). +- You want to talk to yourself — use your own reasoning. + +## Important +- Each member starts working immediately when created via AgentCreate. \ +When a member finishes its task, it will call ``TeamSay`` to report results \ +back to you. You do NOT need to prompt them — just wait for their reply. +- **DO NOT** reply to a member's report message unless you have further \ +questions or requirements. ``TeamSay`` is for coordination, not chit-chat — \ +your top priority is to complete the overall task. +""" + +_WORKER_DESCRIPTION = """Send a message to the team leader or broadcast to \ +all team members. + +## When to Use This Tool +- **IMPORTANT**: When you finish your assigned task, you MUST call this \ +tool to report your results back to the leader. The leader is waiting \ +for your report — do not end your turn without sending it. +- Share intermediate findings or ask the leader for clarification. +- Broadcast information that other members might need. + +## When NOT to Use This Tool +- You want to talk to yourself — use your own reasoning. +- The message is a transient internal thought with no value to others. +""" + + +class TeamSay(_TeamToolBase): + """Send a message to a teammate (or broadcast to all teammates). + + Resolves the team membership at ``__call__`` time from storage, + so a member added moments earlier in the same chat run is + addressable immediately. + + The ``description`` shown to the agent differs by role: leaders + are reminded not to poll members, workers are reminded to report + results when done. The role is passed at construction time via + the ``role`` parameter. + """ + + name: str = "TeamSay" + description: str + + is_concurrency_safe: bool = True + is_read_only: bool = True + + input_schema: dict = _TeamSayParams.model_json_schema() + + def __init__( + self, + *args: Any, + role: str = "leader", + **kwargs: Any, + ) -> None: + """Initialise with role-specific description. + + Args: + role (`str`, defaults to ``"leader"``): + Either ``"leader"`` or ``"worker"``. Determines which + description the agent sees for this tool. + *args: + Forwarded to :class:`_TeamToolBase.__init__`. + **kwargs: + Forwarded to :class:`_TeamToolBase.__init__`. + """ + super().__init__(*args, **kwargs) + self.description = ( + _LEADER_DESCRIPTION if role == "leader" else _WORKER_DESCRIPTION + ) + + async def __call__( + self, + content: str, + to: str | None = None, + ) -> ToolChunk: + """Deliver the message directly via storage + message bus. + + Reads the current session record from storage to resolve the + team_id (the agent's team membership may have changed since + agent assembly), builds the team's (agent_id, session_id) + directory, and pushes a HintBlock + wakeup to each recipient. + + Args: + content (`str`): + Message body. + to (`str | None`, defaults to ``None``): + Display name of a specific team member to target, or + ``None`` for broadcast. Routing is by name (not + agent id) so workers can address the leader by the + name they see in ````. + Name uniqueness within a team is enforced at + ``AgentCreate`` time. + + Returns: + `ToolChunk`: + A confirmation containing the recipient count, or an + error chunk on failure. + """ + try: + session = await self._storage.get_session( + self._user_id, + self._agent_id, + self._session_id, + ) + if session is None or session.team_id is None: + return ToolChunk( + content=[ + TextBlock( + text=( + "TeamSay: this session is not in any " + "team — call TeamCreate first if you " + "want to start one." + ), + ), + ], + state=ToolResultState.ERROR, + ) + + team = await self._storage.get_team( + self._user_id, + session.team_id, + ) + if team is None: + return ToolChunk( + content=[ + TextBlock( + text=( + f"TeamSay: team {session.team_id} no longer " + f"exists." + ), + ), + ], + state=ToolResultState.ERROR, + ) + + leader_session = await self._storage.get_session( + self._user_id, + "", + team.session_id, + ) + if leader_session is None: + return ToolChunk( + content=[ + TextBlock( + text=( + f"TeamSay: leader session " + f"{team.session_id} missing for team " + f"{team.id}." + ), + ), + ], + state=ToolResultState.ERROR, + ) + + # Build a (name -> (session_id, agent_id)) directory in one + # pass over the team. Routing is by **name** rather than + # agent_id so workers can address the leader (they receive + # the leader's name in the hint + # but never see the leader's agent_id). Uniqueness of names + # within the team is enforced at AgentCreate time. + leader_agent = await self._storage.get_agent( + self._user_id, + leader_session.agent_id, + ) + leader_name = ( + leader_agent.data.name + if leader_agent is not None + else leader_session.agent_id + ) + directory: dict[str, tuple[str, str]] = { + leader_name: (leader_session.id, leader_session.agent_id), + } + for worker_agent_id in team.data.member_ids: + worker_agent = await self._storage.get_agent( + self._user_id, + worker_agent_id, + ) + if worker_agent is None: + continue + sessions = await self._storage.list_sessions( + self._user_id, + worker_agent_id, + ) + if sessions: + directory[worker_agent.data.name] = ( + sessions[0].id, + worker_agent_id, + ) + + own_session_ids = {sid for sid, _aid in directory.values()} + if self._session_id not in own_session_ids: + return ToolChunk( + content=[ + TextBlock( + text=( + f"TeamSay: this session " + f"({self._session_id}) is not part of " + f"team {team.id}." + ), + ), + ], + state=ToolResultState.ERROR, + ) + + if to is None: + recipients: list[tuple[str, str]] = [ + (sid, aid) + for sid, aid in directory.values() + if sid != self._session_id + ] + else: + resolved = directory.get(to) + if resolved is None: + known = sorted(directory.keys()) + return ToolChunk( + content=[ + TextBlock( + text=( + f"TeamSay: no team member is named " + f"{to!r}. Known members: {known}." + ), + ), + ], + state=ToolResultState.ERROR, + ) + target_session_id, target_agent_id = resolved + if target_session_id == self._session_id: + return ToolChunk( + content=[ + TextBlock( + text=( + "TeamSay: cannot send a message to " + "yourself; talk to yourself in your " + "own reasoning instead." + ), + ), + ], + state=ToolResultState.ERROR, + ) + recipients = [(target_session_id, target_agent_id)] + + # Resolve sender display name once. + sender_agent = await self._storage.get_agent( + self._user_id, + self._agent_id, + ) + sender_name = ( + sender_agent.data.name + if sender_agent is not None + else self._agent_id + ) + + hint = HintBlock( + hint=( + f'\n' + f"{content}\n" + f"" + ), + source=sender_name, + ) + payload = hint.model_dump(mode="json") + + for sid, aid in recipients: + await self._message_bus.queue_push( + MessageBusKeys.inbox(sid), + payload, + ) + await enqueue_run_trigger( + self._message_bus, + user_id=self._user_id, + session_id=sid, + agent_id=aid, + ) + + count = len(recipients) + target = "broadcast" if to is None else f"member {to!r}" + return ToolChunk( + content=[ + TextBlock( + text=( + f"Delivered to {count} recipient(s) " + f"({target})." + ), + ), + ], + ) + except Exception as e: # pylint: disable=broad-except + return ToolChunk( + content=[TextBlock(text=f"TeamSay failed: {e}")], + state=ToolResultState.ERROR, + ) diff --git a/src/agentscope/app/_tool/_team_tool_base.py b/src/agentscope/app/_tool/_team_tool_base.py new file mode 100644 index 0000000000000000000000000000000000000000..72c86df97b7654f0a765924c91301ddc0ad06418 --- /dev/null +++ b/src/agentscope/app/_tool/_team_tool_base.py @@ -0,0 +1,99 @@ +# -*- coding: utf-8 -*- +"""Base class shared by the team tools.""" +from typing import Any, TYPE_CHECKING + +from ...permission import ( + PermissionBehavior, + PermissionContext, + PermissionDecision, +) +from ...tool import ToolBase + +if TYPE_CHECKING: + from ..message_bus import MessageBus + from ..storage import StorageBase + + +class _TeamToolBase(ToolBase): + """Shared base for the team tools. + + All team tools are constructed at agent assembly time (in + :func:`get_toolkit`) with the request-scoped ``user_id``, + ``session_id``, and ``agent_id`` plus ``storage`` + ``message_bus`` + references. Each tool's ``__call__`` does its work directly via + those two dependencies — there is no intermediate service layer. + + Permissions: all team tools allow themselves unconditionally — the + agent's authority to call them is already gated by the + role/source-aware logic inside :func:`get_toolkit` that decides + which team tools to attach in the first place. + """ + + name: str + description: str + input_schema: dict[str, Any] + is_concurrency_safe: bool = False + is_read_only: bool = True + is_state_injected: bool = False + is_external_tool: bool = False + is_mcp: bool = False + mcp_name: str | None = None + + def __init__( + self, + storage: "StorageBase", + message_bus: "MessageBus", + user_id: str, + session_id: str, + agent_id: str, + ) -> None: + """Bind request-scoped identifiers and shared dependencies. + + Args: + storage (`StorageBase`): + Application storage. Each tool reads the current + session / team records at ``__call__`` time to do + runtime precondition checks (am I in a team? am I the + leader?) — this is what lets all four team tools be + attached unconditionally to ``source='user'`` agents + without depending on a stale snapshot of ``team_id`` + taken at agent assembly time. + message_bus (`MessageBus`): + Application message bus. Tools that deliver + inter-session messages (``AgentCreate``, ``TeamSay``) + push HintBlocks + wakeups through it. + user_id (`str`): + The owner user id of the calling agent. + session_id (`str`): + The current session id of the calling agent. + agent_id (`str`): + The id of the agent invoking the tool. + """ + self._storage = storage + self._message_bus = message_bus + self._user_id = user_id + self._session_id = session_id + self._agent_id = agent_id + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: PermissionContext, + ) -> PermissionDecision: + """Always allow — gating is done by tool-attachment logic. + + Args: + tool_input (`dict[str, Any]`): + The arguments the agent passed; ignored here. + context (`PermissionContext`): + The active permission context; ignored here. + + Returns: + `PermissionDecision`: + An ``ALLOW`` decision with a brief explanation. + """ + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message=f"{self.name} is always allowed when attached to the " + f"agent.", + ) diff --git a/src/agentscope/app/_types.py b/src/agentscope/app/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..d5a0f9ecf0a16e2b943868efe69fdcd0c40c7055 --- /dev/null +++ b/src/agentscope/app/_types.py @@ -0,0 +1,189 @@ +# -*- coding: utf-8 -*- +"""Shared type aliases for the agentscope app layer.""" +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Protocol + +from pydantic import BaseModel, Field + +from ..agent import ContextConfig, ReActConfig +from ..event import AgentEvent +from ..middleware import MiddlewareBase +from ..permission import PermissionContext +from ..state import TaskContext +from ..tool import ToolBase + +if TYPE_CHECKING: + from ._service._session_projection import SessionProjection + from .storage import AgentRecord, SessionRecord + + +AgentMiddlewareFactory = Callable[ + [str, str, str], + Awaitable[list[MiddlewareBase]], +] +# Async factory signature: ``(user_id, agent_id, session_id)`` → +# awaitable of :class:`~agentscope.middleware.MiddlewareBase` instances. + +AgentToolFactory = Callable[ + [str, str, str], + Awaitable[list[ToolBase]], +] +# Async factory signature: ``(user_id, agent_id, session_id)`` → +# awaitable of :class:`~agentscope.tool.ToolBase` instances. + + +class EventProjector(Protocol): + """Strategy that projects events from one session onto another. + + Each projector owns one cross-session UI feed (HITL, progress, + errors, …). :class:`~agentscope.app._service.ChatService` calls + every registered projector once per event produced by a run; a + projector decides whether the event is relevant and, if so, mirrors + it onto the owning session's target via the shared + :class:`SessionProjection` primitive. Adding a feed means adding a + projector — no new bus wrapper, dispatcher, or manager. + + Service dependencies a projector needs (e.g. storage, to resolve the + team a worker belongs to) are its own concern, injected at + construction. Only per-run request data is passed to + :meth:`maybe_project`. + """ + + async def maybe_project( + self, + user_id: str, + session_record: "SessionRecord", + agent_record: "AgentRecord", + event: AgentEvent, + projection: "SessionProjection", + ) -> None: + """Project ``event`` if it is relevant to this feed. + + Implementations must be a no-op for irrelevant events and + sessions. They should not raise for ordinary "not applicable" + cases; ``ChatService`` guards the call with a catch-all so a + projection failure never tears down the producing run, but + projectors should not rely on that for control flow. + + Args: + user_id (`str`): + The owner of the running session. + session_record (`SessionRecord`): + The currently-running session's record. + agent_record (`AgentRecord`): + The currently-running agent's record. + event (`AgentEvent`): + The event just published to this session's channel. + projection (`SessionProjection`): + Shared primitive used to write the durable entry and the + live notification. + """ + + +class SubAgentTemplate(BaseModel): + """A reusable blueprint for sub-agent creation within a team. + + Developers register one or more templates at the ``create_app`` entry + point. When the leader agent calls ``AgentCreate`` with a matching + ``subagent_type``, the template's configuration is used instead of the + built-in defaults. + + The :attr:`type` field serves as the routing key — it becomes an enum + value of the ``subagent_type`` parameter exposed to the LLM. This is + distinct from the ``name`` parameter in ``AgentCreate``, which is the + per-instance identifier the leader assigns to each worker (used for + ``TeamSay(to=name)``). + + All fields are pure data (no callables), so the template is fully + serializable for future config-driven startup. + """ + + type: str = Field( + description=( + "Template type identifier, e.g. ``'researcher'`` or " + "``'coder'``. Used as the enum value for the " + "``subagent_type`` parameter in ``AgentCreate``." + ), + ) + + description: str = Field( + description=( + "Agent-readable description of this sub-agent type. " + "Exposed to the LLM in the ``AgentCreate`` tool schema " + "so it can choose the appropriate type." + ), + ) + + system_prompt_template: str = Field( + description=( + "A Python format string for the worker's system prompt. " + "Available placeholders: ``{team_name}``, " + "``{team_description}``, ``{member_name}``, " + "``{member_description}``, ``{leader_name}``." + ), + ) + + context_config: ContextConfig = Field( + default_factory=ContextConfig, + description="Context configuration for the sub-agent.", + ) + + react_config: ReActConfig = Field( + default_factory=ReActConfig, + description="ReAct loop configuration for the sub-agent.", + ) + + permission_context: PermissionContext = Field( + default_factory=PermissionContext, + description=( + "Permission context applied to the sub-agent at " + "creation time. Controls what the worker is allowed " + "to do (e.g. read-only vs. full access)." + ), + ) + + override_leader_mode: bool = Field( + default=False, + description=( + "Whether the template's :attr:`permission_context.mode` " + "should override the leader session's mode for the worker. " + "``True`` — the worker runs in the template's mode " + "(typical for templates that pin a specific posture, e.g. " + "a read-only research worker). ``False`` (default) — the " + "worker inherits the leader's current mode." + ), + ) + + extend_leader_permission_rules: bool = Field( + default=True, + description=( + "Whether the leader session's allow/deny/ask permission " + "rules should be merged on top of the template's. " + "``True`` (default) — leader rules are appended after " + "the template's rules for each tool, so the worker " + "doesn't re-prompt for permissions the user has already " + "confirmed; the template's rules take precedence on " + "evaluation order. ``False`` — the template's rules are " + "the worker's complete rule set." + ), + ) + + extend_leader_working_directories: bool = Field( + default=True, + description=( + "Whether the leader session's working directories should " + "be merged into the template's. ``True`` (default) — " + "leader directories are added for keys not already in the " + "template (template wins on key collisions). ``False`` — " + "the template's working directories are the worker's " + "complete set." + ), + ) + + tasks_context: TaskContext = Field( + default_factory=TaskContext, + description=( + "Pre-defined task context for the sub-agent, allowing " + "the template to seed an initial workflow." + ), + ) diff --git a/src/agentscope/app/deps.py b/src/agentscope/app/deps.py new file mode 100644 index 0000000000000000000000000000000000000000..a6a1d24d04f674c86fb5072f61559c0a80090576 --- /dev/null +++ b/src/agentscope/app/deps.py @@ -0,0 +1,293 @@ +# -*- coding: utf-8 -*- +"""Shared FastAPI dependencies for the agentscope app.""" +from fastapi import Header, HTTPException, Request, status + +from .workspace_manager import WorkspaceManagerBase +from ._manager import ( + BackgroundTaskManager, + ChatRunRegistry, + SchedulerManager, +) +from ._service import ChatService, KnowledgeBaseService, SessionService +from ._types import AgentMiddlewareFactory, AgentToolFactory +from .message_bus import MessageBus +from .rag.blob_store import BlobStoreBase +from .rag.knowledge_base_manager import KnowledgeBaseManagerBase +from .storage import StorageBase +from ..rag import ParserBase + + +async def get_current_user_id( + x_user_id: str = Header( + description="Caller's user ID. " + "Temporary header-based identity; will be replaced by JWT auth.", + ), +) -> str: + """Return the caller's user ID from the ``X-User-ID`` request header. + + Args: + x_user_id (`str`): Value of the ``X-User-ID`` header. + + Returns: + `str`: The authenticated user ID. + + Raises: + `HTTPException`: 401 if the header is missing or empty. + """ + if not x_user_id: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="X-User-ID header is required.", + ) + return x_user_id + + +async def get_storage(request: Request) -> StorageBase: + """Return the application-wide storage backend. + + Args: + request (`Request`): The incoming FastAPI request. + + Returns: + `StorageBase`: The storage instance stored in ``app.state``. + """ + return request.app.state.storage + + +async def get_message_bus(request: Request) -> MessageBus: + """Return the application-wide message bus. + + Args: + request (`Request`): The incoming FastAPI request. + + Returns: + `MessageBus`: The message bus instance stored in ``app.state``. + """ + return request.app.state.message_bus + + +async def get_chat_service(request: Request) -> ChatService: + """Return the application-wide chat service. + + Args: + request (`Request`): The incoming FastAPI request. + + Returns: + `ChatService`: The chat service instance stored in ``app.state``. + """ + return request.app.state.chat_service + + +async def get_session_service(request: Request) -> SessionService: + """Return the application-wide session service. + + Args: + request (`Request`): The incoming FastAPI request. + + Returns: + `SessionService`: The session service instance stored in + ``app.state``. + """ + return request.app.state.session_service + + +async def get_chat_run_registry(request: Request) -> ChatRunRegistry: + """Return the per-process chat-run registry. + + Args: + request (`Request`): The incoming FastAPI request. + + Returns: + `ChatRunRegistry`: The registry stored in ``app.state``. + """ + return request.app.state.chat_run_registry + + +async def get_scheduler_manager(request: Request) -> SchedulerManager: + """Return the application-wide scheduler manager. + + Args: + request (`Request`): The incoming FastAPI request. + + Returns: + `SchedulerManager`: The scheduler manager stored in ``app.state``. + """ + return request.app.state.scheduler_manager + + +async def get_background_task_manager( + request: Request, +) -> BackgroundTaskManager: + """Return the application-wide background task manager. + + Args: + request (`Request`): The incoming FastAPI request. + + Returns: + `BackgroundTaskManager`: The background task manager stored in + ``app.state``. + """ + return request.app.state.background_task_manager + + +async def get_workspace_manager(request: Request) -> WorkspaceManagerBase: + """Return the application-wide workspace manager. + + Args: + request (`Request`): The incoming FastAPI request. + + Returns: + `WorkspaceManagerBase`: The workspace manager stored in ``app.state``. + """ + return request.app.state.workspace_manager + + +async def get_extra_agent_middlewares( + request: Request, +) -> AgentMiddlewareFactory | None: + """Return the caller-supplied agent middleware factory, if any. + + Args: + request (`Request`): The incoming FastAPI request. + + Returns: + `AgentMiddlewareFactory | None`: The factory passed to + :func:`~agentscope.app.create_app`, or ``None`` if not configured. + """ + return request.app.state.extra_agent_middlewares + + +async def get_extra_agent_tools( + request: Request, +) -> AgentToolFactory | None: + """Return the caller-supplied agent tool factory, if any. + + Args: + request (`Request`): The incoming FastAPI request. + + Returns: + `AgentToolFactory | None`: The factory passed to + :func:`~agentscope.app.create_app`, or ``None`` if not configured. + """ + return request.app.state.extra_agent_tools + + +async def get_knowledge_base_service( + request: Request, +) -> KnowledgeBaseService: + """Return the application-wide knowledge base service. + + Args: + request (`Request`): + The incoming FastAPI request. + + Returns: + `KnowledgeBaseService`: + The service stored in ``app.state``. + + Raises: + `HTTPException`: + ``503`` when the app was created without a + ``knowledge_base_manager`` and therefore exposes no + knowledge base endpoints. + """ + service = getattr(request.app.state, "knowledge_base_service", None) + if service is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=( + "Knowledge base feature is disabled — pass a " + "knowledge_base_manager to create_app() to enable it." + ), + ) + return service + + +async def get_knowledge_base_manager( + request: Request, +) -> KnowledgeBaseManagerBase: + """Return the application-wide knowledge base manager. + + Args: + request (`Request`): + The incoming FastAPI request. + + Returns: + `KnowledgeBaseManagerBase`: + The manager stored in ``app.state``. + + Raises: + `HTTPException`: + ``503`` when the app was created without a + ``knowledge_base_manager``. + """ + manager = getattr(request.app.state, "knowledge_base_manager", None) + if manager is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=( + "Knowledge base feature is disabled — pass a " + "knowledge_base_manager to create_app() to enable it." + ), + ) + return manager + + +async def get_blob_store(request: Request) -> BlobStoreBase: + """Return the application-wide blob store. + + Args: + request (`Request`): + The incoming FastAPI request. + + Returns: + `BlobStoreBase`: + The blob store instance stored in ``app.state``. + + Raises: + `HTTPException`: + ``503`` when no blob store is configured (e.g. the KB + feature was disabled at app-creation time). + """ + blob_store = getattr(request.app.state, "blob_store", None) + if blob_store is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=( + "Blob store is not configured — pass a " + "knowledge_base_manager (and optionally a blob_store) " + "to create_app() to enable knowledge base features." + ), + ) + return blob_store + + +async def get_knowledge_parsers( + request: Request, +) -> list[ParserBase] | dict[str, ParserBase]: + """Return the parser registry configured on the app. + + Args: + request (`Request`): + The incoming FastAPI request. + + Returns: + `list[ParserBase] | dict[str, ParserBase]`: + The parser registry stored in ``app.state.knowledge_parsers`` + — the same value the index worker uses to dispatch uploads. + + Raises: + `HTTPException`: + ``503`` when the KB feature is disabled (no parsers + configured). + """ + parsers = getattr(request.app.state, "knowledge_parsers", None) + if not parsers: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=( + "Knowledge base feature is disabled — pass a " + "knowledge_base_manager to create_app() to enable it." + ), + ) + return parsers diff --git a/src/agentscope/app/message_bus/__init__.py b/src/agentscope/app/message_bus/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e132a8f215fde3cab29f7623899efedc1dcd4480 --- /dev/null +++ b/src/agentscope/app/message_bus/__init__.py @@ -0,0 +1,14 @@ +# -*- coding: utf-8 -*- +"""The message bus module — live transport for cross-session messages.""" + +from ._base import MessageBus +from ._in_memory_message_bus import InMemoryMessageBus +from ._keys import MessageBusKeys +from ._redis_message_bus import RedisMessageBus + +__all__ = [ + "InMemoryMessageBus", + "MessageBus", + "MessageBusKeys", + "RedisMessageBus", +] diff --git a/src/agentscope/app/message_bus/_base.py b/src/agentscope/app/message_bus/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..cb48d1d8157569e29f2fc26c05323c85bcd220a2 --- /dev/null +++ b/src/agentscope/app/message_bus/_base.py @@ -0,0 +1,810 @@ +# -*- coding: utf-8 -*- +"""The message bus abstract base class. + +The message bus is the *live* transport layer used to coordinate work +across sessions and processes. It is intentionally separate from +:class:`StorageBase`, which owns *persistent* records: storage may live +on a relational database while the bus stays on a push-capable backend +(Redis, NATS, …) where waking idle consumers and fanning out events is +cheap. + +The interface is grouped by **consumption semantics** — i.e. how a +payload's lifetime ends — rather than by business use case. Callers map +their own concepts (agent inbox, session SSE replay, idle wake-up, …) +onto the right mode plus a key naming convention they own. + +Three orthogonal modes are exposed: + +============================ =========================================== +Mode A — drain queue Mode C — replay log +``queue_push`` / ``log_append`` / +``queue_drain`` ``log_read`` / ``log_trim`` + +Single-consumer, ack-on-read. Multi-consumer, externally bounded. +Each entry returned at most Each reader tracks its own cursor; +once; storage drops it the entries persist until trimmed, +moment it is read. TTL bounds ``max_len`` truncates from the head, +orphaned data when the consumer or TTL expires the whole key. +disappears. +============================ =========================================== + +Mode D — transient broadcast: ``publish`` / ``subscribe``. Fire-and-forget +pub/sub; only currently-subscribed listeners receive a payload, no +history. Use for wake-up signals where missed-while-offline is fine. + +Counted broadcast (one entry consumed by N distinct readers) is +intentionally not exposed as a primitive: in practice it requires +consumer-group coordination (Redis Streams' XREADGROUP/XACK semantics) +and adds substantial state. Producers wanting "fan out to N members" +should fan out at write time — push one entry per recipient inbox using +Mode A. The bus stays simple; deduplication is the producer's +responsibility. +""" +from abc import ABC, abstractmethod +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from typing import Any, Callable, Self + +from typing_extensions import deprecated + +from ._keys import MessageBusKeys + + +class MessageBus(ABC): # pylint: disable=too-many-public-methods + """Abstract base class for live message transport. + + Implementations expose three consumption modes (drain queue, replay + log, transient broadcast) over arbitrary string keys and JSON-style + dict payloads. Callers own key naming and payload schemas. + """ + + async def __aenter__(self) -> Self: + """Open underlying transport resources (connection pools, …). + + Returns: + `Self`: + The bus instance, for use as an async context manager. + """ + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: Any, + ) -> None: + """Release transport resources on context exit. + + Args: + exc_type (`type[BaseException] | None`): + The exception class raised inside the context, if any. + exc_value (`BaseException | None`): + The exception instance raised inside the context, if any. + traceback (`Any`): + The traceback associated with the exception, if any. + """ + await self.aclose() + + async def aclose(self) -> None: + """Release underlying transport resources. Default is a no-op.""" + + # ------------------------------------------------------------------ + # Mode A — drain queue (single consumer, ack-on-read) + # ------------------------------------------------------------------ + + @abstractmethod + async def queue_push( + self, + key: str, + payload: dict, + *, + ttl_secs: int | None = None, + ) -> str: + """Append ``payload`` to the drain queue at ``key``. + + Drain queues are single-consumer, ack-on-read: a subsequent + :meth:`queue_drain` returns each entry exactly once and then + deletes it. ``ttl_secs`` bounds the queue's lifetime so a key + whose consumer disappears does not accumulate entries forever. + + Args: + key (`str`): + Queue identifier. Caller-defined naming convention; the + bus treats it as opaque. + payload (`dict`): + JSON-serializable dict to enqueue. Schema is the + caller's responsibility. + ttl_secs (`int | None`, optional): + If set, refresh the queue key's expiry to this many + seconds on every push (sliding TTL). When ``None``, the + key never expires and must be drained or deleted + explicitly. + + Returns: + `str`: + Transport-level entry id (e.g. the Redis Stream entry + id). Useful for tracing; not required for normal + consumption. + """ + + @abstractmethod + async def queue_drain( + self, + key: str, + max_count: int = 100, + ) -> list[tuple[str, dict]]: + """Drain up to ``max_count`` entries from the queue at ``key``. + + Returned entries are removed from the queue in the same + operation, so a subsequent call returns only entries that + arrived after this one. Safe under the single-consumer-per-key + invariant. + + Args: + key (`str`): + Queue identifier. + max_count (`int`, defaults to ``100``): + Maximum number of entries to return in one call. + Older entries are returned first; remaining entries + stay in the queue for the next call. + + Returns: + `list[tuple[str, dict]]`: + ``(entry_id, payload)`` pairs in arrival order. Empty + list when the queue is empty. + """ + + @abstractmethod + async def queue_delete(self, key: str) -> None: + """Delete the drain queue at ``key`` and all of its entries. + + Idempotent: a no-op when the key does not exist. Used by + :meth:`session_purge` to drop a session's inbox during + cascade-delete. + + Args: + key (`str`): + Queue identifier. + """ + + # ------------------------------------------------------------------ + # Mode C — replay log (multi-consumer, externally bounded) + # ------------------------------------------------------------------ + + @abstractmethod + async def log_append( + self, + key: str, + payload: dict, + *, + ttl_secs: int | None = None, + max_len: int | None = None, + ) -> str: + """Append ``payload`` to the replay log at ``key``. + + Replay logs are append-only; readers track their own cursor and + may join at any time. Lifetime is bounded externally: by + ``ttl_secs`` (whole key expires), by ``max_len`` (oldest + entries trimmed once the log exceeds the cap), or by explicit + :meth:`log_trim`. + + Args: + key (`str`): + Log identifier. + payload (`dict`): + JSON-serializable dict to append. + ttl_secs (`int | None`, optional): + If set, refresh the key's expiry to this many seconds + on every append (sliding TTL). ``None`` means no TTL. + max_len (`int | None`, optional): + If set, cap the log at approximately this many entries + — older entries are trimmed when the cap is exceeded. + The cap is approximate so the operation can use the + backend's efficient near-trim mode (e.g. ``XADD MAXLEN + ~N``). ``None`` means no cap. + + Returns: + `str`: + Transport-level entry id, useful as a cursor for + subsequent :meth:`log_read` calls. + """ + + @abstractmethod + async def log_read( + self, + key: str, + since: str | None = None, + max_count: int = 100, + ) -> list[tuple[str, dict]]: + """Read up to ``max_count`` entries from the replay log at + ``key``, starting after ``since``. + + Reads are non-destructive: the same entries can be returned to + any number of readers, or to the same reader on retry. Each + reader is responsible for tracking its own cursor. + + Args: + key (`str`): + Log identifier. + since (`str | None`, optional): + Cursor — return entries strictly newer than this id. + Pass the last ``entry_id`` from the previous read. + ``None`` reads from the beginning of the log. + max_count (`int`, defaults to ``100``): + Maximum number of entries to return. + + Returns: + `list[tuple[str, dict]]`: + ``(entry_id, payload)`` pairs in append order. Empty + list when no entries are newer than ``since``. + """ + + @abstractmethod + async def log_trim( + self, + key: str, + before_id: str | None = None, + ) -> None: + """Trim the replay log at ``key``. + + Args: + key (`str`): + Log identifier. + before_id (`str | None`, optional): + Drop all entries with id strictly older than this. + ``None`` drops the entire log (i.e. deletes the key). + """ + + # ------------------------------------------------------------------ + # Mode D — transient broadcast (fire-and-forget) + # ------------------------------------------------------------------ + + @abstractmethod + async def publish( + self, + key: str, + payload: dict, + ) -> None: + """Publish ``payload`` on the broadcast channel ``key``. + + Only subscribers connected at the moment of publish receive the + payload — no history is retained. Use for wake-up signals or + short-lived notifications where missed-while-offline is + acceptable. + + Args: + key (`str`): + Channel identifier. + payload (`dict`): + JSON-serializable dict delivered as-is to subscribers. + """ + + @abstractmethod + async def subscribe( + self, + key: str, + *, + on_ready: Callable[[], None] | None = None, + ) -> AsyncGenerator[dict, None]: + """Yield broadcast payloads for ``key`` until the consumer + closes the generator. + + Subscriptions are best-effort: only payloads published *after* + the subscription is established are delivered. Callers own the + generator's lifetime — closing it releases the underlying + subscription. + + Args: + key (`str`): + Channel identifier. + on_ready (`Callable[[], None] | None`, optional): + If supplied, invoked exactly once after the underlying + subscription is established and before any payload is + yielded. Used by callers that need to block a + bootstrapping step (e.g. + ``SessionTriggerListenerManager.start``) until the + subscription is live, so a publish-immediately-after + race is impossible. + + Yields: + `dict`: + Each payload originally passed to :meth:`publish`. + """ + # The empty `yield` makes Python treat this as an async generator + # function (return type AsyncGenerator) rather than a coroutine + # returning an AsyncGenerator. Subclasses override it; this body + # never runs. + if False: # pylint: disable=using-constant-test + yield # pylint: disable=unreachable + + # ------------------------------------------------------------------ + # Mode E — distributed lock (cluster-wide mutex) + # ------------------------------------------------------------------ + + @abstractmethod + @asynccontextmanager + async def acquire_lock( + self, + key: str, + *, + ttl_secs: int = 600, + ) -> AsyncGenerator[None, None]: + """Acquire a distributed mutex on ``key``. + + Blocks until the lock is acquired, then yields. The + implementation maintains the lock across long-running + bodies (typically with a heartbeat task that renews the + TTL) so the lease only expires if the holding process + actually crashes — at which point another acquirer may + take over after at most ``ttl_secs``. + + Args: + key (`str`): + Lock identifier. + ttl_secs (`int`, defaults to ``600``): + Lease duration in seconds. The implementation + should renew this periodically while the body + runs; callers do not need to. + + Yields: + `None`: while the lock is held. + """ + # The decorator-based abstract method requires a body for + # @asynccontextmanager to work; subclasses override it. + if False: # pylint: disable=using-constant-test + yield # pylint: disable=unreachable + + @abstractmethod + async def is_locked(self, key: str) -> bool: + """Return whether ``key`` currently holds a lock. + + Args: + key (`str`): + Lock identifier (same key passed to + :meth:`acquire_lock`). + + Returns: + `bool`: + ``True`` if some process holds the lock right now. + """ + + # ------------------------------------------------------------------ + # Mode F — registry map (hash-keyed namespace) + # ------------------------------------------------------------------ + + @abstractmethod + async def registry_set( + self, + namespace: str, + field: str, + value: str, + *, + ttl_secs: int | None = None, + ) -> None: + """Set ``field`` to ``value`` in the registry at ``namespace``. + + If the namespace does not exist it is created. When + ``ttl_secs`` is supplied, the namespace's TTL is refreshed + (sliding) — individual fields do not carry independent TTLs. + + Args: + namespace (`str`): + Registry key (e.g. ``"agentscope:bg_tasks:sess123"``). + field (`str`): + Field name within the registry. + value (`str`): + Serialized value to store. + ttl_secs (`int | None`, optional): + Refresh the namespace expiry to this many seconds. + """ + + @abstractmethod + async def registry_del(self, namespace: str, field: str) -> None: + """Remove ``field`` from the registry at ``namespace``. + + A no-op when the field or namespace does not exist. + + Args: + namespace (`str`): + Registry key. + field (`str`): + Field to remove. + """ + + @abstractmethod + async def registry_exists(self, namespace: str, field: str) -> bool: + """Return whether ``field`` exists in the registry at + ``namespace``. + + Args: + namespace (`str`): + Registry key. + field (`str`): + Field to check. + + Returns: + `bool`: + ``True`` if the field is present. + """ + + @abstractmethod + async def registry_getall( + self, + namespace: str, + ) -> dict[str, str]: + """Return all field-value pairs in the registry at + ``namespace``. + + Args: + namespace (`str`): + Registry key. + + Returns: + `dict[str, str]`: + All entries. Empty dict when the namespace is absent. + """ + + @abstractmethod + async def registry_drop(self, namespace: str) -> None: + """Delete the entire registry at ``namespace``. + + Idempotent: a no-op when the namespace does not exist. + + Args: + namespace (`str`): + Registry key to delete. + """ + + # ================================================================== + # Deprecated domain helpers + # + # These thin shells delegate to the generic primitives above. They + # exist so that code written against the old API keeps working for + # one release cycle; new code should use the primitives + MessageBusKeys + # (or the standalone functions in agentscope.app._service) directly. + # + # The _XXX_KEY class-level constants are kept as well — some tests + # reference them — but new code should use MessageBusKeys instead. + # ================================================================== + + # Key constants (kept for backward compat) ------------------------- + + _SESSION_LOCK_KEY = "agentscope:session:lock:{sid}" + _SESSION_EVENTS_KEY = "agentscope:session:events:{sid}" + _SESSION_CANCEL_KEY = "agentscope:session:cancel" + _SESSION_RUN_TTL_SECS = 600 + _SESSION_REPLAY_MAX_LEN = 1000 + _INBOX_KEY = "agentscope:inbox:{sid}" + _WAKEUP_QUEUE_KEY = "agentscope:wakeups" + _WAKEUP_SIGNAL_KEY = "agentscope:wakeup_signal" + _BG_TASKS_KEY = "agentscope:bg_tasks:{sid}" + _BG_TASKS_TTL_SECS = 86400 + _TASK_CANCEL_KEY = "agentscope:task:cancel" + + # Session run coordination ----------------------------------------- + + @deprecated( + "Use acquire_lock(MessageBusKeys.session_lock(sid), ...) directly.", + ) + @asynccontextmanager + async def session_run(self, session_id: str) -> AsyncGenerator[None, None]: + """Acquire the session lock, yield, then trim the replay log.""" + async with self.acquire_lock( + self._SESSION_LOCK_KEY.format(sid=session_id), + ttl_secs=self._SESSION_RUN_TTL_SECS, + ): + try: + yield + finally: + await self.log_trim( + self._SESSION_EVENTS_KEY.format(sid=session_id), + ) + + @deprecated( + "Use is_locked(MessageBusKeys.session_lock(sid)) directly.", + ) + async def session_is_running(self, session_id: str) -> bool: + """Check whether some process holds the session lock.""" + return await self.is_locked( + self._SESSION_LOCK_KEY.format(sid=session_id), + ) + + @deprecated( + "Use publish_session_event(bus, sid, event) from " + "agentscope.app._bus_ops directly.", + ) + async def session_publish_event( + self, + session_id: str, + event: dict, + ) -> str: + """Append + fan-out a session event.""" + key = self._SESSION_EVENTS_KEY.format(sid=session_id) + entry_id = await self.log_append( + key, + event, + max_len=self._SESSION_REPLAY_MAX_LEN, + ) + await self.publish(key, {**event, "_entry_id": entry_id}) + return entry_id + + @deprecated( + "Use log_read(MessageBusKeys.session_events(sid), ...) directly.", + ) + async def session_read_events( + self, + session_id: str, + since: str | None = None, + max_count: int = 1000, + ) -> list[tuple[str, dict]]: + """Read events from the session's replay log.""" + return await self.log_read( + self._SESSION_EVENTS_KEY.format(sid=session_id), + since=since, + max_count=max_count, + ) + + @deprecated( + "Use subscribe(MessageBusKeys.session_events(sid), ...) directly, " + "stripping _entry_id from each payload.", + ) + async def session_subscribe_events( + self, + session_id: str, + *, + on_ready: Callable[[], None] | None = None, + ) -> AsyncGenerator[dict, None]: + """Live-subscribe to session events, stripping _entry_id.""" + key = self._SESSION_EVENTS_KEY.format(sid=session_id) + async for payload in self.subscribe(key, on_ready=on_ready): + yield {k: v for k, v in payload.items() if k != "_entry_id"} + + # Cross-process cancel --------------------------------------------- + + @deprecated( + "Use publish(MessageBusKeys.session_cancel_channel(), " + "{'session_id': sid}) directly.", + ) + async def session_publish_cancel(self, session_id: str) -> None: + """Broadcast a session cancel request.""" + await self.publish( + self._SESSION_CANCEL_KEY, + {"session_id": session_id}, + ) + + @deprecated( + "Use subscribe(MessageBusKeys.session_cancel_channel(), ...) " + "directly, extracting session_id from the payload.", + ) + async def session_subscribe_cancel( + self, + *, + on_ready: Callable[[], None] | None = None, + ) -> AsyncGenerator[str, None]: + """Subscribe to session cancel broadcasts, yielding session ids.""" + async for payload in self.subscribe( + self._SESSION_CANCEL_KEY, + on_ready=on_ready, + ): + sid = payload.get("session_id") + if isinstance(sid, str): + yield sid + + # Purge ------------------------------------------------------------- + + @deprecated( + "Call log_trim / queue_delete / registry_drop with " + "MessageBusKeys directly.", + ) + async def session_purge(self, session_id: str) -> None: + """Delete all per-session bus state.""" + await self.log_trim(self._SESSION_EVENTS_KEY.format(sid=session_id)) + await self.queue_delete(self._INBOX_KEY.format(sid=session_id)) + await self.registry_drop(self._BG_TASKS_KEY.format(sid=session_id)) + + # Inbox ----------------------------------------------------------- + + @deprecated( + "Use queue_push(MessageBusKeys.inbox(sid), ...) directly.", + ) + async def inbox_push( + self, + session_id: str, + msg: dict, + *, + ttl_secs: int | None = None, + ) -> str: + """Push a message to a session's inbox.""" + return await self.queue_push( + self._INBOX_KEY.format(sid=session_id), + msg, + ttl_secs=ttl_secs, + ) + + @deprecated( + "Use queue_drain(MessageBusKeys.inbox(sid), ...) directly.", + ) + async def inbox_drain( + self, + session_id: str, + max_count: int = 100, + ) -> list[tuple[str, dict]]: + """Drain pending inbox messages for a session.""" + return await self.queue_drain( + self._INBOX_KEY.format(sid=session_id), + max_count=max_count, + ) + + # Wakeup ---------------------------------------------------------- + + @deprecated( + "Use enqueue_run_trigger(bus, ...) from " + "agentscope.app._bus_ops directly.", + ) + async def enqueue_wakeup( + self, + user_id: str, + session_id: str, + agent_id: str, + ) -> None: + """Enqueue an idle-session wake-up.""" + await self.queue_push( + self._WAKEUP_QUEUE_KEY, + { + "user_id": user_id, + "session_id": session_id, + "agent_id": agent_id, + "kind": MessageBusKeys.WAKEUP_KIND_WAKE, + "input": None, + }, + ) + await self.publish(self._WAKEUP_SIGNAL_KEY, {}) + + @deprecated( + "Use enqueue_run_trigger(bus, ...) from " + "agentscope.app._bus_ops directly.", + ) + async def enqueue_input( + self, + user_id: str, + session_id: str, + agent_id: str, + *, + kind: str, + inputs: dict | None = None, + ) -> None: + """Enqueue a typed run trigger.""" + await self.queue_push( + self._WAKEUP_QUEUE_KEY, + { + "user_id": user_id, + "session_id": session_id, + "agent_id": agent_id, + "kind": kind, + "input": inputs, + }, + ) + await self.publish(self._WAKEUP_SIGNAL_KEY, {}) + + @deprecated( + "Use queue_drain(MessageBusKeys.wakeup_queue(), ...) directly.", + ) + async def dequeue_wakeups( + self, + max_count: int = 64, + ) -> list[dict]: + """Drain pending run-trigger entries.""" + entries = await self.queue_drain( + self._WAKEUP_QUEUE_KEY, + max_count=max_count, + ) + return [payload for _entry_id, payload in entries] + + @deprecated( + "Use subscribe(MessageBusKeys.wakeup_signal(), ...) directly.", + ) + async def subscribe_wakeup_signal( + self, + *, + on_ready: Callable[[], None] | None = None, + ) -> AsyncGenerator[dict, None]: + """Subscribe to the shared wake-up signal channel.""" + async for payload in self.subscribe( + self._WAKEUP_SIGNAL_KEY, + on_ready=on_ready, + ): + yield payload + + # Background task registry ------------------------------------------- + + @deprecated( + "Use registry_set(MessageBusKeys.bg_tasks(sid), ...) directly.", + ) + async def bg_task_register( + self, + session_id: str, + task_id: str, + metadata: str, + ) -> None: + """Register a background task.""" + await self.registry_set( + self._BG_TASKS_KEY.format(sid=session_id), + task_id, + metadata, + ttl_secs=self._BG_TASKS_TTL_SECS, + ) + + @deprecated( + "Use registry_del(MessageBusKeys.bg_tasks(sid), tid) directly.", + ) + async def bg_task_unregister( + self, + session_id: str, + task_id: str, + ) -> None: + """Unregister a background task.""" + await self.registry_del( + self._BG_TASKS_KEY.format(sid=session_id), + task_id, + ) + + @deprecated( + "Use registry_exists(MessageBusKeys.bg_tasks(sid), tid) directly.", + ) + async def bg_task_exists( + self, + session_id: str, + task_id: str, + ) -> bool: + """Check whether a background task is registered.""" + return await self.registry_exists( + self._BG_TASKS_KEY.format(sid=session_id), + task_id, + ) + + @deprecated( + "Use registry_getall(MessageBusKeys.bg_tasks(sid)) directly.", + ) + async def bg_task_list( + self, + session_id: str, + ) -> dict[str, str]: + """List all background tasks for a session.""" + return await self.registry_getall( + self._BG_TASKS_KEY.format(sid=session_id), + ) + + @deprecated( + "Use registry_drop(MessageBusKeys.bg_tasks(sid)) directly.", + ) + async def bg_task_purge(self, session_id: str) -> None: + """Delete all background task entries for a session.""" + await self.registry_drop( + self._BG_TASKS_KEY.format(sid=session_id), + ) + + @deprecated( + "Use publish(MessageBusKeys.task_cancel_channel(), " + "{'task_id': tid}) directly.", + ) + async def task_publish_cancel(self, task_id: str) -> None: + """Broadcast a cancel request for a single background task.""" + await self.publish(self._TASK_CANCEL_KEY, {"task_id": task_id}) + + @deprecated( + "Use subscribe(MessageBusKeys.task_cancel_channel(), ...) " + "directly, extracting task_id from the payload.", + ) + async def task_subscribe_cancel( + self, + *, + on_ready: Callable[[], None] | None = None, + ) -> AsyncGenerator[str, None]: + """Subscribe to task cancel broadcasts, yielding task ids.""" + async for payload in self.subscribe( + self._TASK_CANCEL_KEY, + on_ready=on_ready, + ): + tid = payload.get("task_id") + if isinstance(tid, str): + yield tid diff --git a/src/agentscope/app/message_bus/_in_memory_message_bus.py b/src/agentscope/app/message_bus/_in_memory_message_bus.py new file mode 100644 index 0000000000000000000000000000000000000000..78ed7746314f345b8b2c08dc668660cc79b84af6 --- /dev/null +++ b/src/agentscope/app/message_bus/_in_memory_message_bus.py @@ -0,0 +1,470 @@ +# -*- coding: utf-8 -*- +"""In-memory message bus implementation. + +A pure-Python :class:`MessageBus` backed by :mod:`asyncio` primitives, +Python dicts and lists. Designed for **single-process** use — local +development, unit tests, and examples that want to avoid a Redis +dependency. + +.. note:: + + **Not suitable for production multiprocess deployments.** All state + lives inside the process; there is no persistence, no cross-process + pub/sub, and the "distributed" lock is just an :class:`asyncio.Lock`. + For real deployments use :class:`RedisMessageBus` (or another + networked backend). +""" +from __future__ import annotations + +import asyncio +import uuid +from collections import defaultdict +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from typing import Callable, Self + +from ._base import MessageBus + + +class InMemoryMessageBus(MessageBus): + """In-memory implementation of :class:`MessageBus`. + + Mapping of bus modes to in-memory structures: + + - **Mode A (drain queue)** — each key maps to a + :class:`list[tuple[str, dict]]` of ``(entry_id, payload)`` pairs. + ``queue_push`` appends; ``queue_drain`` pops from the front (FIFO) + and deletes the returned entries. + - **Mode C (replay log)** — same underlying list structure, but + ``log_read`` is non-destructive. ``log_trim`` removes entries + in-place. + - **Mode D (transient broadcast)** — each channel keeps a set of + :class:`asyncio.Queue` subscribers. ``publish`` pushes to all + currently-subscribed queues; ``subscribe`` yields from one. + - **Mode E (distributed lock)** — :class:`asyncio.Lock` per key, + suitable only for single-process concurrency. + - **Mode F (registry map)** — ``dict[str, dict[str, str]]``, one + nested dict per namespace. + + Entry ids are monotonic ``"-0"`` strings (e.g. ``"1-0"``, + ``"2-0"``, …). They are **not** lexicographically sortable once + the sequence exceeds single digits; comparison must parse the + numeric prefix (as :meth:`log_read` does internally). + """ + + def __init__(self) -> None: + """Initialise empty in-memory stores.""" + # Global auto-increment counter for entry ids. + self._seq: int = 0 + + # Mode A — drain queues: key -> [(entry_id, payload), ...] + self._queues: dict[str, list[tuple[str, dict]]] = defaultdict(list) + + # Mode C — replay logs: key -> [(entry_id, payload), ...] + self._logs: dict[str, list[tuple[str, dict]]] = defaultdict(list) + + # Mode D — pub/sub: channel -> set of asyncio.Queue subscribers + self._subscribers: dict[ + str, + set[asyncio.Queue[dict | None]], + ] = defaultdict(set) + + # Mode E — locks: key -> asyncio.Lock + self._locks: dict[str, asyncio.Lock] = defaultdict(asyncio.Lock) + # Track which key is currently held so is_locked() works. + self._lock_holders: dict[str, str] = {} + + # Mode F — registry maps: namespace -> {field: value} + self._registries: dict[str, dict[str, str]] = defaultdict(dict) + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def __aenter__(self) -> Self: + """No-op — nothing to open for in-memory transport. + + Returns: + `Self`: + The bus, ready for use. + """ + return self + + async def aclose(self) -> None: + """Signal all open subscribers so their generators terminate.""" + for subs in self._subscribers.values(): + for q in subs: + q.put_nowait(None) + self._subscribers.clear() + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _next_id(self) -> str: + """Return a monotonic ``"-0"`` id. + + The format matches Redis Stream entry ids so callers that parse + ids (e.g. ``_exclusive_start``) keep working. + + Returns: + `str`: + A unique, monotonically increasing entry id. + """ + self._seq += 1 + return f"{self._seq}-0" + + # ------------------------------------------------------------------ + # Mode A — drain queue + # ------------------------------------------------------------------ + + async def queue_push( + self, + key: str, + payload: dict, + *, + ttl_secs: int | None = None, + ) -> str: + """Append ``payload`` to the in-memory drain queue at ``key``. + + ``ttl_secs`` is accepted for API compatibility but ignored — the + in-memory implementation does not expire keys. + + Args: + key (`str`): + Queue identifier. + payload (`dict`): + JSON-serializable dict to enqueue. + ttl_secs (`int | None`, optional): + Ignored (no-op). + + Returns: + `str`: + The synthetic entry id assigned to this entry. + """ + entry_id = self._next_id() + self._queues[key].append((entry_id, payload)) + return entry_id + + async def queue_drain( + self, + key: str, + max_count: int = 100, + ) -> list[tuple[str, dict]]: + """Drain up to ``max_count`` entries from the queue at ``key``. + + Returned entries are removed from the internal list. + + Args: + key (`str`): + Queue identifier. + max_count (`int`, defaults to ``100``): + Maximum entries to return. + + Returns: + `list[tuple[str, dict]]`: + ``(entry_id, payload)`` pairs in arrival order. + """ + q = self._queues.get(key) + if not q: + return [] + drained = q[:max_count] + del q[:max_count] + return drained + + async def queue_delete(self, key: str) -> None: + """Delete the drain queue at ``key``. + + Args: + key (`str`): + Queue identifier. + """ + self._queues.pop(key, None) + + # ------------------------------------------------------------------ + # Mode C — replay log + # ------------------------------------------------------------------ + + async def log_append( + self, + key: str, + payload: dict, + *, + ttl_secs: int | None = None, + max_len: int | None = None, + ) -> str: + """Append ``payload`` to the replay log at ``key``. + + Args: + key (`str`): + Log identifier. + payload (`dict`): + JSON-serializable dict to append. + ttl_secs (`int | None`, optional): + Ignored (no-op — in-memory does not expire keys). + max_len (`int | None`, optional): + If set, trim the log to approximately this many entries + after appending. + + Returns: + `str`: + The entry id assigned to the new entry. + """ + entry_id = self._next_id() + log = self._logs[key] + log.append((entry_id, payload)) + if max_len is not None and len(log) > max_len: + del log[: len(log) - max_len] + return entry_id + + async def log_read( + self, + key: str, + since: str | None = None, + max_count: int = 100, + ) -> list[tuple[str, dict]]: + """Read up to ``max_count`` entries newer than ``since``. + + Reads are non-destructive. + + Args: + key (`str`): + Log identifier. + since (`str | None`, optional): + Exclusive cursor — return entries whose id is strictly + greater than ``since``. ``None`` reads from the start. + max_count (`int`, defaults to ``100``): + Maximum entries to return. + + Returns: + `list[tuple[str, dict]]`: + ``(entry_id, payload)`` pairs in append order. + """ + log = self._logs.get(key) + if not log: + return [] + if since is None: + return log[:max_count] + # Find the first entry with id > since. Entry ids are + # "-0" strings; integer comparison on the seq prefix is + # sufficient because our ids are monotonic. + since_seq = int(since.split("-")[0]) + start = 0 + for i, (eid, _) in enumerate(log): + if int(eid.split("-")[0]) > since_seq: + start = i + break + else: + # All entries are <= since. + return [] + return log[start : start + max_count] + + async def log_trim( + self, + key: str, + before_id: str | None = None, + ) -> None: + """Trim the replay log at ``key``. + + Args: + key (`str`): + Log identifier. + before_id (`str | None`, optional): + Drop all entries with id strictly older than this. + ``None`` drops the entire log. + """ + if before_id is None: + self._logs.pop(key, None) + return + log = self._logs.get(key) + if not log: + return + before_seq = int(before_id.split("-")[0]) + self._logs[key] = [ + (eid, p) for eid, p in log if int(eid.split("-")[0]) >= before_seq + ] + + # ------------------------------------------------------------------ + # Mode D — transient broadcast + # ------------------------------------------------------------------ + + async def publish( + self, + key: str, + payload: dict, + ) -> None: + """Publish ``payload`` to all current subscribers of ``key``. + + Args: + key (`str`): + Channel identifier. + payload (`dict`): + JSON-serializable dict delivered to each subscriber. + """ + for q in self._subscribers.get(key, set()): + q.put_nowait(payload) + + async def subscribe( + self, + key: str, + *, + on_ready: Callable[[], None] | None = None, + ) -> AsyncGenerator[dict, None]: + """Yield broadcast payloads for ``key`` until the consumer + closes the generator. + + Args: + key (`str`): + Channel identifier. + on_ready (`Callable[[], None] | None`, optional): + Called once after the subscription is established and + before any payload is yielded. + + Yields: + `dict`: + Each payload from :meth:`publish`. + """ + q: asyncio.Queue[dict | None] = asyncio.Queue() + self._subscribers[key].add(q) + try: + if on_ready is not None: + on_ready() + while True: + item = await q.get() + if item is None: + # Sentinel from aclose() — shut down gracefully. + break + yield item + finally: + self._subscribers[key].discard(q) + + # ------------------------------------------------------------------ + # Mode E — distributed lock (process-local asyncio.Lock) + # ------------------------------------------------------------------ + + @asynccontextmanager + async def acquire_lock( + self, + key: str, + *, + ttl_secs: int = 600, + ) -> AsyncGenerator[None, None]: + """Acquire a process-local mutex on ``key``. + + In-memory equivalent of a distributed lock. ``ttl_secs`` is + accepted for API compatibility but does **not** expire the lock + automatically — the lock is held until the context exits. + + Args: + key (`str`): + Lock identifier. + ttl_secs (`int`, defaults to ``600``): + Ignored (no automatic expiry). + + Yields: + `None`: while the lock is held. + """ + lock = self._locks[key] + token = uuid.uuid4().hex + async with lock: + self._lock_holders[key] = token + try: + yield + finally: + self._lock_holders.pop(key, None) + + async def is_locked(self, key: str) -> bool: + """Return whether ``key`` currently holds a lock. + + Args: + key (`str`): + Lock identifier. + + Returns: + `bool`: + ``True`` if some coroutine holds the lock. + """ + return key in self._lock_holders + + # ------------------------------------------------------------------ + # Mode F — registry map + # ------------------------------------------------------------------ + + async def registry_set( + self, + namespace: str, + field: str, + value: str, + *, + ttl_secs: int | None = None, + ) -> None: + """Set ``field`` to ``value`` in the registry at ``namespace``. + + ``ttl_secs`` is accepted for API compatibility but ignored. + + Args: + namespace (`str`): + Registry key. + field (`str`): + Field name. + value (`str`): + Value to store. + ttl_secs (`int | None`, optional): + Ignored (no TTL support). + """ + self._registries[namespace][field] = value + + async def registry_del(self, namespace: str, field: str) -> None: + """Remove ``field`` from the registry at ``namespace``. + + Args: + namespace (`str`): + Registry key. + field (`str`): + Field to remove. + """ + reg = self._registries.get(namespace) + if reg is not None: + reg.pop(field, None) + + async def registry_exists(self, namespace: str, field: str) -> bool: + """Return whether ``field`` exists in the registry at + ``namespace``. + + Args: + namespace (`str`): + Registry key. + field (`str`): + Field to check. + + Returns: + `bool`: + ``True`` if the field is present. + """ + return field in self._registries.get(namespace, {}) + + async def registry_getall( + self, + namespace: str, + ) -> dict[str, str]: + """Return all field-value pairs in the registry at + ``namespace``. + + Args: + namespace (`str`): + Registry key. + + Returns: + `dict[str, str]`: + All entries (shallow copy). Empty dict when absent. + """ + return dict(self._registries.get(namespace, {})) + + async def registry_drop(self, namespace: str) -> None: + """Delete the entire registry at ``namespace``. + + Args: + namespace (`str`): + Registry key to delete. + """ + self._registries.pop(namespace, None) diff --git a/src/agentscope/app/message_bus/_keys.py b/src/agentscope/app/message_bus/_keys.py new file mode 100644 index 0000000000000000000000000000000000000000..8767c5fd56252146e6c5f5b35c6d63cb8146de70 --- /dev/null +++ b/src/agentscope/app/message_bus/_keys.py @@ -0,0 +1,228 @@ +# -*- coding: utf-8 -*- +"""Centralised registry of message-bus key/namespace conventions used +by application-layer services. + +:class:`~agentscope.app.message_bus.MessageBus` itself stays +domain-agnostic — it exposes only generic primitives +(``publish`` / ``subscribe`` / ``queue_*`` / ``log_*`` / ``registry_*`` +/ ``acquire_lock``). All business-specific key formats live here so +they can be audited, migrated, and (eventually) ported off from the +current scattered ``_BASE_…_KEY`` constants on ``MessageBus``. + +Add new business keys here as needed. As legacy keys are migrated off +``MessageBus``, they should move into this class as well. +""" + +from typing import Final + + +class MessageBusKeys: + """Application-layer key conventions for the message bus.""" + + # ------------------------------------------------------------------ + # Run-trigger queue — the discriminator carried by each entry on the + # shared trigger queue, telling the dispatcher how to spawn the run. + # Centralised here (rather than on ``MessageBus``) so the bus stays + # free of business vocabulary. + # ------------------------------------------------------------------ + + WAKEUP_KIND_WAKE: Final = "wake" + """Trigger kind: wake an *idle* session to drain pending inbox + content. The dispatcher spawns the run with ``input_msg=None`` and + skips the session entirely while it is already running.""" + + WAKEUP_KIND_RESUME: Final = "resume" + """Trigger kind: resume a session parked on an awaiting tool call by + feeding it a human-in-the-loop result. The dispatcher spawns the run + with the carried ``input`` event and — unlike ``wake`` — must *not* + drop the entry while the session is running; it re-queues until the + parked run releases its lock.""" + + # ------------------------------------------------------------------ + # Cross-session UI projection — a generic per-session Redis-hash + # store onto which one session can project UI cards owned by another + # (e.g. a team member's pending HITL request projected onto its + # leader). The ``kind`` prefix on each field lets a single target + # session carry several independent projection feeds without key + # collisions, so new projection features reuse this store rather + # than minting their own. + # ------------------------------------------------------------------ + + _PROJECTION_NS = "agentscope:session:projection:{sid}" + """Redis-hash namespace key template (per *target* session id).""" + + @classmethod + def projection_namespace(cls, target_session_id: str) -> str: + """Return the registry namespace for a session's projections. + + Args: + target_session_id (`str`): + The session the entries are projected onto (the session + whose UI renders them). + + Returns: + `str`: + The Redis-hash namespace key. + """ + return cls._PROJECTION_NS.format(sid=target_session_id) + + @staticmethod + def projection_field(kind: str, entry_id: str) -> str: + """Return the hash field key for a single projected entry. + + The ``kind`` prefix namespaces the field so different projection + feeds sharing one target session never collide, and so a feed + can be listed/purged by scanning for its prefix. + + Args: + kind (`str`): + The projection feed this entry belongs to (e.g. + ``"subagent_hitl"``). + entry_id (`str`): + The entry's identity within the feed, unique per + ``kind`` (e.g. ``"{worker_session_id}:{reply_id}"``). + + Returns: + `str`: + The hash field key, ``"{kind}:{entry_id}"``. + """ + return f"{kind}:{entry_id}" + + @staticmethod + def projection_field_prefix(kind: str) -> str: + """Return the field-key prefix that identifies one feed. + + Used to filter :meth:`MessageBus.registry_getall` down to a + single ``kind`` when listing or purging. + + Args: + kind (`str`): + The projection feed. + + Returns: + `str`: + The field-key prefix, ``"{kind}:"``. + """ + return f"{kind}:" + + # ------------------------------------------------------------------ + # Session event stream (replay log + live pub/sub) + # ------------------------------------------------------------------ + + _SESSION_EVENTS = "agentscope:session:events:{sid}" + + SESSION_REPLAY_MAX_LEN = 1000 + """Replay log length cap; older events are trimmed on append.""" + + @classmethod + def session_events(cls, session_id: str) -> str: + """Replay log + live pub/sub channel key for a session.""" + return cls._SESSION_EVENTS.format(sid=session_id) + + # ------------------------------------------------------------------ + # Session run lock + # ------------------------------------------------------------------ + + _SESSION_LOCK = "agentscope:session:lock:{sid}" + + SESSION_RUN_TTL_SECS = 600 + """Default lock lease for a chat run (10 minutes).""" + + @classmethod + def session_lock(cls, session_id: str) -> str: + """Per-session distributed-lock key.""" + return cls._SESSION_LOCK.format(sid=session_id) + + # ------------------------------------------------------------------ + # Session inbox + # ------------------------------------------------------------------ + + _INBOX = "agentscope:inbox:{sid}" + + @classmethod + def inbox(cls, session_id: str) -> str: + """Per-session inbox drain-queue key.""" + return cls._INBOX.format(sid=session_id) + + # ------------------------------------------------------------------ + # Run trigger queue (wakeup / resume) + # ------------------------------------------------------------------ + + _WAKEUP_QUEUE = "agentscope:wakeups" + _WAKEUP_SIGNAL = "agentscope:wakeup_signal" + + @classmethod + def wakeup_queue(cls) -> str: + """Shared run-trigger queue key.""" + return cls._WAKEUP_QUEUE + + @classmethod + def wakeup_signal(cls) -> str: + """Shared signal channel that nudges dispatchers to drain.""" + return cls._WAKEUP_SIGNAL + + # ------------------------------------------------------------------ + # Cross-process cancel + # ------------------------------------------------------------------ + + _SESSION_CANCEL = "agentscope:session:cancel" + _TASK_CANCEL = "agentscope:task:cancel" + + @classmethod + def session_cancel_channel(cls) -> str: + """Global session-cancel broadcast channel.""" + return cls._SESSION_CANCEL + + @classmethod + def task_cancel_channel(cls) -> str: + """Single-task cancel broadcast channel.""" + return cls._TASK_CANCEL + + # ------------------------------------------------------------------ + # Background task registry + # ------------------------------------------------------------------ + + _BG_TASKS = "agentscope:bg_tasks:{sid}" + + BG_TASKS_TTL_SECS = 86400 + """Fallback TTL for the per-session BG task registry (24 h).""" + + @classmethod + def bg_tasks(cls, session_id: str) -> str: + """Per-session background task registry key.""" + return cls._BG_TASKS.format(sid=session_id) + + # ------------------------------------------------------------------ + # Knowledge-base indexing pipeline + # ------------------------------------------------------------------ + + _INDEX_TASKS_QUEUE = "agentscope:index:tasks" + _INDEX_TASKS_SIGNAL = "agentscope:index:tasks:wake" + + @classmethod + def index_tasks_queue(cls) -> str: + """Shared, durable index-task queue. + + The producer-side + :class:`~agentscope.app._service.KnowledgeBaseService` + ``queue_push``\\ es here through + :func:`~agentscope.app._bus_ops.enqueue_index_task`; the + consumer-side + :class:`~agentscope.app._service.IndexTaskConsumer` + ``queue_drain``\\ s it on each signal — plus once eagerly on + consumer start-up so tasks queued while every worker was down + are picked up immediately. + """ + return cls._INDEX_TASKS_QUEUE + + @classmethod + def index_tasks_signal(cls) -> str: + """Shared pub/sub channel that nudges every running index-task + consumer to drain the queue. + + Payload is opaque — only its arrival matters. Redis pub/sub is + fire-and-forget, so a published signal does not guarantee + delivery; the durable queue keeps work safe in case every + subscriber happens to be offline. + """ + return cls._INDEX_TASKS_SIGNAL diff --git a/src/agentscope/app/message_bus/_redis_message_bus.py b/src/agentscope/app/message_bus/_redis_message_bus.py new file mode 100644 index 0000000000000000000000000000000000000000..f15abb9697038943c64290f251a561f1ae8902dd --- /dev/null +++ b/src/agentscope/app/message_bus/_redis_message_bus.py @@ -0,0 +1,638 @@ +# -*- coding: utf-8 -*- +"""The Redis-backed message bus implementation.""" +import asyncio +import json +import uuid +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from typing import Any, Callable, Self, TYPE_CHECKING + +from ._base import MessageBus + +if TYPE_CHECKING: + from redis.asyncio import ConnectionPool, Redis +else: + ConnectionPool = Any + Redis = Any + + +class RedisMessageBus(MessageBus): + """Redis-backed implementation of :class:`MessageBus`. + + Mapping of bus modes to Redis primitives: + + - **Mode A (drain queue)** uses a Redis Stream per key. ``XADD`` + appends a payload whose single field ``payload`` carries the + JSON-serialised dict. ``queue_drain`` performs ``XRANGE`` followed + by per-id ``XDEL`` so the read is destructive and idempotent + under the single-consumer-per-key invariant. ``ttl_secs`` is + enforced via ``EXPIRE`` after each push (sliding TTL). + - **Mode C (replay log)** also uses a Redis Stream, but never + ``XDEL``s on read. Trimming happens via ``XADD … MAXLEN ~N`` + (approximate, for performance) on append, via the ``ttl_secs`` + ``EXPIRE`` on the whole key, or explicitly via + :meth:`log_trim`. ``log_read`` uses ``XRANGE`` with an exclusive + start id derived from ``since``. + - **Mode D (transient broadcast)** rides Redis Pub/Sub. Wake-ups + are best-effort: payloads published before a subscription exists + are not delivered. + + The bus owns its own connection pool by default; an external pool + may be supplied for tests or for sharing a pool across services. + """ + + def __init__( + self, + host: str = "localhost", + port: int = 6379, + db: int = 0, + password: str | None = None, + connection_pool: ConnectionPool | None = None, + **kwargs: Any, + ) -> None: + """Store connection parameters; the actual pool is created in + :meth:`__aenter__`. + + Args: + host (`str`, defaults to ``"localhost"``): + Redis server host. + port (`int`, defaults to ``6379``): + Redis server port. + db (`int`, defaults to ``0``): + Redis logical database index. + password (`str | None`, optional): + Redis password if required by the server. + connection_pool (`ConnectionPool | None`, optional): + An externally managed connection pool. When provided + the pool is used as-is and **not** closed by + :meth:`aclose` — the caller retains ownership of its + lifecycle. When omitted a pool is created from + *host*/*port*/*db*/*password* on :meth:`__aenter__` + and closed on :meth:`aclose`. + **kwargs (`Any`): + Extra keyword arguments forwarded to + ``redis.asyncio.ConnectionPool`` when the pool is + created internally (e.g. ``max_connections=20``, + ``socket_timeout=5``). + """ + self._host = host + self._port = port + self._db = db + self._password = password + self._external_pool: ConnectionPool | None = connection_pool + self._kwargs = kwargs + + # Populated in __aenter__; None until the context is entered. + self._client: Redis | None = None + self._owned_pool: ConnectionPool | None = None + + async def __aenter__(self) -> Self: + """Create the connection pool and Redis client. + + If an external pool was supplied at construction time it is + used directly and its lifecycle remains the caller's + responsibility. Otherwise, an internal pool is created from the + stored host/port/db parameters and will be closed by + :meth:`aclose`. + + Returns: + `Self`: + The bus, ready for use as an async context manager. + """ + try: + import redis.asyncio as aioredis + except ImportError as e: + raise ImportError( + "The 'redis' package is required for RedisMessageBus. " + "Install it with: pip install redis[async]", + ) from e + + if self._external_pool is not None: + pool = self._external_pool + else: + self._owned_pool = aioredis.ConnectionPool( + host=self._host, + port=self._port, + db=self._db, + password=self._password, + decode_responses=True, + **self._kwargs, + ) + pool = self._owned_pool + + self._client = aioredis.Redis(connection_pool=pool) + return self + + async def aclose(self) -> None: + """Close the connection pool if it was created internally. + + Externally supplied pools are left open — the caller owns them. + """ + if self._owned_pool is not None: + await self._owned_pool.aclose() + self._owned_pool = None + self._client = None + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: Any, + ) -> None: + """Exit the async context manager and release resources. + + Args: + exc_type (`type[BaseException] | None`): + Exception class raised inside the context, if any. + exc_value (`BaseException | None`): + Exception instance raised inside the context, if any. + traceback (`Any`): + Traceback associated with the exception, if any. + """ + await self.aclose() + + def get_client(self) -> Redis: + """Return the underlying Redis client. + + Only valid inside the async context (between :meth:`__aenter__` + and :meth:`aclose`). + + Returns: + `Redis`: + The asyncio Redis client instance. + """ + return self._client + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + @staticmethod + def _exclusive_start(since: str | None) -> str: + """Compute the exclusive start id for ``XRANGE`` given a cursor. + + Redis Streams 6.2+ support the ``(`` prefix for exclusive + ranges, so the simplest portable encoding of "after ``since``" + is ``f"({since}"``. ``None`` resolves to ``"-"`` (the absolute + minimum id). + + Args: + since (`str | None`): + The previous cursor returned by :meth:`log_read`, or + ``None`` to read from the beginning. + + Returns: + `str`: + The start argument for ``XRANGE``. + """ + if since is None: + return "-" + return f"({since}" + + # ------------------------------------------------------------------ + # Mode A — drain queue + # ------------------------------------------------------------------ + + async def queue_push( + self, + key: str, + payload: dict, + *, + ttl_secs: int | None = None, + ) -> str: + """Append ``payload`` to the drain queue at ``key``. + + Args: + key (`str`): + Stream key for this drain queue. + payload (`dict`): + JSON-serializable dict; encoded into the ``payload`` + field of a Redis Stream entry. + ttl_secs (`int | None`, optional): + If set, refresh the key's expiry on every push + (sliding TTL). ``None`` means no TTL — the queue + persists until drained or deleted explicitly. + + Returns: + `str`: + The Redis Stream entry id assigned by ``XADD``. + """ + entry_id = await self._client.xadd( + key, + {"payload": json.dumps(payload)}, + ) + if ttl_secs is not None: + await self._client.expire(key, ttl_secs) + return entry_id + + async def queue_drain( + self, + key: str, + max_count: int = 100, + ) -> list[tuple[str, dict]]: + """Drain up to ``max_count`` entries from the queue at ``key``. + + Implementation: ``XRANGE`` followed by ``XDEL`` of the + returned ids, so the operation is destructive in a single + round-trip pair. Entries that arrive between ``XRANGE`` and + ``XDEL`` are not affected. + + Args: + key (`str`): + Stream key for the drain queue. + max_count (`int`, defaults to ``100``): + Maximum entries to return in one call. + + Returns: + `list[tuple[str, dict]]`: + ``(entry_id, payload)`` pairs in arrival order. Empty + list when the queue is empty or absent. + """ + entries = await self._client.xrange(key, count=max_count) + if not entries: + return [] + + results: list[tuple[str, dict]] = [] + ids_to_delete: list[str] = [] + for entry_id, fields in entries: + ids_to_delete.append(entry_id) + raw = fields.get("payload") + if raw is None: + continue + results.append((entry_id, json.loads(raw))) + + if ids_to_delete: + await self._client.xdel(key, *ids_to_delete) + + return results + + async def queue_delete(self, key: str) -> None: + """Delete the drain queue at ``key``. + + Args: + key (`str`): + Stream key for the drain queue. ``DEL`` is a no-op + when the key does not exist. + """ + await self._client.delete(key) + + # ------------------------------------------------------------------ + # Mode C — replay log + # ------------------------------------------------------------------ + + async def log_append( + self, + key: str, + payload: dict, + *, + ttl_secs: int | None = None, + max_len: int | None = None, + ) -> str: + """Append ``payload`` to the replay log at ``key``. + + Args: + key (`str`): + Stream key for this replay log. + payload (`dict`): + JSON-serializable dict; encoded into the ``payload`` + field of a Redis Stream entry. + ttl_secs (`int | None`, optional): + If set, refresh the key's expiry on every append + (sliding TTL). ``None`` means no TTL. + max_len (`int | None`, optional): + If set, cap the log at approximately this many entries + using ``XADD MAXLEN ~N``. The cap is approximate so + Redis can use its O(1) trim-by-radix-tree-node mode. + ``None`` means no cap. + + Returns: + `str`: + The Redis Stream entry id, suitable as a cursor for + later :meth:`log_read` calls. + """ + kwargs: dict[str, Any] = {} + if max_len is not None: + kwargs["maxlen"] = max_len + kwargs["approximate"] = True + entry_id = await self._client.xadd( + key, + {"payload": json.dumps(payload)}, + **kwargs, + ) + if ttl_secs is not None: + await self._client.expire(key, ttl_secs) + return entry_id + + async def log_read( + self, + key: str, + since: str | None = None, + max_count: int = 100, + ) -> list[tuple[str, dict]]: + """Read up to ``max_count`` entries newer than ``since``. + + Args: + key (`str`): + Stream key for the replay log. + since (`str | None`, optional): + Exclusive cursor — return entries strictly newer than + this id (typically the last id from a previous read). + ``None`` reads from the beginning. + max_count (`int`, defaults to ``100``): + Maximum entries to return. + + Returns: + `list[tuple[str, dict]]`: + ``(entry_id, payload)`` pairs in append order. + """ + start = self._exclusive_start(since) + entries = await self._client.xrange( + key, + min=start, + count=max_count, + ) + results: list[tuple[str, dict]] = [] + for entry_id, fields in entries: + raw = fields.get("payload") + if raw is None: + continue + results.append((entry_id, json.loads(raw))) + return results + + async def log_trim( + self, + key: str, + before_id: str | None = None, + ) -> None: + """Trim the replay log at ``key``. + + Args: + key (`str`): + Stream key for the replay log. + before_id (`str | None`, optional): + Drop all entries with id strictly older than this. + ``None`` deletes the entire key. + """ + if before_id is None: + await self._client.delete(key) + return + # XTRIM MINID drops entries with id < before_id. + await self._client.xtrim(key, minid=before_id) + + # ------------------------------------------------------------------ + # Mode F — registry map (hash-keyed namespace) + # ------------------------------------------------------------------ + + async def registry_set( + self, + namespace: str, + field: str, + value: str, + *, + ttl_secs: int | None = None, + ) -> None: + """Set ``field`` in the Redis Hash at ``namespace``. + + Args: + namespace (`str`): + Hash key. + field (`str`): + Hash field. + value (`str`): + Value to store. + ttl_secs (`int | None`, optional): + Refresh the key's expiry (sliding TTL). + """ + await self._client.hset(namespace, field, value) + if ttl_secs is not None: + await self._client.expire(namespace, ttl_secs) + + async def registry_del(self, namespace: str, field: str) -> None: + """Remove ``field`` from the Redis Hash at ``namespace``. + + Args: + namespace (`str`): + Hash key. + field (`str`): + Hash field to remove. + """ + await self._client.hdel(namespace, field) + + async def registry_exists(self, namespace: str, field: str) -> bool: + """Return whether ``field`` exists in the Hash at ``namespace``. + + Args: + namespace (`str`): + Hash key. + field (`str`): + Hash field to check. + + Returns: + `bool`: + ``True`` if the field exists. + """ + return bool(await self._client.hexists(namespace, field)) + + async def registry_getall( + self, + namespace: str, + ) -> dict[str, str]: + """Return all field-value pairs from the Hash at ``namespace``. + + Args: + namespace (`str`): + Hash key. + + Returns: + `dict[str, str]`: + All entries. Empty dict when the key is absent. + """ + return await self._client.hgetall(namespace) or {} + + async def registry_drop(self, namespace: str) -> None: + """Delete the entire Hash at ``namespace``. + + Args: + namespace (`str`): + Hash key to delete. + """ + await self._client.delete(namespace) + + # ------------------------------------------------------------------ + # Mode D — transient broadcast + # ------------------------------------------------------------------ + + # Poll interval for the pub/sub read loop. Bounding each read keeps + # long-lived idle subscriptions resilient: an idle ``socket_timeout`` + # read (raised when the connection defines one, or when the server + # drops idle connections) surfaces as a benign per-poll timeout that + # we ignore, instead of a fatal error that tears down the generator. + _SUBSCRIBE_POLL_TIMEOUT_SECS = 1.0 + + async def publish( + self, + key: str, + payload: dict, + ) -> None: + """Publish ``payload`` on the broadcast channel ``key``. + + Args: + key (`str`): + Pub/Sub channel name. + payload (`dict`): + JSON-serializable dict; encoded as the channel + message body. + """ + await self._client.publish(key, json.dumps(payload)) + + async def subscribe( + self, + key: str, + *, + on_ready: Callable[[], None] | None = None, + ) -> AsyncGenerator[dict, None]: + """Yield broadcast payloads on ``key`` until the consumer + closes the generator. + + Args: + key (`str`): + Pub/Sub channel name. + on_ready (`Callable[[], None] | None`, optional): + Called once after ``SUBSCRIBE`` has been issued (and + before any payload is yielded). Lets callers block on + a "subscription is live" event so they can publish + immediately after starting a subscriber without losing + the first message to a SUBSCRIBE/PUBLISH race. + + Yields: + `dict`: + Each payload originally passed to :meth:`publish`. + """ + from redis import exceptions as redis_exceptions + + pubsub = self._client.pubsub() + try: + await pubsub.subscribe(key) + if on_ready is not None: + on_ready() + while True: + try: + message = await pubsub.get_message( + ignore_subscribe_messages=True, + timeout=self._SUBSCRIBE_POLL_TIMEOUT_SECS, + ) + except redis_exceptions.TimeoutError: + # Idle read timeout (e.g. the connection defines a + # ``socket_timeout`` or the server drops idle + # connections). No message arrived in this window; + # keep listening rather than crashing the loop. + continue + if message is None: + # No payload within the poll window — keep waiting. + continue + if message.get("type") != "message": + # Skip ``subscribe`` ack and similar control frames. + continue + data = message.get("data") + if data is None: + continue + yield json.loads(data) + finally: + await pubsub.unsubscribe(key) + await pubsub.aclose() + + # ------------------------------------------------------------------ + # Mode E — distributed lock + # ------------------------------------------------------------------ + + # Poll interval while waiting for a contested lock. + _LOCK_RETRY_DELAY_SECS = 0.1 + + @asynccontextmanager + async def acquire_lock( + self, + key: str, + *, + ttl_secs: int = 600, + ) -> AsyncGenerator[None, None]: + """Acquire ``key`` as a distributed mutex. + + Implementation: + + - ``SET key NX EX ttl_secs`` to claim the + lock atomically. Retries every + ``_LOCK_RETRY_DELAY_SECS`` until acquired. + - A heartbeat task renews the TTL every ``ttl_secs / 2`` + seconds while the body runs, so a long-running holder + does not lose the lease. + - On exit, the heartbeat is cancelled and the lock is + released by GET-then-DEL guarded on the random token — + we never delete a key whose value isn't ours, so a + process whose lease has already expired (and been re- + acquired by someone else) cannot accidentally release + the new holder's lock. The non-atomic GET+DEL race + window opens only after the heartbeat is cancelled, in + a sub-millisecond span; the ``ttl_secs`` lease still has + plenty of time left so the race effectively never + materialises. + + Args: + key (`str`): + Lock identifier. + ttl_secs (`int`, defaults to ``600``): + Lease duration; auto-renewed via heartbeat. + + Yields: + `None`: while the lock is held. + """ + token = uuid.uuid4().hex + # Acquire (poll until success). + while True: + ok = await self._client.set(key, token, nx=True, ex=ttl_secs) + if ok: + break + await asyncio.sleep(self._LOCK_RETRY_DELAY_SECS) + + # Heartbeat: renew TTL every ttl/2 seconds. + async def _heartbeat() -> None: + while True: + await asyncio.sleep(max(1.0, ttl_secs / 2)) + await self._client.expire(key, ttl_secs) + + hb_task = asyncio.create_task( + _heartbeat(), + name=f"lock-heartbeat:{key}", + ) + + try: + yield + finally: + hb_task.cancel() + try: + await hb_task + except asyncio.CancelledError: + pass + # Release: only delete if the value still matches our + # token. Failures are swallowed (the lease may have + # expired naturally, in which case there is nothing to + # do). + try: + current = await self._client.get(key) + if current == token: + await self._client.delete(key) + except Exception: # pylint: disable=broad-except + pass + + async def is_locked(self, key: str) -> bool: + """Return whether ``key`` currently holds a lock. + + Args: + key (`str`): + Lock identifier. + + Returns: + `bool`: + ``True`` if Redis has a value at this key. + """ + result = await self._client.exists(key) + return bool(result) diff --git a/src/agentscope/app/middleware/__init__.py b/src/agentscope/app/middleware/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1a622193d975e541ae95766ec211fb98f516c880 --- /dev/null +++ b/src/agentscope/app/middleware/__init__.py @@ -0,0 +1,16 @@ +# -*- coding: utf-8 -*- +"""The middlewares module.""" + +from ._inbox_middleware import InboxMiddleware +from ._protocol import ProtocolMiddlewareBase, AGUIProtocolMiddleware +from ._state_change_middleware import StateChangeMiddleware +from ._tool_offload_middleware import ToolOffloadMiddleware + + +__all__ = [ + "InboxMiddleware", + "ProtocolMiddlewareBase", + "AGUIProtocolMiddleware", + "StateChangeMiddleware", + "ToolOffloadMiddleware", +] diff --git a/src/agentscope/app/middleware/_inbox_middleware.py b/src/agentscope/app/middleware/_inbox_middleware.py new file mode 100644 index 0000000000000000000000000000000000000000..b5db66f28529dcdb99263f3a185cd72685fb4aaf --- /dev/null +++ b/src/agentscope/app/middleware/_inbox_middleware.py @@ -0,0 +1,136 @@ +# -*- coding: utf-8 -*- +"""Generic middleware that drains the message bus inbox before reasoning. + +Producers push :class:`~agentscope.message.HintBlock` payloads into +the per-session inbox via :class:`~agentscope.app._message_bus.MessageBus`. +This middleware drains the inbox at the start of each reasoning step +and injects the HintBlocks into ``agent.state.context`` — appended to +the last assistant message's content list (same pattern as +:class:`ToolOffloadMiddleware`). + +Each injected HintBlock also yields a one-shot ``HintBlockEvent`` +so the front-end SSE stream can render it in real time. +""" +from typing import Any, AsyncGenerator, Callable + +from ..message_bus import MessageBus, MessageBusKeys +from ..._logging import logger +from ...agent import Agent +from ...event import HintBlockEvent +from ...message import AssistantMsg, HintBlock +from ...middleware import MiddlewareBase + + +class InboxMiddleware(MiddlewareBase): # pylint: disable=abstract-method + """Drain the session's inbox and inject HintBlocks before each + reasoning step. + + Each entry in the inbox is a serialised + :class:`~agentscope.message.HintBlock`. The middleware + deserializes them, appends to the last assistant message in + ``agent.state.context``, and yields a one-shot ``HintBlockEvent`` + for each so the front-end sees them. + + Args: + message_bus (`MessageBus`): + The application message bus to read from. + max_count (`int`, defaults to ``100``): + Maximum number of entries drained per reasoning step. + """ + + def __init__( + self, + message_bus: MessageBus, + max_count: int = 100, + ) -> None: + """Initialise the middleware. + + Args: + message_bus (`MessageBus`): + The application-level message bus. + max_count (`int`, defaults to ``100``): + Maximum entries drained per reasoning step. + """ + self._bus = message_bus + self._max_count = max_count + + async def on_reasoning( # type: ignore[override] + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator[Any, None]: + """Drain the inbox, inject HintBlocks into context, yield + events, then continue with downstream reasoning. + + Args: + agent (`Agent`): + The executing agent. ``agent.state.session_id`` selects + the inbox to drain. + input_kwargs (`dict`): + Reasoning input kwargs (contains ``tool_choice``); + forwarded unchanged to ``next_handler``. + next_handler (`Callable[..., AsyncGenerator]`): + The downstream middleware or core reasoning logic. + + Yields: + `Any`: + One ``HintBlockEvent`` per drained inbox entry, + followed by events from downstream. + """ + entries = await self._bus.queue_drain( + MessageBusKeys.inbox(agent.state.session_id), + max_count=self._max_count, + ) + + if entries: + hint_blocks = [ + HintBlock.model_validate(payload) + for _entry_id, payload in entries + ] + + logger.info( + "InboxMiddleware: injecting %d HintBlock(s) into context " + "for session %s", + len(hint_blocks), + agent.state.session_id, + ) + + # Inject into agent context (same pattern as + # ToolOffloadMiddleware). + if len(agent.state.context) > 0: + last_msg = agent.state.context[-1] + if ( + last_msg.role == "assistant" + and last_msg.name == agent.name + ): + last_msg.content.extend(hint_blocks) + else: + agent.state.context.append( + AssistantMsg( + id=agent.state.reply_id, + name=agent.name, + content=list(hint_blocks), + ), + ) + else: + agent.state.context.append( + AssistantMsg( + id=agent.state.reply_id, + name=agent.name, + content=list(hint_blocks), + ), + ) + + # Yield one-shot events so the front-end SSE stream sees + # each HintBlock. + for hint in hint_blocks: + yield HintBlockEvent( + reply_id=agent.state.reply_id, + block_id=hint.id, + source=hint.source, + hint=hint.hint, + ) + + async for evt in next_handler(**input_kwargs): + yield evt diff --git a/src/agentscope/app/middleware/_protocol/__init__.py b/src/agentscope/app/middleware/_protocol/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e6e0180ebd27ac727bafb6decffdcd74e6b72074 --- /dev/null +++ b/src/agentscope/app/middleware/_protocol/__init__.py @@ -0,0 +1,10 @@ +# -*- coding: utf-8 -*- +"""The middleware used for agent protocol.""" + +from ._base import ProtocolMiddlewareBase +from ._agui import AGUIProtocolMiddleware + +__all__ = [ + "ProtocolMiddlewareBase", + "AGUIProtocolMiddleware", +] diff --git a/src/agentscope/app/middleware/_protocol/_agui.py b/src/agentscope/app/middleware/_protocol/_agui.py new file mode 100644 index 0000000000000000000000000000000000000000..e2de09cad1e091d1808fc9db654bd1199a3cd621 --- /dev/null +++ b/src/agentscope/app/middleware/_protocol/_agui.py @@ -0,0 +1,257 @@ +# -*- coding: utf-8 -*- +"""The AGUI middleware class.""" +from typing import TYPE_CHECKING, Any + +from starlette.types import ASGIApp + +from ._base import ProtocolMiddlewareBase +from ....event import ( + AgentEvent, + DataBlockDeltaEvent, + DataBlockEndEvent, + DataBlockStartEvent, + ExceedMaxItersEvent, + ExternalExecutionResultEvent, + ModelCallEndEvent, + ModelCallStartEvent, + ReplyEndEvent, + ReplyStartEvent, + RequireExternalExecutionEvent, + RequireUserConfirmEvent, + TextBlockDeltaEvent, + TextBlockEndEvent, + TextBlockStartEvent, + ThinkingBlockDeltaEvent, + ThinkingBlockEndEvent, + ThinkingBlockStartEvent, + ToolCallDeltaEvent, + ToolCallEndEvent, + ToolCallStartEvent, + ToolResultDataDeltaEvent, + ToolResultEndEvent, + ToolResultStartEvent, + ToolResultTextDeltaEvent, + UserConfirmResultEvent, +) + +if TYPE_CHECKING: + from ag_ui.core.events import BaseEvent as AGUIBaseEvent +else: + AGUIBaseEvent = Any + + +class AGUIProtocolMiddleware(ProtocolMiddlewareBase): + """The middleware that converts the AgentScope events into AGUI + protocol.""" + + def __init__(self, app: ASGIApp) -> None: + """Initialize the AGUI protocol middleware. + + Args: + app: The ASGI application to wrap. + """ + super().__init__(app) + # Per-instance state; safe under typical single-stream usage + # but not across concurrent requests. Use contextvars if + # concurrency is needed. + self._last_model_name: str = "model_call" + self._tool_result_buffers: dict[str, list[str]] = {} + + def _convert_to_protocol(self, event: AgentEvent) -> dict: + """Convert the AgentScope events into AGUI protocol.""" + agui_event = self._to_agui_event(event) + return agui_event.model_dump( + mode="json", + exclude_none=True, + by_alias=True, + ) + + # pylint: disable=too-many-return-statements + def _to_agui_event( # noqa: C901 + self, + event: AgentEvent, + ) -> "AGUIBaseEvent": + """Convert an AgentScope event to an AGUI event.""" + + from ag_ui.core.events import ( + CustomEvent as AGUICustomEvent, + ReasoningMessageContentEvent as AGUIReasoningMessageContentEvent, + ReasoningMessageEndEvent as AGUIReasoningMessageEndEvent, + ReasoningMessageStartEvent as AGUIReasoningMessageStartEvent, + RunErrorEvent as AGUIRunErrorEvent, + RunFinishedEvent as AGUIRunFinishedEvent, + RunStartedEvent as AGUIRunStartedEvent, + StepFinishedEvent as AGUIStepFinishedEvent, + StepStartedEvent as AGUIStepStartedEvent, + TextMessageContentEvent as AGUITextMessageContentEvent, + TextMessageEndEvent as AGUITextMessageEndEvent, + TextMessageStartEvent as AGUITextMessageStartEvent, + ToolCallArgsEvent as AGUIToolCallArgsEvent, + ToolCallEndEvent as AGUIToolCallEndEvent, + ToolCallResultEvent as AGUIToolCallResultEvent, + ToolCallStartEvent as AGUIToolCallStartEvent, + ) + + if isinstance(event, ReplyStartEvent): + return AGUIRunStartedEvent( + thread_id=event.session_id, + run_id=event.reply_id, + ) + + if isinstance(event, ReplyEndEvent): + return AGUIRunFinishedEvent( + thread_id=event.session_id, + run_id=event.reply_id, + ) + + if isinstance(event, ExceedMaxItersEvent): + return AGUIRunErrorEvent( + message=(f"Agent '{event.name}' exceeded max iterations"), + code="exceed_max_iters", + ) + + if isinstance(event, ModelCallStartEvent): + self._last_model_name = event.model_name + return AGUIStepStartedEvent( + step_name=event.model_name, + ) + + if isinstance(event, ModelCallEndEvent): + return AGUIStepFinishedEvent( + step_name=self._last_model_name, + ) + + if isinstance(event, TextBlockStartEvent): + return AGUITextMessageStartEvent( + message_id=event.block_id, + ) + + if isinstance(event, TextBlockDeltaEvent): + return AGUITextMessageContentEvent( + message_id=event.block_id, + delta=event.delta, + ) + + if isinstance(event, TextBlockEndEvent): + return AGUITextMessageEndEvent( + message_id=event.block_id, + ) + + # AGUI has a two-level reasoning structure (ReasoningStart/End wrapping + # ReasoningMessage*), but _convert_to_protocol returns a single dict + # per input event, so only the inner ReasoningMessage* events are + # emitted. Most AGUI consumers render correctly with message-level + # events alone. + if isinstance(event, ThinkingBlockStartEvent): + return AGUIReasoningMessageStartEvent( + message_id=event.block_id, + role="reasoning", + ) + + if isinstance(event, ThinkingBlockDeltaEvent): + return AGUIReasoningMessageContentEvent( + message_id=event.block_id, + delta=event.delta, + ) + + if isinstance(event, ThinkingBlockEndEvent): + return AGUIReasoningMessageEndEvent( + message_id=event.block_id, + ) + + if isinstance(event, ToolCallStartEvent): + return AGUIToolCallStartEvent( + tool_call_id=event.tool_call_id, + tool_call_name=event.tool_call_name, + parent_message_id=event.reply_id, + ) + + if isinstance(event, ToolCallDeltaEvent): + return AGUIToolCallArgsEvent( + tool_call_id=event.tool_call_id, + delta=event.delta, + ) + + if isinstance(event, ToolCallEndEvent): + return AGUIToolCallEndEvent( + tool_call_id=event.tool_call_id, + ) + + if isinstance(event, ToolResultStartEvent): + return AGUICustomEvent( + name="tool_result_start", + value=event.model_dump(exclude_none=True), + ) + + if isinstance(event, ToolResultTextDeltaEvent): + self._tool_result_buffers.setdefault( + event.tool_call_id, + [], + ).append(event.delta) + return AGUICustomEvent( + name="tool_result_text_delta", + value=event.model_dump(exclude_none=True), + ) + + if isinstance(event, ToolResultDataDeltaEvent): + return AGUICustomEvent( + name="tool_result_data_delta", + value=event.model_dump(exclude_none=True), + ) + + if isinstance(event, ToolResultEndEvent): + content = "".join( + self._tool_result_buffers.pop(event.tool_call_id, []), + ) + return AGUIToolCallResultEvent( + tool_call_id=event.tool_call_id, + message_id=event.reply_id, + content=content or str(event.state), + ) + + if isinstance(event, DataBlockStartEvent): + return AGUICustomEvent( + name="data_block_start", + value=event.model_dump(exclude_none=True), + ) + + if isinstance(event, DataBlockDeltaEvent): + return AGUICustomEvent( + name="data_block_delta", + value=event.model_dump(exclude_none=True), + ) + + if isinstance(event, DataBlockEndEvent): + return AGUICustomEvent( + name="data_block_end", + value=event.model_dump(exclude_none=True), + ) + + if isinstance(event, RequireUserConfirmEvent): + return AGUICustomEvent( + name="require_user_confirm", + value=event.model_dump(exclude_none=True), + ) + + if isinstance(event, RequireExternalExecutionEvent): + return AGUICustomEvent( + name="require_external_execution", + value=event.model_dump(exclude_none=True), + ) + + if isinstance(event, UserConfirmResultEvent): + return AGUICustomEvent( + name="user_confirm_result", + value=event.model_dump(exclude_none=True), + ) + + if isinstance(event, ExternalExecutionResultEvent): + return AGUICustomEvent( + name="external_execution_result", + value=event.model_dump(exclude_none=True), + ) + + return AGUICustomEvent( + name="unknown", + value=event.model_dump(exclude_none=True), + ) diff --git a/src/agentscope/app/middleware/_protocol/_base.py b/src/agentscope/app/middleware/_protocol/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..64816dfa1993d48c82999578c2b68fe3e7361fad --- /dev/null +++ b/src/agentscope/app/middleware/_protocol/_base.py @@ -0,0 +1,244 @@ +# -*- coding: utf-8 -*- +"""Protocol middleware base class for converting AgentEvent stream to +various protocols.""" + +import json +from abc import ABC, abstractmethod +from typing import AsyncGenerator, Callable + +from fastapi import Request, Response +from fastapi.responses import StreamingResponse +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.types import ASGIApp + +from agentscope.event import AgentEvent + + +class ProtocolMiddlewareBase(BaseHTTPMiddleware, ABC): + """Base middleware for converting AgentEvent stream to protocol format. + + This middleware intercepts ``text/event-stream`` responses, deserializes + AgentEvent objects from SSE ``data:`` frames, and converts them to a + specific protocol format. + + Subclasses should implement the `_convert_to_protocol` method to define + the conversion logic for their specific protocol (e.g., AGUI, A2A). + + Example: + ```python + class AGUIMiddleware(ProtocolMiddlewareBase): + def _convert_to_protocol(self, event: AgentEvent) -> dict: + # Implement AGUI-specific conversion logic + return {...} + + app = FastAPI() + app.add_middleware(AGUIMiddleware) + ``` + """ + + def __init__(self, app: ASGIApp) -> None: + """Initialize the protocol middleware. + + Args: + app: The ASGI application to wrap. + """ + super().__init__(app) + + async def dispatch( + self, + request: Request, + call_next: Callable, + ) -> Response: + """Process the request and convert AgentEvent stream to protocol + format. + + Args: + request: The incoming HTTP request. + call_next: The next middleware or endpoint handler. + + Returns: + The response, potentially with converted stream content. + """ + # Call the next middleware or endpoint + response = await call_next(request) + + content_type = response.headers.get("content-type", "") + body_iterator = getattr(response, "body_iterator", None) + + if ( + content_type.startswith("text/event-stream") + and body_iterator is not None + ): + # Wrap the original stream with our conversion logic + converted_stream = self._convert_stream(body_iterator) + + # Create a new StreamingResponse with the converted stream + return StreamingResponse( + content=converted_stream, + status_code=response.status_code, + headers=dict(response.headers), + media_type=response.media_type, + ) + + return response + + async def _convert_stream( + self, + original_stream: AsyncGenerator, + ) -> AsyncGenerator[bytes, None]: + """Convert AgentEvent stream to protocol format. + + Args: + original_stream: The original stream yielding serialized + AgentEvent objects. + + Yields: + Bytes in protocol format. + """ + async for chunk in original_stream: + if isinstance(chunk, bytes): + chunk_str = chunk.decode("utf-8") + else: + chunk_str = chunk + + converted = self._convert_sse_frame(chunk_str) + if converted is not None: + yield converted + continue + + # Fallback for subclasses that may override dispatch() to handle + # non-SSE streams while still reusing this converter. + converted = self._convert_event_json(chunk_str) + if converted is not None: + yield converted + continue + + if isinstance(chunk, bytes): + yield chunk + else: + yield chunk.encode("utf-8") + + def _convert_sse_frame(self, frame: str) -> bytes | None: + """Convert AgentEvent payloads inside an SSE frame. + + Note: + This method targets the AgentScope service's SSE stream shape: + each ``data:`` line contains a complete JSON payload, and each + input ``frame`` contains one or more complete SSE frames. SSE + multi-line ``data:`` concatenation and cross-chunk frame + reassembly are intentionally out of scope here. + + Args: + frame: A server-sent event frame. + + Returns: + Converted frame bytes if at least one ``data:`` payload was + converted, otherwise ``None``. + """ + lines = frame.splitlines(keepends=True) + converted_lines: list[str] = [] + converted_any = False + + for line in lines: + if not line.startswith("data:"): + converted_lines.append(line) + continue + + line_content, line_ending = self._split_line_ending(line) + payload = line_content.removeprefix("data:") + if payload.startswith(" "): + payload = payload[1:] + + converted = self._convert_event_json(payload) + if converted is None: + converted_lines.append(line) + continue + + converted_json = converted.decode("utf-8").rstrip("\n") + converted_lines.append(f"data: {converted_json}{line_ending}") + converted_any = True + + if not converted_any: + return None + + return "".join(converted_lines).encode("utf-8") + + @staticmethod + def _split_line_ending(line: str) -> tuple[str, str]: + """Split a line into content and its original line ending.""" + if line.endswith("\r\n"): + return line[:-2], "\r\n" + if line.endswith("\n"): + return line[:-1], "\n" + if line.endswith("\r"): + return line[:-1], "\r" + return line, "" + + def _convert_event_json(self, chunk_str: str) -> bytes | None: + """Convert a serialized AgentEvent JSON string. + + Args: + chunk_str: Serialized AgentEvent JSON. + + Returns: + Converted protocol JSON bytes with trailing newline, or ``None`` + when ``chunk_str`` is not a valid AgentEvent payload. + """ + try: + event_dict = json.loads(chunk_str) + agent_event = self._deserialize_event(event_dict) + protocol_data = self._convert_to_protocol(agent_event) + return ( + json.dumps(protocol_data, ensure_ascii=False).encode( + "utf-8", + ) + + b"\n" + ) + except (json.JSONDecodeError, KeyError, TypeError, ValueError): + return None + + def _deserialize_event(self, event_dict: dict) -> AgentEvent: + """Deserialize event dictionary to AgentEvent object. + + Args: + event_dict: Dictionary containing event data with 'type' field. + + Returns: + Deserialized AgentEvent object. + + Raises: + ValueError: If event type is unknown or deserialization fails. + """ + from pydantic import Field, TypeAdapter + from typing import Annotated + + # Use Pydantic's discriminated union to automatically deserialize + # based on the 'type' field + adapter = TypeAdapter( + Annotated[AgentEvent, Field(discriminator="type")], + ) + return adapter.validate_python(event_dict) + + @abstractmethod + def _convert_to_protocol(self, event: AgentEvent) -> dict: + """Convert AgentEvent to protocol format. + + This is an abstract method that must be implemented by subclasses + to define the conversion logic for their specific protocol. + + Args: + event: The AgentEvent object to convert. + + Returns: + Dictionary in the target protocol format. + + Example: + ```python + class AGUIMiddleware(ProtocolMiddlewareBase): + def _convert_to_protocol(self, event: AgentEvent) -> dict: + # Convert to AGUI format + agui_data = event.model_dump() + agui_data["agui_version"] = "1.0" + return agui_data + ``` + """ diff --git a/src/agentscope/app/middleware/_state_change_middleware.py b/src/agentscope/app/middleware/_state_change_middleware.py new file mode 100644 index 0000000000000000000000000000000000000000..676dc24af37889eeece6648a29b6418fcd7d4a1a --- /dev/null +++ b/src/agentscope/app/middleware/_state_change_middleware.py @@ -0,0 +1,193 @@ +# -*- coding: utf-8 -*- +"""Middleware that detects agent state / team changes after each tool +call and pushes a :class:`CustomEvent` notification to the session's +event stream. + +Two kinds of change are detected: + +- **State change** — ``tasks_context`` or ``permission_context`` + modified (detected via hash comparison). Checked both around each + tool call (``on_acting``, for incremental updates during a turn) + and around the whole reply (``on_reply``, to catch changes made + outside the tool-execution window — e.g. permission rules added + while handling a user confirmation). Pushes + ``CustomEvent(name="state_updated", value={...})``. +- **Team change** — the tool that just ran is one of the team tools + (``TeamCreate``, ``AgentCreate``, ``TeamDelete``). These tools + directly mutate storage (``TeamRecord``, ``SessionRecord.team_id``), + so we don't need to check storage; the fact that the tool ran is + the trigger. Pushes ``CustomEvent(name="team_updated", value={})``. + +Both events are published directly to the bus (via +``session_publish_event``) instead of being yielded through the agent's +event chain, because ``on_acting`` yields ``ToolChunk | ToolResponse`` +— not ``AgentEvent``. The SSE ``/stream`` endpoint picks them up from +the bus like any other session event. +""" +import hashlib +from typing import Any, AsyncGenerator, Callable + +from ..message_bus import MessageBus +from .._bus_ops import publish_session_event +from ...event import CustomEvent +from ...middleware import MiddlewareBase + +_TEAM_TOOL_NAMES = frozenset({"TeamCreate", "AgentCreate", "TeamDelete"}) +# Tool names whose execution implies a team membership change. + + +class StateChangeMiddleware(MiddlewareBase): # pylint: disable=abstract-method + """Detect state / team changes after each tool call and push + notifications to the session event stream. + + Args: + message_bus (`MessageBus`): + Used to publish ``CustomEvent`` to the session's event + stream via :meth:`MessageBus.session_publish_event`. + session_id (`str`): + The session whose event stream to publish to. + """ + + def __init__( + self, + message_bus: MessageBus, + session_id: str, + ) -> None: + """Initialise the middleware. + + Args: + message_bus (`MessageBus`): + Application message bus. + session_id (`str`): + The session id to publish events for. + """ + self._bus = message_bus + self._session_id = session_id + + @staticmethod + def _state_hash(agent: Any) -> str: + """Compute a fast hash of the state fields we track. + + Only ``tasks_context`` and ``permission_context`` are included; + ``context`` (the message history) is intentionally excluded + because it changes on every reasoning step and is not what + this middleware cares about. + + Args: + agent: The agent instance. + + Returns: + `str`: A hex digest that changes when the tracked fields + change. + """ + raw = ( + agent.state.tasks_context.model_dump_json() + + agent.state.permission_context.model_dump_json() + ) + return hashlib.md5(raw.encode()).hexdigest() + + async def _publish_state(self, agent: Any) -> None: + """Push a ``state_updated`` event with the current tracked state. + + Args: + agent: The agent instance whose state to publish. + """ + event = CustomEvent( + name="state_updated", + value={ + "tasks_context": agent.state.tasks_context.model_dump( + mode="json", + ), + "permission_context": ( + agent.state.permission_context.model_dump( + mode="json", + ) + ), + }, + ) + await publish_session_event( + self._bus, + self._session_id, + event.model_dump(mode="json"), + ) + + async def on_reply( + self, + agent: Any, + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Wrap the whole reply turn to catch state changes that happen + **outside** the ``on_acting`` tool-execution window. + + Permission rules added while handling a + ``UserConfirmResultEvent`` (the user's "always allow" choice) + mutate ``permission_context`` in ``_handle_incoming_event`` — + which runs at the *start* of the reply turn, before the + confirmed tool's ``on_acting`` snapshot is taken. ``on_acting`` + therefore sees no diff (the rule is already present in both its + before- and after-hash) and never pushes. Snapshotting around + the entire reply closes that gap. + + Args: + agent: The executing agent. + input_kwargs (`dict`): + The reply inputs (new message(s) or a resumption event). + next_handler (`Callable[..., AsyncGenerator]`): + The downstream middleware or core reply logic. + + Yields: + ``AgentEvent | Msg`` — unchanged from downstream. + """ + hash_before = self._state_hash(agent) + + async for item in next_handler(**input_kwargs): + yield item + + hash_after = self._state_hash(agent) + if hash_before != hash_after: + await self._publish_state(agent) + + async def on_acting( + self, + agent: Any, + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Wrap tool execution: snapshot state hash before, compare + after, and push notifications if anything changed. + + Args: + agent: The executing agent. + input_kwargs (`dict`): + Contains ``tool_call`` (``ToolCallBlock``). + next_handler (`Callable[..., AsyncGenerator]`): + The downstream middleware or core acting logic. + + Yields: + ``ToolChunk | ToolResponse`` — unchanged from downstream. + """ + tool_call = input_kwargs.get("tool_call") + tool_name = tool_call.name if tool_call else "" + + hash_before = self._state_hash(agent) + + async for item in next_handler(**input_kwargs): + yield item + + # Check 1: state fields changed? + hash_after = self._state_hash(agent) + if hash_before != hash_after: + await self._publish_state(agent) + + # Check 2: team tool ran? + if tool_name in _TEAM_TOOL_NAMES: + event = CustomEvent( + name="team_updated", + value={}, + ) + await publish_session_event( + self._bus, + self._session_id, + event.model_dump(mode="json"), + ) diff --git a/src/agentscope/app/middleware/_tool_offload_middleware.py b/src/agentscope/app/middleware/_tool_offload_middleware.py new file mode 100644 index 0000000000000000000000000000000000000000..a80d52c69ebc9a2363855989ad0eab1636bea7c0 --- /dev/null +++ b/src/agentscope/app/middleware/_tool_offload_middleware.py @@ -0,0 +1,397 @@ +# -*- coding: utf-8 -*- +"""Middleware that offloads long-running tool calls to background tasks. + +When a tool times out, this middleware: + +- Lets the underlying asyncio task keep running via + :class:`BackgroundTaskManager` (the task is **never cancelled**). +- Yields a synthetic placeholder :class:`ToolResponse` so the agent + loop unblocks immediately. +- On task completion, pushes the real result as a + :class:`HintBlock` to the session's inbox via the message bus and + enqueues a wakeup. From there the result follows the same path as a + team message: :class:`InboxMiddleware` drains it into context on + the next reasoning step (same run or a fresh wakeup-driven run), + and :class:`WakeupDispatcher` ensures an idle session is started by + some process. + +The middleware no longer keeps its own pending-result store or +retrigger callback — both responsibilities now live in the +bus/wakeup infrastructure, which works correctly across processes. +""" +import asyncio +import json +from copy import deepcopy +from typing import AsyncGenerator, Callable + +from .._manager import BackgroundTaskManager +from ...middleware import MiddlewareBase +from ...tool import ToolChunk, ToolResponse +from ...message import ( + DataBlock, + HintBlock, + TextBlock, + ToolResultState, +) +from ...agent import Agent +from ..message_bus import MessageBus, MessageBusKeys +from .._bus_ops import enqueue_run_trigger +from ..._logging import logger + + +# Sentinel object used to signal end-of-stream in the drain queue. +_QUEUE_SENTINEL = object() + + +class ToolOffloadMiddleware(MiddlewareBase): # pylint: disable=abstract-method + """Middleware that offloads timed-out tool calls to background tasks. + + .. note:: + Tools with ``is_state_injected=True`` receive the live + ``agent.state`` object. Offloading them could race on shared + state, so they are always executed synchronously instead. + + Args: + bg_manager (`BackgroundTaskManager`): + Application-level background task manager. Used to register + the running asyncio task so :class:`ToolStop` can target it. + message_bus (`MessageBus`): + Application message bus. The completion callback uses it to + push the result HintBlock to the session's inbox and to + enqueue a wakeup so an idle session is woken on any process. + user_id (`str`): + User id of the current request — included in the wakeup + payload so the dispatcher can re-invoke ``ChatService.run``. + agent_id (`str`): + Agent record id (not the display name) — same purpose as + ``user_id``. + timeout_secs (`float`, defaults to ``10.0``): + Maximum seconds to wait for a tool execution before + offloading it to the background. + """ + + def __init__( + self, + bg_manager: BackgroundTaskManager, + message_bus: MessageBus, + user_id: str, + agent_id: str, + timeout_secs: float = 10.0, + ) -> None: + """Bind dependencies. + + Args: + bg_manager (`BackgroundTaskManager`): + Application background task manager. + message_bus (`MessageBus`): + Application message bus. + user_id (`str`): + User id of the current request. + agent_id (`str`): + Agent record id (not the display name). + timeout_secs (`float`, defaults to ``10.0``): + Tool execution timeout before offloading to background. + """ + self._bg_manager = bg_manager + self._message_bus = message_bus + self._user_id = user_id + self._agent_id = agent_id + self._timeout_secs = timeout_secs + + # ------------------------------------------------------------------ + # Middleware hooks + # ------------------------------------------------------------------ + + async def on_acting( # type:ignore[override] + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Execute a tool with timeout; offload to background on expiry. + + The inner ``next_handler`` generator is wrapped in an + :mod:`asyncio` task whose output is fed through a + :class:`asyncio.Queue`. Items are consumed with a rolling + deadline. If the deadline fires before the tool finishes: + + - The running task is **not** cancelled. + - It is registered with :attr:`_bg_manager` so that + :class:`ToolStop` can target it and shutdown can cancel it. + - A separate watcher coroutine is spawned to await the task's + completion and then push the result as a + :class:`HintBlock` to the session inbox + enqueue a wakeup. + - A synthetic :class:`~agentscope.tool.ToolResponse` notifying + the agent of the background task id is yielded instead. + + .. note:: + Tools with ``is_state_injected=True`` or ``is_external_tool=True`` + bypass this logic and are always executed synchronously. + + Args: + agent (`Agent`): + The executing agent. + input_kwargs (`dict`): + Acting input kwargs (contains ``tool_call``). + next_handler (`Callable[..., AsyncGenerator]`): + The downstream middleware or ``_acting_impl``. + + Yields: + `ToolChunk | ToolResponse`: + Normal results when the tool finishes in time, or a + synthetic ``ToolResponse`` when offloaded. + """ + tool_call = input_kwargs["tool_call"] + + # ---------------------------------------------------------------- + # Guard: state-injected tools and external tools are never offloaded. + # - is_state_injected: the tool receives the live agent.state object; + # running it concurrently in a background task could cause race + # conditions on agent.state. For now, we fall back to synchronous + # execution. Once background tasks can be given an isolated state + # snapshot this guard should become a hard RuntimeError instead. + # - is_external_tool: external tools wait for a human/external system + # to push a result back; offloading them makes no sense because the + # agent would lose track of the pending confirmation. + tool = await agent.toolkit.get_tool(tool_call.name) + if tool is not None and ( + tool.is_state_injected or tool.is_external_tool + ): + async for item in next_handler(**input_kwargs): + yield item + return + + # ---------------------------------------------------------------- + # Wrap next_handler in a Task, draining output into a Queue + # ---------------------------------------------------------------- + queue: asyncio.Queue = asyncio.Queue() + + async def _drain_to_queue() -> None: + """Drain next_handler output into *queue*.""" + try: + async for item in next_handler(**input_kwargs): + await queue.put(item) + except Exception as exc: # pylint: disable=broad-except + await queue.put(exc) + finally: + await queue.put(_QUEUE_SENTINEL) + + drain_task: asyncio.Task = asyncio.create_task(_drain_to_queue()) + + # ---------------------------------------------------------------- + # Consume items until the deadline or normal completion + # ---------------------------------------------------------------- + loop = asyncio.get_event_loop() + deadline = loop.time() + self._timeout_secs + pre_collected: list = [] + completed = False + + while True: + remaining = deadline - loop.time() + if remaining <= 0: + break + + try: + item = await asyncio.wait_for( + queue.get(), + timeout=remaining, + ) + except asyncio.TimeoutError: + break + + if item is _QUEUE_SENTINEL: + completed = True + break + + if isinstance(item, BaseException): + drain_task.cancel() + raise item + + pre_collected.append(item) + + # ToolResponse is always the terminal item from _acting_impl + if isinstance(item, ToolResponse): + completed = True + break + + if completed: + for item in pre_collected: + yield item + drain_task.cancel() + return + + # ---------------------------------------------------------------- + # Timeout path: spawn deliverer, register drain_task, yield synthetic + # ---------------------------------------------------------------- + session_id = agent.state.session_id + tool_name = tool_call.name + snapshot = list(pre_collected) + + logger.info( + "Tool '%s' timed out after %.1fs, offloading to background: " + "session_id=%s, agent_id=%s", + tool_name, + self._timeout_secs, + session_id, + agent.name, + ) + + async def _deliver_when_done() -> None: + """Wait for the offloaded tool to finish, then push its + result to the session inbox + enqueue a wakeup. + + On cancellation (e.g. ``ToolStop``) or unhandled exception + no result is delivered — the agent is left without a + completion notification in those edge cases. + """ + try: + await drain_task + except asyncio.CancelledError: + logger.info( + "Background tool '%s' cancelled, skipping delivery: " + "session_id=%s", + tool_name, + session_id, + ) + return + except Exception: # pylint: disable=broad-except + logger.warning( + "Background tool '%s' failed, skipping delivery: " + "session_id=%s", + tool_name, + session_id, + exc_info=True, + ) + return + + remaining_items: list = [] + while not queue.empty(): + try: + item = queue.get_nowait() + if isinstance(item, ToolResponse): + remaining_items.append(item) + break + except asyncio.QueueEmpty: + break + + all_items = snapshot + remaining_items + response: ToolResponse | None = next( + (i for i in all_items if isinstance(i, ToolResponse)), + None, + ) + + tool_call_id = tool_call.id + + hint_source = json.dumps( + { + "label": "tool_output", + "sublabel": f"{tool_name} · {tool_call_id}", + }, + ensure_ascii=False, + ) + + if response is None or len(response.content) == 0: + hint = HintBlock( + hint=( + f"" + f"Tool '{tool_name}' running in background " + f"(id={tool_call_id}) has completed with no output." + f"" + ), + source=hint_source, + ) + + else: + # Preserve all content blocks (text + multimodal) from + # the tool response so nothing is lost. + content_blocks: list[TextBlock | DataBlock] = deepcopy( + response.content, + ) + + prefix = ( + f"" + f"Tool '{tool_name}' running in background " + f"(id={tool_call_id}) has completed.\n\n" + f"Result:\n\n" + ) + + if isinstance(content_blocks[0], TextBlock): + content_blocks[0].text = prefix + content_blocks[0].text + + else: + content_blocks.insert(0, TextBlock(text=prefix)) + + suffix = "" + + if isinstance(content_blocks[-1], TextBlock): + content_blocks[-1].text += suffix + else: + content_blocks.append( + TextBlock(text=suffix), + ) + + hint = HintBlock( + hint=content_blocks, + source=hint_source, + ) + + # Deliver via inbox + wakeup — same path as a team message. + # InboxMiddleware drains the inbox into context on the next + # reasoning step, and WakeupDispatcher kicks an idle session + # on whichever process picks up the wakeup. + logger.info( + "Background tool '%s' completed, pushing result to inbox " + "and enqueueing wakeup: session_id=%s", + tool_name, + session_id, + ) + await self._message_bus.queue_push( + MessageBusKeys.inbox(session_id), + hint.model_dump(mode="json"), + ) + await enqueue_run_trigger( + self._message_bus, + user_id=self._user_id, + session_id=session_id, + agent_id=self._agent_id, + ) + + asyncio.create_task(_deliver_when_done()) + task_id = await self._bg_manager.register_task( + asyncio_task=drain_task, + session_id=session_id, + agent_id=self._agent_id, + user_id=self._user_id, + tool_name=tool_name, + ) + + logger.info( + "Synthetic ToolResponse yielded for offloaded tool '%s': " + "task_id=%s, session_id=%s, agent_id=%s", + tool_name, + task_id, + session_id, + agent.name, + ) + + placeholder_text = f"""Tool '{tool_name}' is \ +running in background (id={task_id}) for over {self._timeout_secs}s. \ +You will be notified automatically when it finishes, so **DO NOT** poll, \ +query, or wait for the result yourself. **DO NOT** call any waiting tool \ +such as `bash sleep`. You have exactly two valid options: +1. Continue with other independent tasks and ignore this tool for now; or +2. If there is nothing else to do, simply give a text reply without calling \ +any tool, which ends the current reasoning loop — just do nothing and end \ +this run. +""" + + yield ToolChunk( + content=[TextBlock(text=placeholder_text)], + state=ToolResultState.SUCCESS, + ) + + yield ToolResponse( + content=[TextBlock(text=placeholder_text)], + state=ToolResultState.SUCCESS, + id=tool_call.id, + ) diff --git a/src/agentscope/app/rag/__init__.py b/src/agentscope/app/rag/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0aec3e62524ed3692daa8fecc3998fe42a750af0 --- /dev/null +++ b/src/agentscope/app/rag/__init__.py @@ -0,0 +1,92 @@ +# -*- coding: utf-8 -*- +"""Service-layer RAG building blocks. + +This subpackage groups every RAG-specific service-layer concept under +one roof so the user-facing import surface stays compact: + +- :mod:`.blob_store` — backends storing uploaded document bytes; +- :mod:`.knowledge_base_manager` — knowledge base lifecycle and the + runtime :class:`~agentscope.rag.KnowledgeBase` handle shared between + HTTP and agents; +- :mod:`.index_worker` — out-of-process indexing worker entry point. + +Commonly used symbols are re-exported here so the typical user import +collapses to a single line:: + + from agentscope.app.rag import ( + LocalBlobStore, + S3BlobStore, + CollectionPerKbManager, + run_worker, + ) + +For type-level imports or rare subclasses, use the submodules directly. +""" +from typing import TYPE_CHECKING + +from .blob_store import ( + AsyncReadable, + BlobStoreBase, + LocalBlobStore, + S3BlobStore, +) +from .knowledge_base_manager import ( + CollectionPerKbManager, + DimensionPolicy, + DimensionPolicyError, + DimensionPolicyKind, + KnowledgeBaseError, + KnowledgeBaseManagerBase, + KnowledgeBaseNotFoundError, +) + +if TYPE_CHECKING: + # Re-exported lazily at runtime via ``__getattr__`` to break a + # package-load import cycle; declared here so static analysers + # (pylint, mypy, IDEs) still see the symbol on the package. + from .index_worker import run_worker # noqa: F401 + + +def __getattr__(name: str) -> object: + """Lazy-load attributes that would cause an import cycle. + + ``run_worker`` is lazy-imported because :mod:`.index_worker` pulls + in :mod:`agentscope.app._service`, which itself imports + :mod:`agentscope.app.rag` for the middleware — a direct re-export + would create an import cycle at package load. Users who actually + want the worker (``from agentscope.app.rag import run_worker``) + pay the import cost on first access. + + Args: + name (`str`): + The attribute name requested by the importer. + + Returns: + `object`: + The resolved attribute value. + + Raises: + `AttributeError`: + When ``name`` is not exposed by this package. + """ + if name == "run_worker": + from .index_worker import run_worker as _run_worker + + return _run_worker + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = [ + "AsyncReadable", + "BlobStoreBase", + "CollectionPerKbManager", + "DimensionPolicy", + "DimensionPolicyError", + "DimensionPolicyKind", + "KnowledgeBaseError", + "KnowledgeBaseManagerBase", + "KnowledgeBaseNotFoundError", + "LocalBlobStore", + "S3BlobStore", + "run_worker", +] diff --git a/src/agentscope/app/rag/blob_store/__init__.py b/src/agentscope/app/rag/blob_store/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..29eafd155508f1cf767463ed0f624b5c9bd9d38a --- /dev/null +++ b/src/agentscope/app/rag/blob_store/__init__.py @@ -0,0 +1,12 @@ +# -*- coding: utf-8 -*- +"""Blob storage backends for document uploads.""" +from ._base import AsyncReadable, BlobStoreBase +from ._local import LocalBlobStore +from ._s3 import S3BlobStore + +__all__ = [ + "AsyncReadable", + "BlobStoreBase", + "LocalBlobStore", + "S3BlobStore", +] diff --git a/src/agentscope/app/rag/blob_store/_base.py b/src/agentscope/app/rag/blob_store/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..8421058bd1e7bb929c70a69a6c91a336dccc167c --- /dev/null +++ b/src/agentscope/app/rag/blob_store/_base.py @@ -0,0 +1,158 @@ +# -*- coding: utf-8 -*- +"""Abstract base class for blob storage backends. + +A :class:`BlobStoreBase` is the byte-level home of files uploaded into +the application — knowledge base documents in v1, potentially other +binary payloads later. It is created once at app startup and shared +across requests, mirroring the lifecycle of +:class:`~agentscope.app.storage.StorageBase` and +:class:`~agentscope.rag.VectorStoreBase`. + +The abstraction owes its existence to a single hard requirement: bytes +must never sit in memory for the duration of indexing. The upload +endpoint streams the request body into ``write_stream``; the indexing +worker streams the bytes back out via ``open``. Implementations decide +where those bytes physically live (local disk, S3-compatible object +store, etc.) — neither caller cares. +""" +from abc import ABC, abstractmethod +from contextlib import AbstractAsyncContextManager +from typing import IO, Any, Protocol, Self + + +class AsyncReadable(Protocol): + """Minimal async-read protocol returned by :meth:`BlobStoreBase.open`. + + Backed by ``aiofiles`` handles for the local store and by streaming + response bodies for object-storage backends. The narrow interface + keeps the worker's blob-reading code identical across backends + even though no single concrete type implements both shapes. + """ + + async def read(self, n: int = -1) -> bytes: + """Read up to ``n`` bytes (``-1`` reads to EOF). + + Args: + n (`int`, defaults to ``-1``): + The maximum number of bytes to read; ``-1`` drains the + stream. + + Returns: + `bytes`: + The bytes read; an empty ``bytes`` at EOF. + """ + + +class BlobStoreBase(ABC): + """Abstract base class for blob storage backends. + + Lifecycle is managed via the async context manager protocol so the + app lifespan can open / close connection pools or filesystem + handles uniformly. + """ + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def __aenter__(self) -> Self: + """Enter the async context — open connections / create dirs. + + The default implementation is a no-op. Subclasses that need + explicit setup should override this. + + Returns: + `BlobStoreBase`: + ``self``. + """ + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: Any, + ) -> None: + """Exit the async context — close connections if needed.""" + + # ------------------------------------------------------------------ + # Byte operations + # ------------------------------------------------------------------ + + @abstractmethod + async def write_stream(self, key: str, stream: IO[bytes]) -> str: + """Stream-write a blob and return its URI. + + Implementations MUST copy the bytes in chunks (typically ~1 MB) + and MUST NOT call ``stream.read()`` without a size argument — + the whole point of the abstraction is that the byte payload + never lives in memory all at once. + + Args: + key (`str`): + Backend-relative key, e.g. ``"kb/{kb_id}/{doc_id}"``. + The caller picks the key layout; the backend turns it + into a backend-native location (filesystem path, + ``s3://`` object name, ...). + stream (`IO[bytes]`): + A readable binary stream. Synchronous read API + (``read(n)``) — ``UploadFile.file`` from FastAPI fits + directly, and other binary streams are wrapped via the + same protocol. + + Returns: + `str`: + A scheme-qualified URI (e.g. ``"local://kb/.../uuid"``, + ``"s3://bucket/kb/.../uuid"``) that round-trips back + through :meth:`open` / :meth:`delete` / :meth:`exists`. + """ + + @abstractmethod + async def open( + self, + uri: str, + ) -> AbstractAsyncContextManager[AsyncReadable]: + """Stream-read a blob by URI. + + Returns an async context manager whose ``__aenter__`` yields a + stream with an awaitable ``read(n)`` method. Implementations + decide whether the stream is backed by an ``aiofiles`` handle, + a streaming response body, or any other async source — + callers MUST treat it as forward-only and chunked. + + Args: + uri (`str`): + A URI produced by :meth:`write_stream`. + + Returns: + `AbstractAsyncContextManager[AsyncReadable]`: + The byte stream, wrapped so the backend can release + handles deterministically on context exit. + """ + + @abstractmethod + async def delete(self, uri: str) -> None: + """Delete the blob at ``uri`` if it exists. + + MUST be idempotent — a no-op when the blob is already gone. + The caller (document deletion path) treats blob deletion as + best-effort cleanup; partial failures are recoverable by + re-running the sweep. + + Args: + uri (`str`): + A URI produced by :meth:`write_stream`. + """ + + @abstractmethod + async def exists(self, uri: str) -> bool: + """Return whether the blob at ``uri`` is present. + + Args: + uri (`str`): + A URI produced by :meth:`write_stream`. + + Returns: + `bool`: + ``True`` if the blob is currently retrievable. + """ diff --git a/src/agentscope/app/rag/blob_store/_local.py b/src/agentscope/app/rag/blob_store/_local.py new file mode 100644 index 0000000000000000000000000000000000000000..e49f8c2e23c87494d2d5512522d0de781c9bfb92 --- /dev/null +++ b/src/agentscope/app/rag/blob_store/_local.py @@ -0,0 +1,145 @@ +# -*- coding: utf-8 -*- +"""Local filesystem implementation of :class:`BlobStoreBase`. + +The backend reserves a single root directory and treats keys as +relative paths beneath it. URIs are formatted as ``local://{key}`` — +the scheme acts as a discriminator so a mixed-deployment app can route +``local://`` and ``s3://`` URIs to different backends without +re-parsing the rest of the string. +""" +import os +import shutil +from contextlib import asynccontextmanager +from pathlib import Path +from typing import IO, AsyncIterator + +import aiofiles +import aiofiles.os +import aiofiles.ospath + +from ._base import AsyncReadable, BlobStoreBase + + +_SCHEME = "local://" +_CHUNK_SIZE = 1 << 20 # 1 MiB — keep peak memory bounded, match S3 part size. + + +class LocalBlobStore(BlobStoreBase): + """Store blobs as files beneath a configurable root directory. + + Intended for single-node deployments and development. The lifecycle + of a blob is bound to the document that owns it, not to the request + that uploaded it — this is what distinguishes the store from + FastAPI's per-request ``SpooledTemporaryFile`` and is the reason + the abstraction exists at all. + """ + + def __init__(self, root_dir: str | os.PathLike) -> None: + """Initialize the local backend. + + Args: + root_dir (`str | PathLike`): + Directory that holds all blobs for this app. Created + on :meth:`__aenter__` if it does not exist. + """ + self._root = Path(root_dir).resolve() + + async def __aenter__(self) -> "LocalBlobStore": + """Ensure the root directory exists.""" + await aiofiles.os.makedirs(self._root, exist_ok=True) + return self + + def _path_for(self, key: str) -> Path: + """Resolve a backend-relative key to an absolute filesystem path. + + Rejects keys that try to escape :attr:`_root` via ``..`` or + absolute paths. Callers pick keys server-side so this is + defensive rather than a primary trust boundary, but we still + refuse to write outside the root. + + Args: + key (`str`): + Backend-relative key. + + Returns: + `Path`: + The absolute filesystem path within :attr:`_root`. + """ + if not key or key.startswith("/") or ".." in Path(key).parts: + raise ValueError(f"Invalid blob key: {key!r}") + path = (self._root / key).resolve() + if self._root not in path.parents and path != self._root: + raise ValueError(f"Blob key {key!r} escapes the root directory.") + return path + + def _key_from_uri(self, uri: str) -> str: + """Extract the backend-relative key from a ``local://`` URI.""" + if not uri.startswith(_SCHEME): + raise ValueError(f"Not a local blob URI: {uri!r}") + return uri[len(_SCHEME) :] + + async def write_stream(self, key: str, stream: IO[bytes]) -> str: + """Copy ``stream`` into the blob at ``key`` in 1 MiB chunks. + + Creates intermediate directories as needed. Existing blobs at + the same key are overwritten — keys are generated server-side + from document ids, so collisions only happen on intentional + re-uploads. + + Args: + key (`str`): + Backend-relative key. + stream (`IO[bytes]`): + Synchronous binary source. + + Returns: + `str`: + The ``local://{key}`` URI. + """ + path = self._path_for(key) + await aiofiles.os.makedirs(path.parent, exist_ok=True) + async with aiofiles.open(path, "wb") as out: + while True: + chunk = stream.read(_CHUNK_SIZE) + if not chunk: + break + await out.write(chunk) + return f"{_SCHEME}{key}" + + @asynccontextmanager + async def open( # type: ignore[override] + self, + uri: str, + ) -> AsyncIterator[AsyncReadable]: + """Open the blob at ``uri`` for streaming reads.""" + path = self._path_for(self._key_from_uri(uri)) + async with aiofiles.open(path, "rb") as fp: + yield fp + + async def delete(self, uri: str) -> None: + """Remove the blob at ``uri`` if present (idempotent).""" + path = self._path_for(self._key_from_uri(uri)) + try: + await aiofiles.os.remove(path) + except FileNotFoundError: + return + # Best-effort cleanup of empty parent directories up to the root. + parent = path.parent + while parent != self._root and parent.is_relative_to(self._root): + try: + await aiofiles.os.rmdir(parent) + except OSError: + break + parent = parent.parent + + async def exists(self, uri: str) -> bool: + """Return whether the blob at ``uri`` is present.""" + path = self._path_for(self._key_from_uri(uri)) + return await aiofiles.ospath.isfile(path) + + # Synchronous helper used by tests / cleanup scripts; not part of + # the public abstraction. + def _wipe(self) -> None: + """Delete the entire root directory. Test-only.""" + if self._root.exists(): + shutil.rmtree(self._root) diff --git a/src/agentscope/app/rag/blob_store/_s3.py b/src/agentscope/app/rag/blob_store/_s3.py new file mode 100644 index 0000000000000000000000000000000000000000..45b5286ad3e11615cceee471ef6c38e8d3b14e97 --- /dev/null +++ b/src/agentscope/app/rag/blob_store/_s3.py @@ -0,0 +1,276 @@ +# -*- coding: utf-8 -*- +"""S3-compatible implementation of :class:`BlobStoreBase`. + +Works against any service that implements the S3 wire protocol — +AWS S3, MinIO, Cloudflare R2, Aliyun OSS (S3-compatible), Tencent yun COS, etc. +The discriminator is ``endpoint_url``: ``None`` means real AWS S3 +(``aioboto3`` resolves the regional endpoint); any other value points +at the compatible service's S3 endpoint. + +URIs are formatted as ``s3://{bucket}/{key}``. The bucket is baked +into the URI (rather than carried out-of-band) so a single deployment +can read blobs that were written under a different bucket — for +example, after a bucket migration — without breaking existing +document records. + +The implementation is intentionally minimal: this is a byte store, +not a multipart upload optimiser. Production-scale uploads (>5 GiB) +would need the multipart API explicitly, but our documents are +bounded by ``UPLOAD_MAX_BYTES`` well below that ceiling, so a single +``upload_fileobj`` call (which uses multipart internally when the +body is large enough) is sufficient. +""" +from contextlib import asynccontextmanager +from typing import IO, Any, AsyncIterator + +from ._base import AsyncReadable, BlobStoreBase + + +_SCHEME = "s3://" + + +class _StreamingBody: + """Adapter that exposes :class:`AsyncReadable` over an aioboto3 body. + + ``aioboto3`` returns ``aiohttp``-backed streaming bodies that + already expose ``read(n)`` — but ``n=-1`` is "read until EOF" + which our ``IndexWorker._read_blob`` does not want; it asks for + a bounded chunk every iteration. We wrap the body in a thin + object so the adapter contract is explicit and so the underlying + body type can change without the worker noticing. + """ + + def __init__(self, body: Any) -> None: + """Initialize the body.""" + self._body = body + + async def read(self, n: int = -1) -> bytes: + """Read up to *n* bytes; ``-1`` reads to EOF.""" + if n == -1: + return await self._body.read() + # aiohttp's StreamReader.read(n) returns up to n bytes; if + # the server-side stream is bursty, we may get a short read, + # which is exactly the contract we want for chunked indexing. + return await self._body.read(n) + + +class S3BlobStore(BlobStoreBase): + """Store blobs in an S3-compatible bucket.""" + + def __init__( + self, + bucket: str, + *, + region_name: str | None = None, + endpoint_url: str | None = None, + aws_access_key_id: str | None = None, + aws_secret_access_key: str | None = None, + session_token: str | None = None, + use_ssl: bool = True, + config: Any = None, + ) -> None: + """Initialize an S3-compatible bucket. + + Args: + bucket (`str`): + Bucket name. The bucket MUST exist; this store will not + create it. Production deployments commonly provision the + bucket out-of-band via Terraform / CloudFormation so the + app's IAM role does not need ``s3:CreateBucket``. + region_name (`str | None`, optional): + AWS region for AWS S3. Required for AWS, ignored for + MinIO / R2 / other services that locate the bucket via + ``endpoint_url``. + endpoint_url (`str | None`, optional): + Full URL of the S3-compatible service. ``None`` selects + real AWS S3. Examples: ``http://minio:9000``, + ``https://.r2.cloudflarestorage.com``, + ``https://oss-cn-hangzhou.aliyuncs.com``. + aws_access_key_id (`str | None`, optional): + Static credential. Prefer leaving this ``None`` and letting + the IAM role / environment chain resolve credentials when + running on AWS infrastructure. + aws_secret_access_key (`str | None`, optional): + Paired with ``aws_access_key_id``. + session_token (`str | None`, optional): + For STS-issued temporary credentials. + use_ssl (`bool`, defaults to ``True``): + Force HTTPS. Production deployments must keep this on; + local MinIO with self-signed certs is the only place + ``False`` is reasonable. + config (`Any | None`, optional): + ``aiobotocore.config.AioConfig`` instance for users who + need to tune timeouts, retry mode, signature version + (e.g. ``s3v4`` for Aliyun OSS), or addressing style. + Path-style addressing is needed for MinIO; pass + ``AioConfig(s3={"addressing_style": "path"})``. + """ + try: + import aioboto3 + except ImportError as e: + raise ImportError( + "S3BlobStore requires the optional dependency ``aioboto3``. " + "Install it with ``uv pip install aioboto3`` or with the " + "``[s3]`` extra.", + ) from e + + self._bucket = bucket + self._region_name = region_name + self._endpoint_url = endpoint_url + self._aws_access_key_id = aws_access_key_id + self._aws_secret_access_key = aws_secret_access_key + self._session_token = session_token + self._use_ssl = use_ssl + self._config = config + + # Session is cheap; the actual transport is the ``client`` + # context manager opened per call. We deliberately do NOT + # cache a long-lived client across the lifespan: aiobotocore + # clients hold an aiohttp connection pool tied to a loop, and + # the per-call cost is negligible compared to the network + # round-trip. + self._session: ( + aioboto3.Session | None + ) = None # type: ignore[name-defined] + + async def __aenter__(self) -> "S3BlobStore": + try: + import aioboto3 + except ImportError as e: + raise ImportError( + "S3BlobStore requires the optional dependency ``aioboto3``. " + "Install it with ``uv pip install aioboto3`` or with the " + "``[s3]`` extra.", + ) from e + + self._session = aioboto3.Session() + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: Any, + ) -> None: + self._session = None + + def _client(self) -> Any: + """Open a fresh S3 client context manager for one call.""" + if self._session is None: + raise RuntimeError( + "S3BlobStore is not entered; use it inside " + "``async with`` before calling blob methods.", + ) + return self._session.client( + "s3", + region_name=self._region_name, + endpoint_url=self._endpoint_url, + aws_access_key_id=self._aws_access_key_id, + aws_secret_access_key=self._aws_secret_access_key, + aws_session_token=self._session_token, + use_ssl=self._use_ssl, + config=self._config, + ) + + @staticmethod + def _parse_uri(uri: str) -> tuple[str, str]: + """Split an ``s3://{bucket}/{key}`` URI into ``(bucket, key)``.""" + if not uri.startswith(_SCHEME): + raise ValueError(f"Not an S3 blob URI: {uri!r}") + rest = uri[len(_SCHEME) :] + bucket, _, key = rest.partition("/") + if not bucket or not key: + raise ValueError(f"Malformed S3 blob URI: {uri!r}") + return bucket, key + + @classmethod + def _key_from_uri(cls, uri: str, expected_bucket: str) -> str: + """Return the object key, asserting the bucket matches. + + Used by mutating operations (``delete``, ``exists``) where + crossing into another bucket would be a bug — the configured + bucket is the only place the store owns objects. + """ + bucket, key = cls._parse_uri(uri) + if bucket != expected_bucket: + raise ValueError( + f"Bucket {bucket!r} in URI {uri!r} does not match " + f"configured bucket {expected_bucket!r}.", + ) + return key + + async def write_stream(self, key: str, stream: IO[bytes]) -> str: + """Stream-write a blob and return its ``s3://{bucket}/{key}`` URI. + + Uses ``upload_fileobj`` so aioboto3 picks multipart upload + automatically for bodies above the multipart threshold + (8 MiB by default in botocore). For smaller bodies it + promotes to a single ``PutObject`` call. + """ + async with self._client() as s3: + await s3.upload_fileobj(stream, self._bucket, key) + return f"{_SCHEME}{self._bucket}/{key}" + + @asynccontextmanager + async def open( # type: ignore[override] + self, + uri: str, + ) -> AsyncIterator[AsyncReadable]: + """Stream-read a blob by URI. + + Yields an :class:`AsyncReadable` backed by the response + ``Body``. The S3 GET response is held open for the duration + of the ``async with`` block — exit promptly so the connection + returns to the pool. + + The bucket is read from the URI rather than the configured + bucket so post-migration deployments can still resolve + document records that were written under the previous bucket. + The IAM role must grant ``s3:GetObject`` on the legacy bucket + for this to actually succeed at the wire level. + """ + bucket, key = self._parse_uri(uri) + async with self._client() as s3: + response = await s3.get_object(Bucket=bucket, Key=key) + body = response["Body"] + try: + yield _StreamingBody(body) + finally: + # aiobotocore's body has a ``close`` coroutine; call + # it explicitly so the connection is released rather + # than waiting for GC. + close = getattr(body, "close", None) + if close is not None: + result = close() + if hasattr(result, "__await__"): + await result + + async def delete(self, uri: str) -> None: + """Delete the object at ``uri``. Idempotent.""" + key = self._key_from_uri(uri, self._bucket) + async with self._client() as s3: + # S3 ``DeleteObject`` is idempotent on a missing key — + # no special handling needed for the "already gone" case. + await s3.delete_object(Bucket=self._bucket, Key=key) + + async def exists(self, uri: str) -> bool: + """Return whether the object at ``uri`` is present.""" + key = self._key_from_uri(uri, self._bucket) + async with self._client() as s3: + try: + await s3.head_object(Bucket=self._bucket, Key=key) + return True + except Exception as exc: # noqa: BLE001 + # botocore's ClientError has ``response["Error"]["Code"]`` + # of "404" or "NoSuchKey" for missing objects; we + # treat anything that isn't a successful HEAD as + # "missing" rather than re-raising, matching the + # local backend's permissive contract. A genuinely + # broken backend (auth failure, network) will surface + # on the next read/write. + code = ( + getattr(exc, "response", {}).get("Error", {}).get("Code") + ) + if code in {"404", "NoSuchKey", "NotFound"}: + return False + raise diff --git a/src/agentscope/app/rag/index_worker/__init__.py b/src/agentscope/app/rag/index_worker/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e03a9ad9157824fea5e7457f4494f97f213c362f --- /dev/null +++ b/src/agentscope/app/rag/index_worker/__init__.py @@ -0,0 +1,198 @@ +# -*- coding: utf-8 -*- +"""Out-of-process index worker entry point. + +A worker process owns: + +- a :class:`~agentscope.app._service.IndexWorker` instance, which + runs the parse → chunk → embed pipeline for one document at a time; +- an :class:`~agentscope.app._service.IndexTaskConsumer`, which + subscribes to the shared task channel on the message bus and + feeds the worker. + +It does NOT own the sweeper — the API process keeps running the +sweeper so that documents stuck in ``pending`` (because the API +crashed between the storage write and the dispatcher publish) still +recover, even if every worker process happens to be offline at the +moment the publish was attempted. + +This module is a library: a deployment wires its concrete backends +through :func:`run_worker` from whatever bootstrap script it uses +(systemd unit, Kubernetes Deployment + container entrypoint, +docker-compose ``command``, ...). The shape mirrors +:func:`agentscope.app.create_app` — pass already-constructed +backend instances, the worker manages their lifecycle through an +:class:`AsyncExitStack`. + +Example:: + + import asyncio + import socket + import uuid + + from agentscope.app.rag.blob_store import S3BlobStore + from agentscope.app.rag.knowledge_base_manager import ( + DefaultKnowledgeBaseManager, + ) + from agentscope.app.message_bus import RedisMessageBus + from agentscope.app.storage import RedisStorage + from agentscope.app.rag.index_worker import run_worker + from agentscope.rag import ApproxTokenChunker, TextParser, ... + + async def main() -> None: + storage = RedisStorage(url=os.environ["REDIS_URL"]) + message_bus = RedisMessageBus(url=os.environ["REDIS_URL"]) + blob_store = S3BlobStore( + bucket=os.environ["S3_BUCKET"], + endpoint_url=os.environ.get("S3_ENDPOINT"), + ) + kb_manager = DefaultKnowledgeBaseManager(...) + parsers = [TextParser()] + chunker = ApproxTokenChunker() + await run_worker( + storage=storage, + message_bus=message_bus, + blob_store=blob_store, + knowledge_base_manager=kb_manager, + parsers=parsers, + chunker=chunker, + ) + + if __name__ == "__main__": + asyncio.run(main()) +""" +import asyncio +import signal +import socket +import uuid +from concurrent.futures import ProcessPoolExecutor +from contextlib import AsyncExitStack +from typing import TYPE_CHECKING + +from ..._service import IndexTaskConsumer, IndexWorker +from ...._logging import logger + +if TYPE_CHECKING: + from ..blob_store import BlobStoreBase + from ..knowledge_base_manager import KnowledgeBaseManagerBase + from ...message_bus import MessageBus + from ...storage import StorageBase + from ....rag import ChunkerBase, ParserBase + + +async def run_worker( + *, + storage: "StorageBase", + message_bus: "MessageBus", + blob_store: "BlobStoreBase", + knowledge_base_manager: "KnowledgeBaseManagerBase", + parsers: "list[ParserBase] | dict[str, ParserBase]", + chunker: "ChunkerBase", + node_id: str | None = None, + worker_max_concurrency: int = 4, + consumer_max_batch: int = 32, + parser_executor: ProcessPoolExecutor | None = None, +) -> None: + """Run the out-of-process index worker until cancelled. + + Manages the lifecycle of every passed-in backend through a single + :class:`AsyncExitStack`. The function returns when SIGINT or + SIGTERM is delivered (deployment-controlled shutdown). + + Args: + storage (`StorageBase`): + Persistent storage backend. Lifecycle managed here. + message_bus (`MessageBus`): + Live message bus subscribed to for index-task signals. + Lifecycle managed here. + blob_store (`BlobStoreBase`): + Backend the worker reads document bytes from. Lifecycle + managed here. + knowledge_base_manager (`KnowledgeBaseManagerBase`): + Resolves :class:`KnowledgeBase` runtimes for embedding + + vector store writes. Lifecycle managed here. + parsers (`list[ParserBase] | dict[str, ParserBase]`): + Parsers used to dispatch uploads by IANA media type. + Pass the same registry both to ``create_app`` and to the + worker — list mode expands each parser's + ``supported_media_types`` (later entries override earlier + ones, with a warning); dict mode is used verbatim for + explicit routing. + chunker (`ChunkerBase`): + Shared chunker. + node_id (`str | None`, optional): + Stable identifier for this worker process used on the + storage lease. Defaults to + ``"{hostname}:{uuid-prefix}"`` — good enough for + distinguishing two workers running on the same host; + override when your orchestrator has a more meaningful + id (Kubernetes pod name, ECS task id, ...). + worker_max_concurrency (`int`, defaults to ``4``): + Maximum number of documents the worker processes + concurrently. See + :class:`~agentscope.app._service.IndexWorker`. + consumer_max_batch (`int`, defaults to ``32``): + Maximum entries the consumer drains per signal. See + :class:`~agentscope.app._service.IndexTaskConsumer`. + parser_executor (`ProcessPoolExecutor | None`, optional): + Process pool for CPU-bound parses. ``None`` runs parses + inline (fine for text-only deployments). + """ + resolved_node_id = ( + node_id or f"{socket.gethostname()}:{uuid.uuid4().hex[:8]}" + ) + + async with AsyncExitStack() as stack: + await stack.enter_async_context(storage) + await stack.enter_async_context(message_bus) + await stack.enter_async_context(blob_store) + await stack.enter_async_context(knowledge_base_manager) + + worker = IndexWorker( + storage=storage, + blob_store=blob_store, + knowledge_base_manager=knowledge_base_manager, + parsers=parsers, + chunker=chunker, + node_id=resolved_node_id, + max_concurrency=worker_max_concurrency, + parser_executor=parser_executor, + ) + await stack.enter_async_context( + IndexTaskConsumer( + message_bus=message_bus, + worker=worker, + max_batch=consumer_max_batch, + ), + ) + + logger.info( + "Index worker %s ready (max_concurrency=%d, max_batch=%d)", + resolved_node_id, + worker_max_concurrency, + consumer_max_batch, + ) + + # Block until a signal arrives. We install handlers on the + # running loop rather than using ``signal.signal`` so the + # interaction with asyncio is well-defined (the default + # handler would raise KeyboardInterrupt at an arbitrary + # await point, which is awkward to clean up). + loop = asyncio.get_running_loop() + stop = loop.create_future() + + def _request_stop() -> None: + if not stop.done(): + stop.set_result(None) + + for sig in (signal.SIGINT, signal.SIGTERM): + try: + loop.add_signal_handler(sig, _request_stop) + except NotImplementedError: + # Windows event loop does not implement add_signal_handler; + # tests and Linux/macOS deployments are unaffected. + pass + + try: + await stop + finally: + logger.info("Index worker %s shutting down", resolved_node_id) diff --git a/src/agentscope/app/rag/index_worker/__main__.py b/src/agentscope/app/rag/index_worker/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..c95ae5294ac814614ff1d8ff8b95211f9b693292 --- /dev/null +++ b/src/agentscope/app/rag/index_worker/__main__.py @@ -0,0 +1,108 @@ +# -*- coding: utf-8 -*- +"""Entry point for ``python -m agentscope.app.rag.index_worker``. + +Resolves a deployment-supplied bootstrap callable from the +``AGENTSCOPE_WORKER_BOOTSTRAP`` environment variable, calls it to +obtain the concrete backends, and hands them to :func:`run_worker`. + +The deployment owns the bootstrap because backend selection is a +deployment concern — there is no one-size-fits-all storage, +message bus, or blob store. The bootstrap must be importable, must +take no arguments, and must return a dict whose keys match +:func:`run_worker`'s keyword arguments. + +Example bootstrap (``mydeploy/worker_bootstrap.py``):: + + import os + + from agentscope.app.rag.blob_store import S3BlobStore + from agentscope.app.rag.knowledge_base_manager import ( + DefaultKnowledgeBaseManager, + ) + from agentscope.app.message_bus import RedisMessageBus + from agentscope.app.storage import RedisStorage + from agentscope.rag import ApproxTokenChunker, TextParser + + def bootstrap() -> dict: + return { + "storage": RedisStorage(url=os.environ["REDIS_URL"]), + "message_bus": RedisMessageBus(url=os.environ["REDIS_URL"]), + "blob_store": S3BlobStore( + bucket=os.environ["S3_BUCKET"], + endpoint_url=os.environ.get("S3_ENDPOINT"), + ), + "knowledge_base_manager": DefaultKnowledgeBaseManager(...), + "parsers": [TextParser()], + "chunker": ApproxTokenChunker(), + } + +And launch:: + + AGENTSCOPE_WORKER_BOOTSTRAP=mydeploy.worker_bootstrap:bootstrap \\ + python -m agentscope.app.rag.index_worker +""" +import asyncio +import importlib +import logging +import os +import sys +from typing import Any, Callable + +from . import run_worker + + +def _resolve(dotted: str) -> Callable[[], dict[str, Any]]: + """Import ``module:attribute`` and return the attribute. + + Args: + dotted (`str`): + A ``module:attribute`` reference; must contain a colon. + + Returns: + `Callable[[], dict[str, Any]]`: + The resolved attribute — expected to be a zero-arg + callable that returns the kwargs dict for + :func:`run_worker`. + + Raises: + `ValueError`: + When ``dotted`` does not contain a colon. + """ + if ":" not in dotted: + raise ValueError( + f"AGENTSCOPE_WORKER_BOOTSTRAP must be in 'module:attr' " + f"form, got {dotted!r}.", + ) + module_name, _, attr = dotted.partition(":") + module = importlib.import_module(module_name) + return getattr(module, attr) + + +def main() -> None: + """Resolve the bootstrap callable from the environment and run the worker. + + Reads ``AGENTSCOPE_WORKER_BOOTSTRAP`` (``module:attr`` form), + imports the target, calls it for the kwargs dict, and forwards + them to :func:`run_worker`. Exits with code ``2`` when the + environment variable is missing — the deployment must supply it + because backend selection is a deployment concern. + """ + logging.basicConfig( + level=os.environ.get("LOG_LEVEL", "INFO"), + format="%(asctime)s %(levelname)s %(name)s %(message)s", + ) + bootstrap_path = os.environ.get("AGENTSCOPE_WORKER_BOOTSTRAP") + if not bootstrap_path: + sys.stderr.write( + "AGENTSCOPE_WORKER_BOOTSTRAP is required — set it to " + "'package.module:callable' that returns a kwargs dict for " + "run_worker.\n", + ) + sys.exit(2) + factory = _resolve(bootstrap_path) + kwargs = factory() + asyncio.run(run_worker(**kwargs)) + + +if __name__ == "__main__": + main() diff --git a/src/agentscope/app/rag/knowledge_base_manager/__init__.py b/src/agentscope/app/rag/knowledge_base_manager/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..76751271d6ada029d9bfebb0c3ad880f80eaa99a --- /dev/null +++ b/src/agentscope/app/rag/knowledge_base_manager/__init__.py @@ -0,0 +1,33 @@ +# -*- coding: utf-8 -*- +"""Knowledge base manager classes. + +The manager owns the lifecycle of knowledge bases: + +- creation / deletion / listing of :class:`KnowledgeBaseRecord` rows, +- allocation / drop of the matching vector store storage, +- construction of :class:`~agentscope.rag.KnowledgeBase` runtime handles + used by both the HTTP service and the agent runtime. + +The MVP ships a single isolation strategy +(:class:`CollectionPerKbManager`); future strategies will live +alongside it. +""" + +from ._base import KnowledgeBaseManagerBase +from ._collection_per_kb import CollectionPerKbManager +from ._dimension_policy import DimensionPolicy, DimensionPolicyKind +from ._errors import ( + DimensionPolicyError, + KnowledgeBaseError, + KnowledgeBaseNotFoundError, +) + +__all__ = [ + "CollectionPerKbManager", + "DimensionPolicy", + "DimensionPolicyError", + "DimensionPolicyKind", + "KnowledgeBaseError", + "KnowledgeBaseManagerBase", + "KnowledgeBaseNotFoundError", +] diff --git a/src/agentscope/app/rag/knowledge_base_manager/_base.py b/src/agentscope/app/rag/knowledge_base_manager/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..0825e1726ccae1412787c0a6f4d76f22f25fafcf --- /dev/null +++ b/src/agentscope/app/rag/knowledge_base_manager/_base.py @@ -0,0 +1,302 @@ +# -*- coding: utf-8 -*- +"""Abstract knowledge base manager. + +The manager is the **lifecycle owner** of knowledge bases: + +- it creates / lists / deletes :class:`KnowledgeBaseRecord` rows in + storage, +- it allocates / drops the matching vector store collections, +- it resolves an embedding model from the record's credential and + hands a ready-to-use :class:`KnowledgeBase` runtime back to callers. + +Different subclasses encode different *isolation strategies*: one +collection per knowledge base, a single shared collection scoped by +metadata, native VDB namespaces, etc. All of them share the same +:class:`KnowledgeBaseManagerBase` interface so the rest of the +application can stay strategy-agnostic. + +The manager is created once at application startup and stored on +``app.state.knowledge_base_manager``. See +:func:`~agentscope.app.create_app` for the wiring. +""" +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Self + +from ._dimension_policy import DimensionPolicy + +if TYPE_CHECKING: + from types import TracebackType + + from ...storage import ( + EmbeddingModelConfig, + KnowledgeBaseRecord, + StorageBase, + ) + from ....rag import KnowledgeBase, VectorStoreBase + + +class KnowledgeBaseManagerBase(ABC): + """Abstract base for knowledge base managers. + + Subclasses implement a specific isolation strategy by overriding + :meth:`create_knowledge_base`, :meth:`delete_knowledge_base`, and + :meth:`get_knowledge`. The bookkeeping methods + (:meth:`get_knowledge_base`, :meth:`list_knowledge_bases`) have + default implementations that delegate to the bound storage. + """ + + def __init__( + self, + storage: "StorageBase", + vector_store: "VectorStoreBase", + ) -> None: + """Initialize the manager. + + Args: + storage (`StorageBase`): + The application-wide storage backend used to persist + :class:`KnowledgeBaseRecord` rows and resolve + credentials. + vector_store (`VectorStoreBase`): + The application-wide vector store instance shared by + every knowledge base allocated by this manager. + """ + self._storage = storage + self._vector_store = vector_store + + # ------------------------------------------------------------------ + # Lifecycle hooks + # ------------------------------------------------------------------ + + async def __aenter__(self) -> Self: + """Enter the manager's lifetime. + + Enters the bound vector store's async context so a single + ``create_app`` parameter (the manager) covers the vector store's + lifecycle too. Subclasses that override this MUST call + ``await super().__aenter__()`` first to keep the vector store + ready before subclass-specific setup runs. + """ + await self._vector_store.__aenter__() + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: "TracebackType | None", + ) -> None: + """Exit the manager's lifetime, releasing the vector store. + + Args: + exc_type (`type[BaseException] | None`): + The exception type raised inside the with-block, if any. + exc (`BaseException | None`): + The exception instance raised inside the with-block, + if any. + tb (`TracebackType | None`): + The traceback for the raised exception, if any. + """ + await self._vector_store.__aexit__(exc_type, exc, tb) + + # ------------------------------------------------------------------ + # Capability discovery + # ------------------------------------------------------------------ + + @abstractmethod + async def get_dimension_policy(self) -> DimensionPolicy: + """Return the embedding-dimension policy this manager enforces. + + Surfaced over HTTP so the front-end can soft-filter + incompatible models / dimensions before submission and show a + helpful banner. + + Returns: + `DimensionPolicy`: + The current dimension policy. + """ + + # ------------------------------------------------------------------ + # CRUD + # ------------------------------------------------------------------ + + @abstractmethod + async def create_knowledge_base( + self, + user_id: str, + name: str, + description: str, + embedding_model_config: "EmbeddingModelConfig", + ) -> "KnowledgeBaseRecord": + """Create a new knowledge base for the given user. + + Implementations must: + + 1. validate ``embedding_model_config.dimensions`` against + :meth:`get_dimension_policy`; + 2. allocate the vector store collection (or namespace) the + strategy uses; + 3. persist a :class:`KnowledgeBaseRecord` and return it. + + Args: + user_id (`str`): + The owner user id. + name (`str`): + Display name. + description (`str`): + Free-form description. + embedding_model_config (`EmbeddingModelConfig`): + Embedding model configuration; pinned to the record. + + Returns: + `KnowledgeBaseRecord`: + The newly persisted record. + + Raises: + `DimensionPolicyError`: + If the requested dimension violates the manager's + dimension policy. + """ + + async def get_knowledge_base( + self, + user_id: str, + knowledge_base_id: str, + ) -> "KnowledgeBaseRecord | None": + """Fetch a knowledge base record by id (delegates to storage). + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The knowledge base id. + + Returns: + `KnowledgeBaseRecord | None`: + The record, or ``None`` if not found / not owned by + the user. + """ + return await self._storage.get_knowledge_base( + user_id, + knowledge_base_id, + ) + + async def list_knowledge_bases( + self, + user_id: str, + ) -> "list[KnowledgeBaseRecord]": + """List all knowledge base records owned by the given user. + + Args: + user_id (`str`): + The owner user id. + + Returns: + `list[KnowledgeBaseRecord]`: + All knowledge base records belonging to the user. + """ + return await self._storage.list_knowledge_bases(user_id) + + async def update_knowledge_base( + self, + user_id: str, + knowledge_base_id: str, + name: str | None = None, + description: str | None = None, + ) -> "KnowledgeBaseRecord | None": + """Update mutable fields on an existing knowledge base record. + + Only ``name`` and ``description`` are mutable. The embedding + model configuration and the underlying collection are pinned + for the lifetime of the record because changing either would + invalidate every previously inserted vector. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The knowledge base id. + name (`str | None`, optional): + New display name; ``None`` leaves the name unchanged. + description (`str | None`, optional): + New description; ``None`` leaves the description + unchanged. + + Returns: + `KnowledgeBaseRecord | None`: + The updated record, or ``None`` if the record was not + found / not owned by the user. + """ + record = await self._storage.get_knowledge_base( + user_id, + knowledge_base_id, + ) + if record is None: + return None + if name is not None: + record.name = name + if description is not None: + record.description = description + return await self._storage.upsert_knowledge_base(user_id, record) + + @abstractmethod + async def delete_knowledge_base( + self, + user_id: str, + knowledge_base_id: str, + ) -> bool: + """Delete a knowledge base record and its underlying storage. + + Implementations must: + + 1. authorise the call by looking the record up first; + 2. drop the vector store collection (or scope) the strategy + uses; + 3. remove the :class:`KnowledgeBaseRecord` from storage. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The id of the knowledge base to delete. + + Returns: + `bool`: + ``True`` if the record existed and was deleted, + ``False`` if it was not found. + """ + + # ------------------------------------------------------------------ + # KnowledgeBase runtime + # ------------------------------------------------------------------ + + @abstractmethod + async def get_knowledge( + self, + user_id: str, + knowledge_base_id: str, + ) -> "KnowledgeBase": + """Resolve a runtime :class:`KnowledgeBase` handle for one KB. + + Implementations are responsible for: + + - looking the record up in storage (authorisation); + - resolving the embedding model from the record's credential; + - constructing the :class:`KnowledgeBase` with the strategy's + collection name and metadata filter. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The knowledge base id. + + Returns: + `KnowledgeBase`: + A runtime handle bound to this knowledge base. + + Raises: + `KnowledgeBaseNotFoundError`: + If the record does not exist or does not belong to + the authenticated user. + """ diff --git a/src/agentscope/app/rag/knowledge_base_manager/_collection_per_kb.py b/src/agentscope/app/rag/knowledge_base_manager/_collection_per_kb.py new file mode 100644 index 0000000000000000000000000000000000000000..689aa5993f240570aff8705309232ab648e5ab9d --- /dev/null +++ b/src/agentscope/app/rag/knowledge_base_manager/_collection_per_kb.py @@ -0,0 +1,188 @@ +# -*- coding: utf-8 -*- +"""Collection-per-knowledge-base isolation strategy. + +The simplest correct implementation: every knowledge base gets its own +vector store collection sized to the chosen embedding model. No +cross-KB co-location, no namespace gymnastics — collection names are +the isolation key. + +Because each knowledge base owns its collection outright, the +dimension policy reported by this manager is always +:attr:`DimensionPolicyKind.ANY`: the user is free to pick any +dimension supported by their embedding model. + +The collection name is generated as ``kb_`` (no user id is +encoded, since storage does the user-scoped authorisation). +""" +from typing import TYPE_CHECKING + +from ._base import KnowledgeBaseManagerBase +from ._dimension_policy import DimensionPolicy, DimensionPolicyKind +from ._errors import KnowledgeBaseNotFoundError +from ...._logging import logger +from ....rag import KnowledgeBase +from ..._service._embedding import get_embedding_model +from ...storage import KnowledgeBaseRecord + +if TYPE_CHECKING: + from ...storage import EmbeddingModelConfig + + +class CollectionPerKbManager(KnowledgeBaseManagerBase): + """One-collection-per-KB knowledge base manager. + + Each knowledge base maps to one collection in the bound + :class:`~agentscope.rag.VectorStoreBase`. The collection is + created at :meth:`create_knowledge_base` time and dropped at + :meth:`delete_knowledge_base` time. + """ + + async def get_dimension_policy(self) -> DimensionPolicy: + """Return :attr:`DimensionPolicyKind.ANY` — every KB picks freely. + + Returns: + `DimensionPolicy`: + The dimension policy ``(ANY, None)``. + """ + return DimensionPolicy( + kind=DimensionPolicyKind.ANY, + dimension=None, + ) + + async def create_knowledge_base( + self, + user_id: str, + name: str, + description: str, + embedding_model_config: "EmbeddingModelConfig", + ) -> KnowledgeBaseRecord: + """Allocate a new collection and persist the knowledge base record. + + Args: + user_id (`str`): + The owner user id. + name (`str`): + Display name. + description (`str`): + Free-form description. + embedding_model_config (`EmbeddingModelConfig`): + Embedding model configuration; pinned to the record. + + Returns: + `KnowledgeBaseRecord`: + The newly persisted record. + """ + record = KnowledgeBaseRecord( + user_id=user_id, + name=name, + description=description, + embedding_model_config=embedding_model_config, + collection_name="", + ) + record.collection_name = f"kb_{record.id}" + + await self._vector_store.create_collection( + name=record.collection_name, + dimensions=embedding_model_config.dimensions, + ) + try: + return await self._storage.upsert_knowledge_base(user_id, record) + except Exception: + # Compensating delete must not shadow the original exception: + # if the delete itself raises, the caller would be misdirected + # to investigate the wrong error and the orphan collection + # would leave no trace. + try: + await self._vector_store.delete_collection( + record.collection_name, + ) + except Exception: # noqa: BLE001 — best-effort cleanup + logger.exception( + "Failed to drop orphan collection %r after " + "upsert_knowledge_base failed; collection leaked.", + record.collection_name, + ) + raise + + async def delete_knowledge_base( + self, + user_id: str, + knowledge_base_id: str, + ) -> bool: + """Drop the underlying collection and remove the record. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The id of the knowledge base to delete. + + Returns: + `bool`: + ``True`` if the record existed and was deleted, + ``False`` if it was not found. + """ + record = await self._storage.get_knowledge_base( + user_id, + knowledge_base_id, + ) + if record is None: + return False + + if await self._vector_store.has_collection(record.collection_name): + await self._vector_store.delete_collection(record.collection_name) + return await self._storage.delete_knowledge_base( + user_id, + knowledge_base_id, + ) + + async def get_knowledge( + self, + user_id: str, + knowledge_base_id: str, + ) -> KnowledgeBase: + """Resolve a :class:`KnowledgeBase` runtime for one knowledge base. + + Looks the record up in storage (raising + :class:`KnowledgeBaseNotFoundError` for unknown / un-owned + ids), constructs the embedding model from the pinned + credential, and returns a ready-to-use handle bound to the + record's collection. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The knowledge base id. + + Returns: + `KnowledgeBase`: + A runtime handle bound to this knowledge base. + + Raises: + `KnowledgeBaseNotFoundError`: + If the record does not exist or does not belong to + the authenticated user. + """ + record = await self._storage.get_knowledge_base( + user_id, + knowledge_base_id, + ) + if record is None: + raise KnowledgeBaseNotFoundError( + f"Knowledge base {knowledge_base_id!r} not found.", + ) + + embedding_model = await get_embedding_model( + user_id=user_id, + config=record.embedding_model_config, + storage=self._storage, + ) + return KnowledgeBase( + name=record.name, + description=record.description, + embedding_model=embedding_model, + vector_store=self._vector_store, + collection=record.collection_name, + metadata_filter=None, + ) diff --git a/src/agentscope/app/rag/knowledge_base_manager/_dimension_policy.py b/src/agentscope/app/rag/knowledge_base_manager/_dimension_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..e1bcd760c0bfb4110f0d834b0b6b48fab68d1192 --- /dev/null +++ b/src/agentscope/app/rag/knowledge_base_manager/_dimension_policy.py @@ -0,0 +1,154 @@ +# -*- coding: utf-8 -*- +"""Dimension policy advertised by a knowledge base manager. + +The policy tells the front-end which embedding-model dimensions are +acceptable when creating a new knowledge base. It is the *capability* +side of the contract; the server still hard-validates every create +call against the same rules. + +Three kinds are modelled: + +- ``ANY`` — any positive dimension is acceptable. The user picks + freely from the list of supported dimensions on the chosen + embedding model card. This is the case for one-collection-per-KB + isolation strategies. +- ``FIXED`` — the dimension is fixed by the server (e.g. a single + shared collection sized to a specific dimension). The front-end + hides incompatible models / dimensions. +- ``LOCKED_BY_EXISTING`` — semantically identical to ``FIXED`` from + the front-end's perspective, but communicates *why* the dimension + is fixed: a previously created knowledge base in the same shared + collection has pinned it. +""" +from enum import Enum +from typing import TYPE_CHECKING + +from pydantic import BaseModel, Field, model_validator + +if TYPE_CHECKING: + from ....embedding import EmbeddingModelCard + + +class DimensionPolicyKind(str, Enum): + """The kind of dimension policy the manager publishes.""" + + ANY = "any" + """Any positive dimension is acceptable. + + Used by isolation strategies that allocate a fresh collection per + knowledge base (each one sized to its own embedding model).""" + + FIXED = "fixed" + """A specific dimension is required by the manager configuration. + + Used by isolation strategies that share a single collection across + knowledge bases — every record must agree on the dimension.""" + + LOCKED_BY_EXISTING = "locked_by_existing" + """A specific dimension was locked in by a previously created + knowledge base. + + Identical to ``FIXED`` for clients; the distinction lets the UI + show *why* the dimension is fixed (e.g. "the first KB pinned the + dimension to 768 — switch isolation strategy to use 1536"). + """ + + +class DimensionPolicy(BaseModel): + """The dimension policy a manager exposes to the front-end.""" + + kind: DimensionPolicyKind = Field( + description="What constraint applies to the chosen dimension.", + ) + """The kind of constraint.""" + + dimension: int | None = Field( + default=None, + description=( + "The required dimension when ``kind`` is ``FIXED`` or " + "``LOCKED_BY_EXISTING``. Always ``None`` for ``ANY``." + ), + ) + """The required dimension, or ``None`` when any dimension is fine.""" + + @model_validator(mode="after") + def _enforce_kind_dimension_invariant(self) -> "DimensionPolicy": + """Reject states like ``ANY + dimension=768`` or ``FIXED + None``. + + Without this guard, downstream code silently produces wrong + results (``ANY`` ignores a stray dimension) or crashes + (``FIXED`` with ``None`` makes ``filter_card`` raise + ``TypeError`` on ``target not in card.supported_dimensions``). + """ + if self.kind is DimensionPolicyKind.ANY: + if self.dimension is not None: + raise ValueError( + "DimensionPolicy: kind=ANY requires dimension=None, " + f"got dimension={self.dimension!r}.", + ) + else: + if self.dimension is None or self.dimension <= 0: + raise ValueError( + f"DimensionPolicy: kind={self.kind.value} requires a " + f"positive dimension, got dimension={self.dimension!r}.", + ) + return self + + def accepts(self, dimensions: int) -> bool: + """Check whether a candidate dimension satisfies this policy. + + Args: + dimensions (`int`): + The candidate output dimension. + + Returns: + `bool`: + ``True`` if the dimension is acceptable. + """ + if self.kind is DimensionPolicyKind.ANY: + return dimensions > 0 + return dimensions == self.dimension + + def filter_card( + self, + card: "EmbeddingModelCard", + ) -> "EmbeddingModelCard | None": + """Project an embedding model card through this policy. + + Used by the knowledge base router to pre-filter the embedding + model catalogue exposed at KB-creation time: + + - ``ANY`` returns the card unchanged. + - ``FIXED`` / ``LOCKED_BY_EXISTING``: + + - fixed-dim cards (``supported_dimensions is None``) are + kept iff their default ``dimensions`` matches the locked + value; + - matryoshka cards are kept iff the locked dimension is in + ``supported_dimensions``, and a copy is returned with + ``supported_dimensions`` narrowed to the locked value and + ``dimensions`` set accordingly. This guarantees the + front-end cannot pick an incompatible dimension. + + Args: + card (`EmbeddingModelCard`): + The candidate embedding model card. + + Returns: + `EmbeddingModelCard | None`: + The (possibly narrowed) card, or ``None`` if the card + cannot satisfy this policy. + """ + if self.kind is DimensionPolicyKind.ANY: + return card + target = self.dimension + if card.supported_dimensions is None: + return card if card.dimensions == target else None + if target not in card.supported_dimensions: + return None + return card.model_copy( + update={ + "dimensions": target, + "supported_dimensions": [target], + }, + ) diff --git a/src/agentscope/app/rag/knowledge_base_manager/_errors.py b/src/agentscope/app/rag/knowledge_base_manager/_errors.py new file mode 100644 index 0000000000000000000000000000000000000000..0ecff422b8c042b4f7e62175c5279f52d4821851 --- /dev/null +++ b/src/agentscope/app/rag/knowledge_base_manager/_errors.py @@ -0,0 +1,59 @@ +# -*- coding: utf-8 -*- +"""Knowledge base manager exception hierarchy. + +The router maps each exception to an HTTP status code in +:mod:`agentscope.app._router._knowledge_base`: + +- :class:`KnowledgeBaseNotFoundError` → ``404`` +- :class:`DimensionPolicyError` → ``409`` + +Keeping the mapping inside the router (and not raising +``HTTPException`` from the manager) lets the same manager be reused by +non-HTTP entry points (e.g. CLI tools, the agent middleware) without +pulling in FastAPI. +""" + + +class KnowledgeBaseError(Exception): + """Base class for knowledge base manager errors.""" + + +class KnowledgeBaseNotFoundError(KnowledgeBaseError): + """Raised when a knowledge base record cannot be located. + + The record is missing entirely, or it exists but does not belong + to the authenticated user. The two are reported identically to + avoid leaking existence of other users' knowledge bases. + """ + + +class DimensionPolicyError(KnowledgeBaseError): + """Raised when a create-time embedding model violates the manager's + dimension policy. + + Carries both the offending dimension and the manager's policy so + the router can surface a precise error message without having to + re-derive either side. + """ + + def __init__( + self, + message: str, + *, + requested_dimension: int, + policy_dimension: int | None, + ) -> None: + """Initialize the error. + + Args: + message (`str`): + Human-readable error message. + requested_dimension (`int`): + The dimension the caller asked for. + policy_dimension (`int | None`): + The dimension the policy enforces, or ``None`` if the + policy admits any dimension. + """ + super().__init__(message) + self.requested_dimension = requested_dimension + self.policy_dimension = policy_dimension diff --git a/src/agentscope/app/storage/__init__.py b/src/agentscope/app/storage/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c8448a795706bbc8b549e1b88f9daacb7754e320 --- /dev/null +++ b/src/agentscope/app/storage/__init__.py @@ -0,0 +1,53 @@ +# -*- coding: utf-8 -*- +"""The storage module in agentscope.""" + +from ._base import StorageBase +from ._redis_storage import RedisStorage +from ._model import ( + AgentData, + AgentRecord, + CredentialRecord, + KnowledgeBaseRecord, + KnowledgeDocumentData, + KnowledgeDocumentRecord, + KnowledgeDocumentStatus, + ScheduleData, + ScheduleRecord, + ScheduleSource, + SessionConfig, + SessionKnowledgeConfig, + SessionRecord, + SessionSource, + ChatModelConfig, + TTSModelConfig, + EmbeddingModelConfig, + TeamData, + TeamRecord, + UserRecord, +) + +__all__ = [ + "StorageBase", + "RedisStorage", + # The ORM models + "AgentData", + "AgentRecord", + "CredentialRecord", + "KnowledgeBaseRecord", + "KnowledgeDocumentData", + "KnowledgeDocumentRecord", + "KnowledgeDocumentStatus", + "SessionConfig", + "SessionKnowledgeConfig", + "SessionRecord", + "SessionSource", + "ChatModelConfig", + "TTSModelConfig", + "EmbeddingModelConfig", + "TeamData", + "TeamRecord", + "UserRecord", + "ScheduleData", + "ScheduleRecord", + "ScheduleSource", +] diff --git a/src/agentscope/app/storage/_base.py b/src/agentscope/app/storage/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..aee49cb59e29682da9a0e12cda2551901312490e --- /dev/null +++ b/src/agentscope/app/storage/_base.py @@ -0,0 +1,908 @@ +# -*- coding: utf-8 -*- +# pylint: disable=too-many-public-methods +"""The storage base class.""" +from abc import ABC, abstractmethod +from datetime import datetime, timedelta +from typing import Any, Self + + +from ._model import ( + AgentRecord, + CredentialRecord, + KnowledgeBaseRecord, + KnowledgeDocumentRecord, + KnowledgeDocumentStatus, + ScheduleRecord, + SessionRecord, + SessionConfig, + SessionSource, + TeamRecord, +) +from ...credential import CredentialBase +from ...message import Msg +from ...state import AgentState + + +class StorageBase(ABC): + """The storage abstract base class.""" + + async def __aenter__(self) -> Self: + """Start the storage backend (open connection pool, etc.).""" + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: Any, + ) -> None: + """Shut down the storage backend.""" + await self.aclose() + + async def aclose(self) -> None: + """Release underlying connection resources. Default is a no-op.""" + + @abstractmethod + async def upsert_credential( + self, + user_id: str, + credential_data: CredentialBase, + ) -> str: + """Create or update a credential in the storage. + + Args: + user_id (`str`): + The user id. + credential_data (`CredentialBase`): + The credential data. + + Returns: + `str`: + The credential id. + """ + + @abstractmethod + async def list_credentials(self, user_id: str) -> list[CredentialRecord]: + """List all credentials for a given user. + + Args: + user_id (`str`): + The user id. + + Returns: + `list[CredentialRecord]`: + List of all credentials for a given user. + """ + + @abstractmethod + async def get_credential( + self, + user_id: str, + credential_id: str, + ) -> CredentialRecord | None: + """Fetch a single credential record by id. + + Args: + user_id (`str`): The owner user id. + credential_id (`str`): The credential id. + + Returns: + `CredentialRecord | None`: The record, or ``None`` if not found. + """ + + @abstractmethod + async def delete_credential( + self, + user_id: str, + credential_id: str, + ) -> bool: + """Delete a credential. + + Args: + user_id (`str`): + The user id. + credential_id (`str`): + The credential id. + + Returns: + `bool`: + True if deleted, False if not found. + """ + + @abstractmethod + async def upsert_agent( + self, + user_id: str, + agent_record: AgentRecord, + ) -> str: + """Create an agent record in the storage. + + Args: + user_id (`str`): + The user id. + agent_record (`AgentRecord`): + The agent record. + + Returns: + `str`: + The agent id. + """ + + @abstractmethod + async def list_agents(self, user_id: str) -> list[AgentRecord]: + """List all agents for a given user. + + Args: + user_id (`str`): + The user id. + + Returns: + `list[AgentRecord]`: + List of all agents for a given user. + """ + + @abstractmethod + async def get_agent( + self, + user_id: str, + agent_id: str, + ) -> AgentRecord | None: + """Fetch a single agent record by id. + + Args: + user_id (`str`): The owner user id. + agent_id (`str`): The agent id. + + Returns: + `AgentRecord | None`: The record, or ``None`` if not found. + """ + + @abstractmethod + async def delete_agent(self, user_id: str, agent_id: str) -> bool: + """Delete an agent record. + + Args: + user_id (`str`): + The user id. + agent_id (`str`): + The agent id. + + Returns: + `bool`: + True if deleted, False if not found. + """ + + @abstractmethod + async def upsert_session( + self, + user_id: str, + agent_id: str, + config: SessionConfig, + state: AgentState | None = None, + session_id: str | None = None, + source: SessionSource = SessionSource.USER, + source_schedule_id: str | None = None, + ) -> SessionRecord: + """Create or update a session for a (user, agent) pair. + + Args: + user_id (`str`): The owner user id. + agent_id (`str`): The agent id. + config (`SessionConfig`): Immutable session configuration + (model, workspace). Required on create; passed unchanged on + state-only updates. + state (`AgentState | None`, optional): Runtime state to persist. + Defaults to a fresh ``AgentState()`` when ``None``. + session_id (`str | None`, optional): If provided, update the + existing session with this id. If ``None``, create a new + session. + source (`SessionSource`, optional): The source that created this + session. Defaults to ``SessionSource.USER``. + source_schedule_id (`str | None`, optional): The schedule that + created this session. When set, the session is indexed under + the schedule for execution history queries. + + Returns: + `SessionRecord`: The created or updated record. + """ + + @abstractmethod + async def set_session_team_id( + self, + user_id: str, + session_id: str, + team_id: str | None, + ) -> None: + """Set or clear ``team_id`` on an existing session record. + + Bypasses :meth:`upsert_session` because that method does not + write ``team_id``. Used by team operations (create/dissolve/ + leave) to keep the leader/worker → team relationship consistent. + Idempotent: a no-op if the session does not exist or already + holds the given value. + + Args: + user_id (`str`): + The owner user id. + session_id (`str`): + The session whose ``team_id`` should be updated. + team_id (`str | None`): + The new value. ``None`` detaches the session from any + team. + """ + + @abstractmethod + async def update_session_state( + self, + user_id: str, + agent_id: str, + session_id: str, + state: AgentState, + ) -> None: + """Update only the mutable state of an existing session. + + Convenience method for the hot path (post-chat-turn persistence). + Raises ``KeyError`` if the session does not exist. + + Args: + user_id (`str`): The owner user id. + agent_id (`str`): The agent id. + session_id (`str`): The session id. + state (`AgentState`): The new agent state to persist. + """ + + @abstractmethod + async def list_sessions( + self, + user_id: str, + agent_id: str, + ) -> list[SessionRecord]: + """List all sessions for a given user and agent entity. + + Args: + user_id (`str`): The user id. + agent_id (`str`): The agent id. + + Returns: + `list[SessionRecord]`: List of all sessions for the (user, agent). + """ + + @abstractmethod + async def delete_session( + self, + user_id: str, + agent_id: str, + session_id: str, + ) -> bool: + """Delete a session. + + Args: + user_id (`str`): The user id. + agent_id (`str`): The agent id. + session_id (`str`): The session id. + + Returns: + `bool`: True if deleted, False if not found. + """ + + @abstractmethod + async def get_session( + self, + user_id: str, + agent_id: str, + session_id: str, + ) -> SessionRecord | None: + """Fetch a single session record by id. + + Args: + user_id (`str`): The owner user id. + agent_id (`str`): The agent id. + session_id (`str`): The session id. + + Returns: + `SessionRecord | None`: The record, or ``None`` if not found. + """ + + @abstractmethod + async def list_sessions_by_schedule( + self, + user_id: str, + schedule_id: str, + ) -> list[SessionRecord]: + """Return all sessions created by a given schedule. + + Args: + user_id (`str`): The owner user id. + schedule_id (`str`): The schedule id. + + Returns: + `list[SessionRecord]`: Sessions triggered by this schedule, + ordered by creation time (newest first). + """ + + @abstractmethod + async def upsert_schedule( + self, + user_id: str, + record: ScheduleRecord, + ) -> str: + """Persist a cron task record and register it in the user's index. + + Args: + user_id (`str`): The owner user id. + record (`ScheduleRecord`): The fully-populated record to store. + + Returns: + `str`: The id of the stored record. + """ + + @abstractmethod + async def get_schedule( + self, + user_id: str, + schedule_id: str, + ) -> ScheduleRecord | None: + """Fetch a single cron task record by id. + + Args: + user_id (`str`): The owner user id. + schedule_id (`str`): The task id. + + Returns: + `ScheduleRecord | None`: The record, or ``None`` if not found. + """ + + @abstractmethod + async def list_schedules( + self, + user_id: str, + ) -> list[ScheduleRecord]: + """Return all cron task records belonging to the given user. + + Args: + user_id (`str`): The owner user id. + + Returns: + `list[ScheduleRecord]`: All cron task records for the user. + """ + + @abstractmethod + async def delete_schedule( + self, + user_id: str, + schedule_id: str, + ) -> bool: + """Delete a cron task record and remove it from the user's index. + + Args: + user_id (`str`): The owner user id. + schedule_id (`str`): The id of the task to delete. + + Returns: + `bool`: ``True`` if deleted, ``False`` if not found. + """ + + @abstractmethod + async def list_all_schedules(self) -> list[ScheduleRecord]: + """Return every schedule record across all users. + + Used on startup to restore the in-memory scheduler from persisted + state. Normal per-user listing should use :meth:`list_schedules`. + + Returns: + `list[ScheduleRecord]`: All schedule records in the store. + """ + + # ------------------------------------------------------------------ + # Message persistence + # ------------------------------------------------------------------ + + @abstractmethod + async def upsert_message( + self, + user_id: str, + session_id: str, + msg: Msg, + ) -> None: + """Persist a message to the session's message list. + + If the last message in the list has the same ``id`` as *msg*, it is + replaced (merge/overwrite for the same reply_id across continuation + calls). Otherwise, *msg* is appended as a new entry. + + Args: + user_id (`str`): The owner user id. + session_id (`str`): The session id. + msg (`Msg`): The message to persist. + """ + + @abstractmethod + async def get_message( + self, + user_id: str, + session_id: str, + message_id: str, + ) -> Msg | None: + """Fetch a single message by id from the session's message list. + + Args: + user_id (`str`): The owner user id. + session_id (`str`): The session id. + message_id (`str`): The message id to look up. + + Returns: + `Msg | None`: The message, or ``None`` if not found. + """ + + @abstractmethod + async def list_messages( + self, + user_id: str, + session_id: str, + offset: int = 0, + limit: int = 50, + ) -> list[Msg]: + """Return messages for a session with pagination. + + Args: + user_id (`str`): The owner user id. + session_id (`str`): The session id. + offset (`int`): Starting index (0-based). Defaults to 0. + limit (`int`): Maximum number of messages to return. + + Returns: + `list[Msg]`: Messages in chronological order. + """ + + # ------------------------------------------------------------------ + # Team persistence + # ------------------------------------------------------------------ + + @abstractmethod + async def upsert_team( + self, + user_id: str, + record: TeamRecord, + ) -> TeamRecord: + """Create or update a team record. + + Args: + user_id (`str`): The owner user id. + record (`TeamRecord`): The team record to persist. The record's + ``id`` is used as the primary key; if a record with the same + id already exists it is overwritten. + + Returns: + `TeamRecord`: The stored record (with ``updated_at`` refreshed). + """ + + @abstractmethod + async def get_team( + self, + user_id: str, + team_id: str, + ) -> TeamRecord | None: + """Fetch a single team record by id. + + Args: + user_id (`str`): The owner user id. + team_id (`str`): The team id. + + Returns: + `TeamRecord | None`: The record, or ``None`` if not found. + """ + + @abstractmethod + async def list_teams(self, user_id: str) -> list[TeamRecord]: + """List all teams owned by a given user. + + Args: + user_id (`str`): The user id. + + Returns: + `list[TeamRecord]`: All team records belonging to the user. + """ + + @abstractmethod + async def delete_team(self, user_id: str, team_id: str) -> bool: + """Delete a team record and cascade-delete all of its workers. + + The cascade mirrors SQL's ``ON DELETE CASCADE`` semantics: + + 1. For each ``member_id`` in :attr:`TeamData.member_ids`, call + :meth:`delete_agent` (which cascades that worker's session). + 2. Clear ``team_id`` on the leader session referenced by + :attr:`TeamRecord.session_id` (``ON DELETE SET NULL`` for the + leader's back-reference to the team). Idempotent if the + session has already been deleted. + 3. Delete the :class:`TeamRecord` key and the per-user team + index entry. + + Args: + user_id (`str`): + The owner user id. + team_id (`str`): + The id of the team to delete. + + Returns: + `bool`: + ``True`` if the team record existed and was deleted, + ``False`` if not found. + """ + + # ------------------------------------------------------------------ + # Knowledge base persistence + # ------------------------------------------------------------------ + + @abstractmethod + async def upsert_knowledge_base( + self, + user_id: str, + record: KnowledgeBaseRecord, + ) -> KnowledgeBaseRecord: + """Create or update a knowledge base record. + + The caller is responsible for constructing the full + :class:`KnowledgeBaseRecord` (including ``id`` and + ``collection_name``). If a record with the same id already + exists it is overwritten and ``updated_at`` refreshed. + + Args: + user_id (`str`): + The owner user id. Must match ``record.user_id``. + record (`KnowledgeBaseRecord`): + The fully-populated record to persist. + + Returns: + `KnowledgeBaseRecord`: + The stored record (with ``updated_at`` refreshed). + """ + + @abstractmethod + async def get_knowledge_base( + self, + user_id: str, + knowledge_base_id: str, + ) -> KnowledgeBaseRecord | None: + """Fetch a single knowledge base record by id. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The knowledge base id. + + Returns: + `KnowledgeBaseRecord | None`: + The record, or ``None`` if not found or not owned by + the given user. + """ + + @abstractmethod + async def list_knowledge_bases( + self, + user_id: str, + ) -> list[KnowledgeBaseRecord]: + """List all knowledge base records owned by the given user. + + Args: + user_id (`str`): + The owner user id. + + Returns: + `list[KnowledgeBaseRecord]`: + All knowledge base records belonging to the user. + """ + + @abstractmethod + async def delete_knowledge_base( + self, + user_id: str, + knowledge_base_id: str, + ) -> bool: + """Delete a knowledge base record and remove it from the user index. + + Note: this only removes the metadata record; deletion of the + underlying vector store collection is the caller's + responsibility (typically the knowledge base manager). + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The id of the record to delete. + + Returns: + `bool`: + ``True`` if the record existed and was deleted, + ``False`` if not found. + """ + + # ------------------------------------------------------------------ + # Knowledge document persistence + # ------------------------------------------------------------------ + + @abstractmethod + async def upsert_knowledge_document( + self, + user_id: str, + record: KnowledgeDocumentRecord, + ) -> KnowledgeDocumentRecord: + """Create or update a knowledge document record. + + Used by the upload endpoint to register a freshly arrived + document (``status='pending'``) and by other code paths that + need to overwrite the full record. Phase transitions during + indexing should go through :meth:`update_knowledge_document_status` + instead, which is cheaper and atomic w.r.t. the lease fields. + + Args: + user_id (`str`): + The owner user id. Must match ``record.user_id``. + record (`KnowledgeDocumentRecord`): + The fully-populated record to persist. + + Returns: + `KnowledgeDocumentRecord`: + The stored record (with ``updated_at`` refreshed). + """ + + @abstractmethod + async def get_knowledge_document( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> KnowledgeDocumentRecord | None: + """Fetch a single knowledge document record by id. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The parent knowledge base id. + document_id (`str`): + The document id. + + Returns: + `KnowledgeDocumentRecord | None`: + The record, or ``None`` if not found. + """ + + @abstractmethod + async def list_knowledge_documents( + self, + user_id: str, + knowledge_base_id: str, + ) -> list[KnowledgeDocumentRecord]: + """List all documents in a knowledge base. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The parent knowledge base id. + + Returns: + `list[KnowledgeDocumentRecord]`: + All document records belonging to the knowledge base, + in arbitrary order. + """ + + @abstractmethod + async def delete_knowledge_document( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> bool: + """Delete a knowledge document record. + + Only removes the metadata record; cleanup of the underlying + blob and vector store records is the caller's responsibility. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The parent knowledge base id. + document_id (`str`): + The id of the record to delete. + + Returns: + `bool`: + ``True`` if the record existed and was deleted, + ``False`` if not found. + """ + + @abstractmethod + async def update_knowledge_document_status( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + status: KnowledgeDocumentStatus, + error: str | None = None, + chunk_count: int | None = None, + ) -> None: + """Update only the status-related fields of a document record. + + Used by the indexing worker as it walks the lifecycle + transitions (``parsing`` → ``chunking`` → ``indexing`` → + ``ready`` / ``error``). Cheaper than a full upsert and avoids + races with concurrent lease writes by touching only the status + / error / chunk_count fields. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The parent knowledge base id. + document_id (`str`): + The document being updated. + status (`KnowledgeDocumentStatus`): + The new lifecycle state. + error (`str | None`, optional): + Failure reason, set when ``status == 'error'``. + Ignored otherwise. ``None`` leaves the existing + value unchanged. + chunk_count (`int | None`, optional): + The final chunk count, set when ``status == 'ready'``. + ``None`` leaves the existing value unchanged. + """ + + @abstractmethod + async def acquire_knowledge_document_lease( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + processing_node: str, + lease_ttl: timedelta, + now: datetime | None = None, + ) -> bool: + """Compare-and-swap acquisition of the document processing lease. + + Succeeds only if no other worker currently holds a live lease — + i.e. ``processing_node`` is unset or ``lease_expires_at`` is in + the past relative to ``now``. On success the record's + ``processing_node`` and ``data.lease_expires_at`` fields are + updated. The CAS is what makes the sweeper safe to run on + multiple nodes at once: even if two sweepers redispatch the + same document, only one worker can acquire its lease. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The parent knowledge base id. + document_id (`str`): + The document whose lease is being acquired. + processing_node (`str`): + Stable identifier of the calling worker (e.g. + ``hostname:pid:uuid``). Persisted as + ``processing_node`` on success. + lease_ttl (`timedelta`): + How long the lease should live from ``now``. + now (`datetime | None`, optional): + Reference time for the comparison. Defaults to + ``datetime.now()``. Injectable for testing. + + Returns: + `bool`: + ``True`` when the lease was acquired by this caller, + ``False`` when another worker already holds it. + """ + + @abstractmethod + async def renew_knowledge_document_lease( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + processing_node: str, + lease_ttl: timedelta, + now: datetime | None = None, + ) -> bool: + """Extend an existing lease this worker already holds. + + Required for long-running parses that exceed ``lease_ttl`` — + the worker calls this periodically so the sweeper does not + mistake it for a crash. Updates ``lease_expires_at`` to + ``now + lease_ttl`` only when ``processing_node`` matches the + caller; otherwise returns ``False`` so the worker can abandon + cleanly (its lease was stolen). + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The parent knowledge base id. + document_id (`str`): + The document being renewed. + processing_node (`str`): + The caller's processing node id; must match the + record's current ``processing_node`` for renewal to + succeed. + lease_ttl (`timedelta`): + The new lease duration relative to ``now``. + now (`datetime | None`, optional): + Reference time. Defaults to ``datetime.now()``. + + Returns: + `bool`: + ``True`` when the renewal succeeded, ``False`` when + the lease no longer belongs to the caller. + """ + + @abstractmethod + async def release_knowledge_document_lease( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + processing_node: str, + ) -> None: + """Release the processing lease this worker holds. + + Clears ``processing_node`` and ``data.lease_expires_at`` only + if the current holder matches ``processing_node`` — a stolen + lease (e.g. after sweep) is left untouched. Idempotent: a + no-op when the document is missing or already free. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The parent knowledge base id. + document_id (`str`): + The document whose lease is being released. + processing_node (`str`): + The caller's processing node id; must match the + record's current holder. + """ + + @abstractmethod + async def list_knowledge_documents_with_expired_lease( + self, + now: datetime | None = None, + ) -> list[KnowledgeDocumentRecord]: + """Return non-terminal documents whose lease has expired. + + Scans every user / knowledge base — used by the sweeper, not + by user-facing endpoints. A document is "expired" when + ``data.status`` is not terminal (``ready`` / ``error``), + ``processing_node`` is set, and ``data.lease_expires_at`` is + in the past. Implementations are free to skip records that + match no work and return an unspecified order. + + Args: + now (`datetime | None`, optional): + Reference time. Defaults to ``datetime.now()``. + + Returns: + `list[KnowledgeDocumentRecord]`: + Documents to redispatch. + """ + + @abstractmethod + async def list_knowledge_documents_pending_since( + self, + threshold: datetime, + ) -> list[KnowledgeDocumentRecord]: + """Return documents stuck in ``pending`` older than ``threshold``. + + Catches the corner case where the upload endpoint persisted a + record but the dispatcher (or the process holding it) died + before any worker picked the document up — a crashed lease + sweep would miss it because no lease was ever written. + + Args: + threshold (`datetime`): + Cut-off creation time; only ``pending`` records + ``created_at < threshold`` are returned. + + Returns: + `list[KnowledgeDocumentRecord]`: + Orphan ``pending`` documents to redispatch. + """ diff --git a/src/agentscope/app/storage/_model/__init__.py b/src/agentscope/app/storage/_model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f90e58c1baaef5f0a0fbfcfed762d2cb5efa19ef --- /dev/null +++ b/src/agentscope/app/storage/_model/__init__.py @@ -0,0 +1,46 @@ +# -*- coding: utf-8 -*- +"""Storage models for persisted resources.""" + +from ._agent import AgentRecord, AgentData +from ._credential import CredentialRecord +from ._knowledge_base import KnowledgeBaseRecord +from ._knowledge_document import ( + KnowledgeDocumentData, + KnowledgeDocumentRecord, + KnowledgeDocumentStatus, +) +from ._schedule import ScheduleData, ScheduleRecord, ScheduleSource +from ._session import ( + SessionRecord, + SessionConfig, + SessionKnowledgeConfig, + ChatModelConfig, + TTSModelConfig, + EmbeddingModelConfig, + SessionSource, +) +from ._team import TeamRecord, TeamData +from ._user import UserRecord + +__all__ = [ + "AgentData", + "AgentRecord", + "CredentialRecord", + "KnowledgeBaseRecord", + "KnowledgeDocumentData", + "KnowledgeDocumentRecord", + "KnowledgeDocumentStatus", + "ScheduleData", + "ScheduleRecord", + "ScheduleSource", + "SessionConfig", + "SessionKnowledgeConfig", + "SessionRecord", + "SessionSource", + "ChatModelConfig", + "TTSModelConfig", + "EmbeddingModelConfig", + "TeamData", + "TeamRecord", + "UserRecord", +] diff --git a/src/agentscope/app/storage/_model/_agent.py b/src/agentscope/app/storage/_model/_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..ede8d10f767d033833813bccf05412516f8e3d5c --- /dev/null +++ b/src/agentscope/app/storage/_model/_agent.py @@ -0,0 +1,64 @@ +# -*- coding: utf-8 -*- +"""The agent storage class.""" +from typing import Literal + +from pydantic import Field, BaseModel + +from ...._utils._common import _generate_id +from ._base import _RecordBase +from ....agent import ContextConfig, ReActConfig + + +class AgentData(BaseModel): + """The agent data model.""" + + id: str = Field( + description="Unique agent id", + default_factory=_generate_id, + ) + """The agent id.""" + + name: str = Field( + description="The name of the agent.", + title="Name", + ) + + system_prompt: str = Field( + default="You're a helpful assistant.", + description="The system prompt for the agent.", + title="System Prompt", + # Hint for schema-driven UI renderers; see ``ContextConfig`` for + # the same pattern on long-form prompts. + json_schema_extra={"format": "textarea"}, + ) + + context_config: ContextConfig = Field( + description="The context config for the agent.", + title="Context Config", + ) + + react_config: ReActConfig = Field( + description="The react config for the agent.", + title="React Config", + ) + + +class AgentRecord(_RecordBase): + """The agent ORM model.""" + + user_id: str + """The user id""" + + source: Literal["user", "team"] = "user" + """How this agent was created. + + - ``"user"``: created directly by the user (default). Can have multiple + sessions and is listed in the user's regular agent list. + - ``"team"``: spawned as a team worker by another agent's + ``create_team`` / ``team_add_member`` tool. Has exactly one session. + Team membership itself is session-level and stored on + :class:`SessionRecord.team_id`. + """ + + data: AgentData + """The agent data""" diff --git a/src/agentscope/app/storage/_model/_base.py b/src/agentscope/app/storage/_model/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..92cf21a27713841880d3b5bfd6b756ea0ce3022e --- /dev/null +++ b/src/agentscope/app/storage/_model/_base.py @@ -0,0 +1,26 @@ +# -*- coding: utf-8 -*- +"""The base attributes used in storage.""" +from datetime import datetime + +from pydantic import BaseModel, Field + +from ...._utils._common import _generate_id + + +class _RecordBase(BaseModel): + """The base class for all records.""" + + id: str = Field( + default_factory=_generate_id, + description="Unique identifier for the credential.", + ) + + updated_at: datetime = Field( + default_factory=datetime.now, + ) + """The updated time.""" + + created_at: datetime = Field( + default_factory=datetime.now, + ) + """The created time.""" diff --git a/src/agentscope/app/storage/_model/_credential.py b/src/agentscope/app/storage/_model/_credential.py new file mode 100644 index 0000000000000000000000000000000000000000..a7cb5fafdc2b502c512a6d897673dbbb22b27cf0 --- /dev/null +++ b/src/agentscope/app/storage/_model/_credential.py @@ -0,0 +1,17 @@ +# -*- coding: utf-8 -*- +"""The credential record.""" +from pydantic import Field + +from ...._utils._common import _generate_id +from ._base import _RecordBase + + +class CredentialRecord(_RecordBase): + """The credential model used for storing credentials.""" + + user_id: str = Field( + default_factory=_generate_id, + ) + + data: dict + """The credential data.""" diff --git a/src/agentscope/app/storage/_model/_knowledge_base.py b/src/agentscope/app/storage/_model/_knowledge_base.py new file mode 100644 index 0000000000000000000000000000000000000000..d81d95160323e087057bd146838f8ba7c6e81a55 --- /dev/null +++ b/src/agentscope/app/storage/_model/_knowledge_base.py @@ -0,0 +1,46 @@ +# -*- coding: utf-8 -*- +"""The knowledge base record.""" +from pydantic import Field + +from ._base import _RecordBase +from ._session import EmbeddingModelConfig + + +class KnowledgeBaseRecord(_RecordBase): + """A persisted knowledge base record. + + Stores per-user metadata for a knowledge base; the actual chunks + and vectors live in the configured ``VectorStoreBase`` backend. + Each record is the canonical authorisation gate: HTTP handlers and + middleware look the record up by ``(user_id, id)`` before talking + to the vector store. + """ + + user_id: str = Field(description="The owner user id.") + """The user id that owns this knowledge base.""" + + name: str = Field(description="Display name of the knowledge base.") + """Display name shown in the UI.""" + + description: str = Field( + default="", + description="Free-form description of the knowledge base purpose.", + ) + """Free-form description shown in the UI.""" + + embedding_model_config: EmbeddingModelConfig = Field( + description=( + "Embedding model configuration pinned at creation time. " + "Cannot change for the lifetime of the record because the " + "underlying collection is sized to its dimension." + ), + ) + """Embedding model configuration pinned at creation time.""" + + collection_name: str = Field( + description=( + "The vector store collection that physically backs this " + "knowledge base. Generated server-side; opaque to clients." + ), + ) + """The vector store collection name (e.g. ``kb_``).""" diff --git a/src/agentscope/app/storage/_model/_knowledge_document.py b/src/agentscope/app/storage/_model/_knowledge_document.py new file mode 100644 index 0000000000000000000000000000000000000000..7c202f9808ddd2921615901f18ce88fa4d0a6b15 --- /dev/null +++ b/src/agentscope/app/storage/_model/_knowledge_document.py @@ -0,0 +1,139 @@ +# -*- coding: utf-8 -*- +"""The knowledge document record. + +A :class:`KnowledgeDocumentRecord` is the canonical source of truth +for one uploaded file inside a knowledge base. It owns the document's +**lifecycle** (status, error, lease) and **byte handle** (``blob_uri``) +before any chunks reach the vector store, which is exactly the state +the vector store cannot represent on its own (no chunks means no +``document_id`` to aggregate from). + +Top-level fields are the relational keys (``user_id`` / +``knowledge_base_id`` / ``processing_node``) that the storage backend +indexes on; everything else lives inside ``data`` per the project +record convention. +""" +from datetime import datetime +from typing import Literal + +from pydantic import BaseModel, Field + +from ._base import _RecordBase + + +KnowledgeDocumentStatus = Literal[ + "pending", + "parsing", + "chunking", + "indexing", + "ready", + "error", +] +# The six lifecycle states of a knowledge document. +# +# ``pending`` — bytes are in the blob store, waiting for a worker to +# pick the document up. ``parsing`` / ``chunking`` / ``indexing`` are +# worker-owned transitions; ``ready`` and ``error`` are terminal. + + +class KnowledgeDocumentData(BaseModel): + """The mutable payload of a knowledge document record.""" + + filename: str = Field( + description="Original filename supplied by the uploader.", + ) + """The original filename — used both for citation and as the + ``source`` field on every chunk produced from this document.""" + + size: int = Field( + ge=0, + description="Document size in bytes as observed at upload time.", + ) + """Byte length recorded at upload time. Used by quota checks and + the UI; not authoritative for parsing (the worker reopens the + blob).""" + + content_type: str | None = Field( + default=None, + description="IANA media type used to route the upload to a parser.", + ) + """IANA media type. ``None`` lets the worker fall back to + ``mimetypes.guess_type(filename)``.""" + + blob_uri: str = Field( + description=( + "URI returned by the blob store after the upload was " + "streamed in (e.g. ``local://kb/.../uuid``)." + ), + ) + """Scheme-qualified URI handed back by + :class:`~agentscope.app.rag.blob_store.BlobStoreBase`. The worker + streams bytes back through the same blob store.""" + + status: KnowledgeDocumentStatus = Field( + default="pending", + description="Current lifecycle state.", + ) + """Current lifecycle state. Read by the polling endpoint, written + by the worker as it transitions phases.""" + + error: str | None = Field( + default=None, + description=( + "Human-readable failure reason when ``status == 'error'``. " + "MUST NOT include stack traces or filesystem paths — it is " + "rendered verbatim in the UI." + ), + ) + """Human-readable failure reason. Surfaced directly to the user; + keep it short and free of sensitive content.""" + + chunk_count: int = Field( + default=0, + ge=0, + description="Number of chunks successfully indexed.", + ) + """The final chunk count, written by the worker when the document + reaches ``ready``.""" + + lease_expires_at: datetime | None = Field( + default=None, + description=( + "Wall-clock deadline for the worker that currently holds " + "this document. ``None`` if no worker is processing it." + ), + ) + """Lease deadline. Together with ``processing_node`` lets the + sweeper detect crashed workers and reassign their documents.""" + + +class KnowledgeDocumentRecord(_RecordBase): + """A persisted knowledge document record. + + Top-level fields are relational keys the storage backend needs to + index on directly (per-user listing, per-KB listing, per-node lease + sweeps); the mutable payload lives in :attr:`data`. + """ + + user_id: str = Field(description="The owner user id.") + """The user id that owns the parent knowledge base.""" + + knowledge_base_id: str = Field( + description="The id of the knowledge base this document belongs to.", + ) + """The knowledge base the document is being indexed into.""" + + processing_node: str | None = Field( + default=None, + description=( + "Identifier of the worker process that currently holds the " + "lease on this document. ``None`` if no worker is " + "processing it." + ), + ) + """The current lease holder. Promoted to the top level so the + storage backend can look up "documents owned by node X" or + "documents with no owner" without deserialising every payload.""" + + data: KnowledgeDocumentData + """The mutable document payload.""" diff --git a/src/agentscope/app/storage/_model/_schedule.py b/src/agentscope/app/storage/_model/_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..1a78d0b904e213afc7e949a4dc0068c682b6092c --- /dev/null +++ b/src/agentscope/app/storage/_model/_schedule.py @@ -0,0 +1,117 @@ +# -*- coding: utf-8 -*- +"""The schedule storage model.""" +from datetime import datetime +from enum import Enum + +from pydantic import BaseModel, Field + +from ._base import _RecordBase +from ._session import ChatModelConfig +from ....permission import PermissionMode + + +class ScheduleSource(str, Enum): + """The source that created the schedule. + + Attributes: + USER: Created manually by the user via the UI. + AGENT: Created automatically by an agent, e.g. via a tool call. + """ + + USER = "USER" + AGENT = "AGENT" + + +def _get_local_timezone() -> str: + """Get the local timezone. + + Returns: + `str`: + The local timezone. + """ + try: + from tzlocal import get_localzone + + return str(get_localzone()) + except Exception: + return "UTC" + + +class ScheduleData(BaseModel): + """The schedule configuration data.""" + + name: str = Field(description="Display name of the schedule.") + + description: str = Field( + default="", + description="The description of the schedule, including its purpose, " + "trigger conditions, etc.", + ) + + enabled: bool = Field( + default=True, + description="Whether the schedule is active. Disabled schedules are " + "retained but will not trigger.", + ) + + timezone: str = Field( + default=_get_local_timezone(), + description="IANA timezone name used to evaluate the cron expression, " + "e.g. 'America/New_York' or 'Asia/Shanghai'.", + ) + + cron_expression: str = Field( + description="Standard 5-field cron expression, e.g. '0 9 * * 1-5'.", + ) + + started_at: datetime = Field( + description="The date and time the schedule was started.", + default_factory=datetime.now, + ) + + ended_at: datetime | None = Field( + default=None, + description="The date and time the schedule was ended.", + ) + + chat_model_config: ChatModelConfig = Field( + description="Model configuration for the auto-created session.", + ) + + stateful: bool = Field( + title="Stateful", + default=False, + description="Whether consecutive executions share the same session " + "context. If not, each execution will have its own state.", + ) + + permission_mode: PermissionMode = Field( + title="Permission mode", + default=PermissionMode.DONT_ASK, + description="Permission level for the agent during scheduled " + "execution. Defaults to DONT_ASK since no user is present to " + "answer prompts.", + ) + + source: ScheduleSource = Field( + default=ScheduleSource.USER, + description="Indicates how this schedule was created.", + ) + + source_session_id: str = Field( + default="", + description="The source session identifier, used for resource " + "retrieval.", + ) + + +class ScheduleRecord(_RecordBase): + """Persisted schedule record.""" + + user_id: str = Field(description="Owner user id.") + + agent_id: str = Field( + description="The agent id that will execute the schedule.", + ) + + data: ScheduleData = Field(description="Schedule configuration.") diff --git a/src/agentscope/app/storage/_model/_session.py b/src/agentscope/app/storage/_model/_session.py new file mode 100644 index 0000000000000000000000000000000000000000..0203388eb9f97462f09b93c6651f20cf868492dc --- /dev/null +++ b/src/agentscope/app/storage/_model/_session.py @@ -0,0 +1,168 @@ +# -*- coding: utf-8 -*- +"""The session data class for storage.""" +from datetime import datetime +from enum import Enum + +from pydantic import BaseModel, Field + +from ._base import _RecordBase +from ....state import AgentState + + +class SessionSource(str, Enum): + """The source that created the session.""" + + USER = "user" + SCHEDULE = "schedule" + + +class ChatModelConfig(BaseModel): + """The model configuration class.""" + + type: str + """The provider type.""" + + credential_id: str + """The credential id.""" + + model: str + """The model name.""" + + parameters: dict + """The model parameters.""" + + +class TTSModelConfig(BaseModel): + """The TTS model configuration class.""" + + type: str + """The provider type.""" + + credential_id: str + """The credential id.""" + + model: str + """The TTS model name.""" + + parameters: dict + """TTS parameters (voice, language, etc.).""" + + +class EmbeddingModelConfig(BaseModel): + """Configuration for constructing an embedding model from a credential. + + Mirrors :class:`ChatModelConfig` but targets + :class:`~agentscope.embedding.EmbeddingModelBase` subclasses. + Used by :class:`KnowledgeBaseRecord` to persist the user's + embedding model selection. + """ + + type: str + """The provider type (e.g. ``"openai_credential"``).""" + + credential_id: str + """The credential id to use for authentication.""" + + model: str + """The embedding model name (e.g. ``"text-embedding-3-small"``).""" + + dimensions: int = Field(..., gt=0) + """The output embedding vector dimensions. + + Required and first-class — chosen at config-creation time and + pinned to the resulting :class:`KnowledgeBaseRecord` so subsequent + indexing / retrieval calls are dim-deterministic without any + fallback lookup. + """ + + parameters: dict = Field(default_factory=dict) + """The provider-specific non-dimensional parameters. + + Does **not** carry ``dimensions`` — that field is promoted to a + top-level attribute above. + """ + + +class SessionKnowledgeConfig(BaseModel): + """Session-level knowledge base attachment. + + Persists which knowledge bases the agent should retrieve from for + this session and how the + :class:`~agentscope.middleware.RAGMiddleware` should be + configured. ``parameters`` carries the user-tunable middleware + fields verbatim (mirrors :attr:`ChatModelConfig.parameters`); the + accepted keys and value types are described by + :meth:`RAGMiddleware.Config.model_json_schema`. + """ + + knowledge_base_ids: list[str] = Field(default_factory=list) + """Ids of the knowledge bases attached to this session. + + Empty list means no knowledge base is wired and the middleware is + not installed. + """ + + parameters: dict = Field(default_factory=dict) + """Middleware parameters keyed by ``RAGMiddleware``'s + :class:`Config` model fields (``mode``, ``top_k``, + ``score_threshold``, ``emit_hint_event``, ``persist_hint``, + ``hint_template``). + """ + + +class SessionConfig(BaseModel): + """Session configuration — set at creation, updatable via PATCH.""" + + workspace_id: str + """The workspace id this session is bound to.""" + + name: str = Field( + default_factory=lambda: datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + description="Display name for the session.", + ) + """The session display name.""" + + chat_model_config: ChatModelConfig | None = None + """The chat model config. None means no model has been configured yet.""" + + fallback_chat_model_config: ChatModelConfig | None = None + """The fallback chat model config. Used as a backup when the primary + model fails. None means no fallback configured.""" + + tts_model_config: TTSModelConfig | None = None + """The TTS model config. None means TTS is not enabled.""" + + knowledge_config: SessionKnowledgeConfig | None = None + """Knowledge bases attached to this session and the corresponding + :class:`~agentscope.middleware.RAGMiddleware` parameters. + ``None`` means no knowledge base is wired.""" + + +class SessionRecord(_RecordBase): + """The session record.""" + + user_id: str + """The user id.""" + + agent_id: str + """The agent id.""" + + source: SessionSource = SessionSource.USER + """The source that created this session.""" + + source_schedule_id: str | None = None + """The source schedule Id.""" + + team_id: str | None = None + """The team this session participates in, if any. + + Team membership is session-level: a user agent can lead multiple teams + across different sessions, and each worker session belongs to exactly + one team. ``None`` means the session is not part of any team. + """ + + config: SessionConfig + """Session configuration (workspace, name, model).""" + + state: AgentState = Field(default_factory=AgentState) + """Mutable runtime state, updated after each chat turn.""" diff --git a/src/agentscope/app/storage/_model/_team.py b/src/agentscope/app/storage/_model/_team.py new file mode 100644 index 0000000000000000000000000000000000000000..a70d65f2b4356bd64526f9b4f2676c90dc33af5d --- /dev/null +++ b/src/agentscope/app/storage/_model/_team.py @@ -0,0 +1,55 @@ +# -*- coding: utf-8 -*- +"""The team storage class.""" +from pydantic import BaseModel, Field + +from ._base import _RecordBase + + +class TeamData(BaseModel): + """The team data model.""" + + name: str = Field( + description="Display name of the team.", + title="Name", + ) + + description: str = Field( + default="", + description=( + "What the team is for — its overall goal or shared context. " + "Wired into every member's system prompt so all members share " + "the same high-level understanding of why the team exists." + ), + title="Description", + ) + + member_ids: list[str] = Field( + default_factory=list, + description=( + "Worker agent ids that belong to this team. Each worker has " + "``source='team'`` and exactly one session, so the agent id " + "uniquely identifies the member; the session can be looked up " + "via :meth:`StorageBase.list_sessions`." + ), + title="Member Ids", + ) + + +class TeamRecord(_RecordBase): + """The team ORM model. + + Team membership is session-level: the leader is identified by its + ``session_id`` (since a user agent can lead multiple teams across + different sessions). Workers are identified by their agent id in + :attr:`TeamData.member_ids` (since workers have a 1:1 mapping between + agent and session). + """ + + user_id: str + """The user id.""" + + session_id: str + """The leader session id — the session that called ``create_team``.""" + + data: TeamData + """The team data.""" diff --git a/src/agentscope/app/storage/_model/_user.py b/src/agentscope/app/storage/_model/_user.py new file mode 100644 index 0000000000000000000000000000000000000000..78abf7a871b5b3ae56f98bf1cc13097cab716313 --- /dev/null +++ b/src/agentscope/app/storage/_model/_user.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- +"""The user record for storage.""" + +from ._base import _RecordBase + + +class UserRecord(_RecordBase): + """The user record.""" diff --git a/src/agentscope/app/storage/_redis_storage.py b/src/agentscope/app/storage/_redis_storage.py new file mode 100644 index 0000000000000000000000000000000000000000..cba91a8e7b71ce82aeb1be6d4b0f3deabc81c555 --- /dev/null +++ b/src/agentscope/app/storage/_redis_storage.py @@ -0,0 +1,1732 @@ +# -*- coding: utf-8 -*- +# pylint: disable=too-many-public-methods +"""The Redis storage implementation.""" + +from datetime import datetime, timedelta +from typing import Any, TYPE_CHECKING, Self + +from pydantic import BaseModel + +from ._base import StorageBase +from ._model import ( + AgentRecord, + CredentialRecord, + KnowledgeBaseRecord, + KnowledgeDocumentRecord, + KnowledgeDocumentStatus, + ScheduleRecord, + SessionRecord, + SessionConfig, + SessionSource, + TeamRecord, +) +from ._utils import _dump_with_secrets +from ...credential import CredentialBase +from ...message import Msg +from ...state import AgentState + +if TYPE_CHECKING: + from redis.asyncio import ConnectionPool, Redis +else: + ConnectionPool = Any + Redis = Any + + +def _watch_error() -> type[BaseException]: + """Return the ``WatchError`` class from ``redis.exceptions``. + + Lazy-imported so ``redis`` stays an optional dependency at module + load time — same trick the storage class uses in ``__aenter__``. + """ + from redis.exceptions import WatchError + + return WatchError + + +class RedisStorage(StorageBase): + """The Redis storage implementation.""" + + class KeyConfig(BaseModel): + """Key templates for all Redis keys used by :class:`RedisStorage`. + + Nested on :class:`RedisStorage` because customising key prefixes + is meaningful only for this backend; users tweak them via + ``RedisStorage(key_config=RedisStorage.KeyConfig(...))``. + """ + + # Record keys + credential: str = ( + "agentscope:user:{user_id}:credential:{credential_id}" + ) + agent: str = "agentscope:user:{user_id}:agent:{agent_id}" + session: str = "agentscope:user:{user_id}:session:{session_id}" + + # Index keys (Redis Sets — store all IDs for a given scope) + credential_index: str = "agentscope:user:{user_id}:credentials" + agent_index: str = "agentscope:user:{user_id}:agents" + session_index: str = ( + "agentscope:user:{user_id}:agent:{agent_id}:sessions" + ) + + # Lookup key: maps (user_id, agent_id) → session_id + session_lookup: str = ( + "agentscope:user:{user_id}:agent:{agent_id}:session" + ) + + # Message list key (Redis List — ordered message history per session) + messages: str = ( + "agentscope:user:{user_id}:session:{session_id}:messages" + ) + + schedule: str = "agentscope:user:{user_id}:schedule:{schedule_id}" + schedule_index: str = "agentscope:user:{user_id}:schedules" + schedule_global_index: str = "agentscope:schedules" + schedule_session_index: str = ( + "agentscope:user:{user_id}:schedule:{schedule_id}:sessions" + ) + + team: str = "agentscope:user:{user_id}:team:{team_id}" + team_index: str = "agentscope:user:{user_id}:teams" + + knowledge_base: str = ( + "agentscope:user:{user_id}:knowledge_base:{knowledge_base_id}" + ) + knowledge_base_index: str = "agentscope:user:{user_id}:knowledge_bases" + + # Knowledge document keys + knowledge_document: str = ( + "agentscope:user:{user_id}" + ":knowledge_base:{knowledge_base_id}" + ":document:{document_id}" + ) + knowledge_document_index: str = ( + "agentscope:user:{user_id}" + ":knowledge_base:{knowledge_base_id}:documents" + ) + # Global index of every document key as ``user_id:kb_id:doc_id``; + # used by the lease sweeper, never by per-user listing. + knowledge_document_global_index: str = "agentscope:knowledge_documents" + + def __init__( + self, + host: str = "localhost", + port: int = 6379, + db: int = 0, + password: str | None = None, + connection_pool: ConnectionPool | None = None, + key_ttl: int | None = None, + key_config: "RedisStorage.KeyConfig | None" = None, + **kwargs: Any, + ) -> None: + """Store connection parameters; the actual pool is created in + :meth:`__aenter__`. + + Args: + host (`str`, defaults to `"localhost"`): Redis server host. + port (`int`, defaults to `6379`): Redis server port. + db (`int`, defaults to `0`): Redis database index. + password (`str | None`, optional): Redis password if required. + connection_pool (`ConnectionPool | None`, optional): + An externally managed connection pool. When provided the pool + is used as-is and **not** closed by :meth:`aclose` — the + caller retains ownership of its lifecycle. When omitted a + pool is created from *host*/*port*/*db*/*password* on + :meth:`__aenter__` and closed on :meth:`aclose`. + Extra ``**kwargs`` (e.g. ``max_connections``) are forwarded to + the pool constructor only when the pool is created internally. + key_ttl (`int | None`, optional): + Expire time in seconds for record keys. Refreshed on every + write (sliding TTL). If `None`, keys do not expire. + key_config (`RedisStorage.KeyConfig | None`, optional): + Key template configuration. Defaults to + ``RedisStorage.KeyConfig()``. + **kwargs (`Any`): + Extra keyword arguments forwarded to + ``redis.asyncio.ConnectionPool`` when the pool is created + internally (e.g. ``max_connections=20``, ``socket_timeout=5``). + """ + self._host = host + self._port = port + self._db = db + self._password = password + self._external_pool: ConnectionPool | None = connection_pool + self._kwargs = kwargs + self.key_ttl = key_ttl + self.key_config = key_config or RedisStorage.KeyConfig() + + # Populated in __aenter__; None until the context is entered. + self._client: Redis | None = None + self._owned_pool: ConnectionPool | None = None + + def _key(self, template: str, **kwargs: str) -> str: + """Format a key template with the given keyword arguments.""" + return template.format(**kwargs) + + async def _set_with_ttl(self, key: str, value: str) -> None: + """SET a key and optionally apply the sliding TTL.""" + await self._client.set(key, value) + await self._refresh_key_ttl(key) + + async def _refresh_key_ttl(self, key: str) -> None: + """Apply the sliding TTL to a key, if configured.""" + if self.key_ttl is not None: + await self._client.expire(key, self.key_ttl) + + async def __aenter__(self) -> Self: + """Create the connection pool and Redis client. + + If an external pool was supplied at construction time it is used + directly and its lifecycle remains the caller's responsibility. + Otherwise, an internal pool is created from the stored host/port/db + parameters and will be closed by :meth:`aclose`. + """ + try: + import redis.asyncio as aioredis + except ImportError as e: + raise ImportError( + "The 'redis' package is required for RedisStorage. " + "Install it with: pip install redis[async]", + ) from e + + if self._external_pool is not None: + pool = self._external_pool + else: + self._owned_pool = aioredis.ConnectionPool( + host=self._host, + port=self._port, + db=self._db, + password=self._password, + decode_responses=True, + **self._kwargs, + ) + pool = self._owned_pool + + self._client = aioredis.Redis(connection_pool=pool) + return self + + async def aclose(self) -> None: + """Close the connection pool if it was created internally. + + Externally supplied pools are left open — the caller owns them. + """ + if self._owned_pool is not None: + await self._owned_pool.aclose() + self._owned_pool = None + self._client = None + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: Any, + ) -> None: + """Exit the async context manager.""" + await self.aclose() + + def get_client(self) -> Redis: + """Get the underlying Redis client instance.""" + return self._client + + async def _generate_credential_name( + self, + user_id: str, + credential_data: CredentialBase, + ) -> str: + """Auto-generate a display name for a credential based on its type. + + Produces names like "OpenAI", "OpenAI (2)", "OpenAI (3)", etc. + """ + cred_type = getattr(credential_data, "type", "") + base_name = ( + cred_type.removesuffix("_credential").replace("_", " ").title() + ) + if not base_name: + base_name = "Credential" + + existing = await self.list_credentials(user_id) + same_type_names = [ + c.data.get("name", "") + for c in existing + if c.data.get("type") == cred_type and c.id != credential_data.id + ] + + if base_name not in same_type_names: + return base_name + + idx = 2 + while f"{base_name} ({idx})" in same_type_names: + idx += 1 + return f"{base_name} ({idx})" + + async def upsert_credential( + self, + user_id: str, + credential_data: CredentialBase, + ) -> str: + """Create or update a credential record for the given user. + + If `credential_data.id` is set and the record already exists, the + existing record's `data` field is updated in place (preserving + `created_at`). If the id is set but no record exists, a new record is + created with that id. If `credential_data.id` is ``None``, a new + record with a generated id is always created. + + Args: + user_id (`str`): + The owner user id. + credential_data (`CredentialBase`): + Input data containing an optional `id` and the credential + `data` dict. + + Returns: + `str`: + The id of the created or updated credential record. + """ + if not credential_data.name: + credential_data.name = await self._generate_credential_name( + user_id, + credential_data, + ) + + data_dump = _dump_with_secrets(credential_data) + + if credential_data.id: + key = self._key( + self.key_config.credential, + user_id=user_id, + credential_id=credential_data.id, + ) + raw = await self._client.get(key) + if raw: + record = CredentialRecord.model_validate_json(raw) + record.data = data_dump + record.updated_at = datetime.now() + else: + record = CredentialRecord( + id=credential_data.id, + user_id=user_id, + data=data_dump, + ) + else: + record = CredentialRecord( + user_id=user_id, + data=data_dump, + ) + + key = self._key( + self.key_config.credential, + user_id=user_id, + credential_id=record.id, + ) + index_key = self._key( + self.key_config.credential_index, + user_id=user_id, + ) + await self._set_with_ttl(key, record.model_dump_json()) + await self._client.sadd(index_key, record.id) + return record.id + + async def list_credentials(self, user_id: str) -> list[CredentialRecord]: + """Return all credential records belonging to the given user. + + Reads the per-user credential index Set to obtain all ids, then + fetches each record individually. Records whose keys have expired or + been deleted externally are silently skipped. + + Args: + user_id (`str`): The owner user id. + + Returns: + `list[CredentialRecord]`: All credential records for the user. + """ + index_key = self._key( + self.key_config.credential_index, + user_id=user_id, + ) + ids = await self._client.smembers(index_key) + records = [] + for cred_id in ids: + raw = await self._client.get( + self._key( + self.key_config.credential, + user_id=user_id, + credential_id=cred_id, + ), + ) + if raw: + records.append(CredentialRecord.model_validate_json(raw)) + return records + + async def get_credential( + self, + user_id: str, + credential_id: str, + ) -> CredentialRecord | None: + """Fetch a single credential record by id.""" + key = self._key( + self.key_config.credential, + user_id=user_id, + credential_id=credential_id, + ) + raw = await self._client.get(key) + return CredentialRecord.model_validate_json(raw) if raw else None + + async def delete_credential( + self, + user_id: str, + credential_id: str, + ) -> bool: + """Delete a credential record and remove it from the user's index. + + Args: + user_id (`str`): The owner user id. + credential_id (`str`): The id of the credential to delete. + + Returns: + `bool`: ``True`` if the record existed and was deleted, + ``False`` if it did not exist. + """ + key = self._key( + self.key_config.credential, + user_id=user_id, + credential_id=credential_id, + ) + index_key = self._key( + self.key_config.credential_index, + user_id=user_id, + ) + deleted = await self._client.delete(key) + await self._client.srem(index_key, credential_id) + return deleted > 0 + + async def upsert_agent( + self, + user_id: str, + agent_record: AgentRecord, + ) -> str: + """Persist an agent record and register it in the user's agent index. + + The caller is responsible for constructing the full `AgentRecord` + (including its `id`). If a record with the same id already exists it + will be overwritten. + + Args: + user_id (`str`): + The owner user id. + agent_record (`AgentRecord`): + The fully-populated agent record to store. + + Returns: + `str`: + The id of the stored agent record. + """ + key = self._key( + self.key_config.agent, + user_id=user_id, + agent_id=agent_record.id, + ) + index_key = self._key(self.key_config.agent_index, user_id=user_id) + await self._set_with_ttl(key, agent_record.model_dump_json()) + await self._client.sadd(index_key, agent_record.id) + return agent_record.id + + async def list_agents(self, user_id: str) -> list[AgentRecord]: + """Return user-facing agent records (``source='user'``). + + Reads the per-user agent index Set to obtain all ids, fetches + each record individually, and **filters out team-spawned + workers** (``source='team'``) — those are scoped to a team + and only addressable via team detail / direct id lookup, not + enumerated as part of the user's regular agent list. + + Records whose keys have expired or been deleted externally + are silently skipped. + + Args: + user_id (`str`): The owner user id. + + Returns: + `list[AgentRecord]`: + All ``source='user'`` agent records for the user. + """ + index_key = self._key(self.key_config.agent_index, user_id=user_id) + ids = await self._client.smembers(index_key) + records = [] + for agent_id in ids: + raw = await self._client.get( + self._key( + self.key_config.agent, + user_id=user_id, + agent_id=agent_id, + ), + ) + if raw: + record = AgentRecord.model_validate_json(raw) + if record.source == "user": + records.append(record) + return records + + async def get_agent( + self, + user_id: str, + agent_id: str, + ) -> AgentRecord | None: + """Fetch a single agent record by id.""" + key = self._key( + self.key_config.agent, + user_id=user_id, + agent_id=agent_id, + ) + raw = await self._client.get(key) + return AgentRecord.model_validate_json(raw) if raw else None + + async def delete_agent(self, user_id: str, agent_id: str) -> bool: + """Delete an agent record and cascade-delete its sessions, + schedules, and any team back-references. + + Cascade order: + + 1. **Sessions** — every session belonging to this agent is + deleted via :meth:`delete_session` (which itself cascades + message log, schedule-session index, and — if a session leads + a team — the team). + 2. **Schedules** — every schedule whose ``data.agent_id`` matches + is deleted via :meth:`delete_schedule`. + 3. **Team back-references (defensive)** — if the agent is a team + worker (``source='team'``) but the caller chose to delete it + directly instead of going through :meth:`delete_team`, scan + the user's teams and remove the agent id from every + :attr:`TeamData.member_ids` list it appears in. The normal + path (``delete_team`` iterates ``member_ids`` and calls + ``delete_agent`` for each) does not need this scan, but it + keeps the team record consistent if a caller bypasses it. + 4. **Agent record + index** — finally delete the agent key and + remove from the per-user agent index. + + Args: + user_id (`str`): + The owner user id. + agent_id (`str`): + The id of the agent to delete. + + Returns: + `bool`: + ``True`` if the agent record existed and was deleted, + ``False`` if it did not exist. + """ + # Cascade: sessions + sessions = await self.list_sessions(user_id, agent_id) + for session in sessions: + await self.delete_session(user_id, agent_id, session.id) + + # Cascade: schedules owned by this agent + schedules = await self.list_schedules(user_id) + for schedule in schedules: + if schedule.agent_id == agent_id: + await self.delete_schedule(user_id, schedule.id) + + # Defensive: scrub agent_id from any team's member_ids list. + # The common path (delete_team -> delete_agent) is unaffected + # because the team is being torn down anyway and removed from + # the index in step 4 of delete_team. + teams = await self.list_teams(user_id) + for team in teams: + if agent_id in team.data.member_ids: + team.data.member_ids = [ + mid for mid in team.data.member_ids if mid != agent_id + ] + await self.upsert_team(user_id, team) + + key = self._key( + self.key_config.agent, + user_id=user_id, + agent_id=agent_id, + ) + index_key = self._key(self.key_config.agent_index, user_id=user_id) + deleted = await self._client.delete(key) + await self._client.srem(index_key, agent_id) + return deleted > 0 + + async def upsert_session( + self, + user_id: str, + agent_id: str, + config: SessionConfig, + state: AgentState | None = None, + session_id: str | None = None, + source: SessionSource = SessionSource.USER, + source_schedule_id: str | None = None, + ) -> SessionRecord: + """Create or update a session for a (user, agent) pair. + + When *session_id* is provided the existing session is updated. + When *session_id* is ``None`` a new session is always created. + """ + if session_id: + key = self._key( + self.key_config.session, + user_id=user_id, + session_id=session_id, + ) + raw = await self._client.get(key) + if raw: + record = SessionRecord.model_validate_json(raw) + record.config = config + if state is not None: + record.state = state + record.updated_at = datetime.now() + await self._set_with_ttl(key, record.model_dump_json()) + return record + + # Use the caller-provided ``session_id`` when given so a + # "create-if-missing under this id" call (e.g. scheduler's + # stateful-mode session) lands at the expected key. + new_id_kwargs = {"id": session_id} if session_id else {} + record = SessionRecord( + user_id=user_id, + agent_id=agent_id, + config=config, + source=source, + source_schedule_id=source_schedule_id, + state=state if state is not None else AgentState(), + **new_id_kwargs, + ) + key = self._key( + self.key_config.session, + user_id=user_id, + session_id=record.id, + ) + index_key = self._key( + self.key_config.session_index, + user_id=user_id, + agent_id=agent_id, + ) + await self._set_with_ttl(key, record.model_dump_json()) + await self._client.sadd(index_key, record.id) + + if source_schedule_id: + schedule_session_key = self._key( + self.key_config.schedule_session_index, + user_id=user_id, + schedule_id=source_schedule_id, + ) + await self._client.sadd(schedule_session_key, record.id) + + return record + + async def update_session_state( + self, + user_id: str, + agent_id: str, + session_id: str, + state: AgentState, + ) -> None: + """Update only the mutable state of an existing session. + + Raises: + KeyError: If the session does not exist. + """ + key = self._key( + self.key_config.session, + user_id=user_id, + session_id=session_id, + ) + raw = await self._client.get(key) + if not raw: + raise KeyError(f"Session {session_id!r} not found.") + record = SessionRecord.model_validate_json(raw) + record.state = state + record.updated_at = datetime.now() + await self._set_with_ttl(key, record.model_dump_json()) + + async def list_sessions( + self, + user_id: str, + agent_id: str, + ) -> list[SessionRecord]: + """Return all session records for a given (user, agent) pair. + + Reads the per-agent session index Set to obtain all session ids, then + fetches each record individually. Records whose keys have expired or + been deleted externally are silently skipped. + + Args: + user_id (`str`): The owner user id. + agent_id (`str`): The agent id whose sessions to list. + + Returns: + `list[SessionRecord]`: All session records for the (user, agent) + pair. + """ + index_key = self._key( + self.key_config.session_index, + user_id=user_id, + agent_id=agent_id, + ) + ids = await self._client.smembers(index_key) + records = [] + for session_id in ids: + raw = await self._client.get( + self._key( + self.key_config.session, + user_id=user_id, + session_id=session_id, + ), + ) + if raw: + records.append(SessionRecord.model_validate_json(raw)) + records.sort(key=lambda r: r.created_at, reverse=True) + return records + + async def get_session( + self, + user_id: str, + agent_id: str, + session_id: str, + ) -> SessionRecord | None: + """Fetch a single session record by id.""" + key = self._key( + self.key_config.session, + user_id=user_id, + session_id=session_id, + ) + raw = await self._client.get(key) + if not raw: + return None + return SessionRecord.model_validate_json(raw) + + async def delete_session( + self, + user_id: str, + agent_id: str, + session_id: str, + ) -> bool: + """Delete a session record and cascade clean-up. + + Cascades: + + - Existing: per-session message log, schedule-session index entry. + - **NEW**: if this session is the leader of a team (``team_id`` + set AND a :class:`TeamRecord` exists with + ``session_id == this session_id``), call :meth:`delete_team` + first. ``delete_team`` will recursively cascade workers and + clear ``team_id`` on this session — that clear is idempotent + and the session itself is deleted right after, so the order is + safe. + + Worker sessions (``team_id`` set, but the team's + ``leader_session_id`` is **not** this session) are deleted + without dissolving the team — the team and the surviving leader + keep their member_ids list pointing to the now-orphaned worker + agent. This intentional asymmetry mirrors SQL: there is no FK + from :class:`SessionRecord` back to the agent that owns it, so + deleting a session doesn't automatically delete the agent. + + Args: + user_id (`str`): + The owner user id. + agent_id (`str`): + The id of the agent that owns the session (used to + clean up the per-agent session index). + session_id (`str`): + The id of the session to delete. + + Returns: + `bool`: + ``True`` if the session existed and was deleted, + ``False`` if no record was found. + """ + key = self._key( + self.key_config.session, + user_id=user_id, + session_id=session_id, + ) + raw = await self._client.get(key) + if not raw: + return False + + record = SessionRecord.model_validate_json(raw) + + # Cascade: if this session leads a team, dissolve it first. + if record.team_id: + team = await self.get_team(user_id, record.team_id) + if team is not None and team.session_id == session_id: + await self.delete_team(user_id, record.team_id) + + index_key = self._key( + self.key_config.session_index, + user_id=user_id, + agent_id=agent_id, + ) + msg_key = self._key( + self.key_config.messages, + user_id=user_id, + session_id=session_id, + ) + await self._client.delete(key) + await self._client.srem(index_key, session_id) + await self._client.delete(msg_key) + + if record.source_schedule_id: + schedule_session_key = self._key( + self.key_config.schedule_session_index, + user_id=user_id, + schedule_id=record.source_schedule_id, + ) + await self._client.srem(schedule_session_key, session_id) + + return True + + async def list_sessions_by_schedule( + self, + user_id: str, + schedule_id: str, + ) -> list[SessionRecord]: + """Return all sessions created by a given schedule.""" + schedule_session_key = self._key( + self.key_config.schedule_session_index, + user_id=user_id, + schedule_id=schedule_id, + ) + ids = await self._client.smembers(schedule_session_key) + records = [] + for session_id in ids: + raw = await self._client.get( + self._key( + self.key_config.session, + user_id=user_id, + session_id=session_id, + ), + ) + if raw: + records.append(SessionRecord.model_validate_json(raw)) + records.sort(key=lambda r: r.created_at, reverse=True) + return records + + async def upsert_schedule( + self, + user_id: str, + record: ScheduleRecord, + ) -> str: + """Persist a cron task record and register it in the user and global + indexes.""" + key = self._key( + self.key_config.schedule, + user_id=user_id, + schedule_id=record.id, + ) + index_key = self._key(self.key_config.schedule_index, user_id=user_id) + await self._set_with_ttl(key, record.model_dump_json()) + await self._client.sadd(index_key, record.id) + await self._client.sadd( + self.key_config.schedule_global_index, + f"{user_id}:{record.id}", + ) + return record.id + + async def get_schedule( + self, + user_id: str, + schedule_id: str, + ) -> ScheduleRecord | None: + """Fetch a single cron task record by id.""" + key = self._key( + self.key_config.schedule, + user_id=user_id, + schedule_id=schedule_id, + ) + raw = await self._client.get(key) + if not raw: + return None + return ScheduleRecord.model_validate_json(raw) + + async def list_schedules(self, user_id: str) -> list[ScheduleRecord]: + """Return all cron task records belonging to the given user.""" + index_key = self._key( + self.key_config.schedule_index, + user_id=user_id, + ) + ids = await self._client.smembers(index_key) + records = [] + for schedule_id in ids: + raw = await self._client.get( + self._key( + self.key_config.schedule, + user_id=user_id, + schedule_id=schedule_id, + ), + ) + if raw: + records.append(ScheduleRecord.model_validate_json(raw)) + return records + + async def delete_schedule(self, user_id: str, schedule_id: str) -> bool: + """Delete a cron task record, cascade-delete its execution sessions, + and remove it from the user and global indexes.""" + key = self._key( + self.key_config.schedule, + user_id=user_id, + schedule_id=schedule_id, + ) + raw = await self._client.get(key) + if not raw: + return False + + record = ScheduleRecord.model_validate_json(raw) + + # Cascade: delete all sessions created by this schedule + sessions = await self.list_sessions_by_schedule(user_id, schedule_id) + for session in sessions: + await self.delete_session( + user_id, + record.agent_id, + session.id, + ) + + # Clean up the schedule session index key itself + schedule_session_key = self._key( + self.key_config.schedule_session_index, + user_id=user_id, + schedule_id=schedule_id, + ) + await self._client.delete(schedule_session_key) + + # Delete the schedule record and its index entries + index_key = self._key(self.key_config.schedule_index, user_id=user_id) + await self._client.delete(key) + await self._client.srem(index_key, schedule_id) + await self._client.srem( + self.key_config.schedule_global_index, + f"{user_id}:{schedule_id}", + ) + return True + + async def list_all_schedules(self) -> list[ScheduleRecord]: + """Return every schedule record across all users. + + Reads the global schedule index (a Redis Set of ``user_id:schedule_id`` + pairs) and fetches each record individually. Records whose keys have + expired or been deleted externally are silently skipped. + + Returns: + `list[ScheduleRecord]`: All schedule records in the store. + """ + entries = await self._client.smembers( + self.key_config.schedule_global_index, + ) + records = [] + for entry in entries: + user_id, schedule_id = entry.split(":", 1) + raw = await self._client.get( + self._key( + self.key_config.schedule, + user_id=user_id, + schedule_id=schedule_id, + ), + ) + if raw: + records.append(ScheduleRecord.model_validate_json(raw)) + return records + + # ------------------------------------------------------------------ + # Message persistence + # ------------------------------------------------------------------ + + def _message_key(self, user_id: str, session_id: str) -> str: + """Return the Redis List key for a session's messages.""" + return self._key( + self.key_config.messages, + user_id=user_id, + session_id=session_id, + ) + + async def upsert_message( + self, + user_id: str, + session_id: str, + msg: Msg, + ) -> None: + """Persist a message to the session's message list.""" + key = self._message_key(user_id, session_id) + last_raw = await self._client.lindex(key, -1) + if last_raw: + last_msg = Msg.model_validate_json(last_raw) + if last_msg.id == msg.id: + await self._client.lset(key, -1, msg.model_dump_json()) + await self._refresh_key_ttl(key) + return + await self._client.rpush(key, msg.model_dump_json()) + await self._refresh_key_ttl(key) + + async def get_message( + self, + user_id: str, + session_id: str, + message_id: str, + ) -> Msg | None: + """Fetch a single message by id from the session's message list.""" + key = self._message_key(user_id, session_id) + length = await self._client.llen(key) + for i in range(length - 1, -1, -1): + raw = await self._client.lindex(key, i) + if raw: + msg = Msg.model_validate_json(raw) + if msg.id == message_id: + return msg + return None + + async def list_messages( + self, + user_id: str, + session_id: str, + offset: int = 0, + limit: int = 50, + ) -> list[Msg]: + """Return messages for a session with pagination.""" + key = self._message_key(user_id, session_id) + raw_list = await self._client.lrange(key, offset, offset + limit - 1) + return [Msg.model_validate_json(raw) for raw in raw_list] + + # ------------------------------------------------------------------ + # Team persistence + # ------------------------------------------------------------------ + + async def upsert_team( + self, + user_id: str, + record: TeamRecord, + ) -> TeamRecord: + """Persist a team record and register it in the user's team index. + + Args: + user_id (`str`): + The owner user id. Used to scope both the record key and + the per-user team index. + record (`TeamRecord`): + The team record to persist. Its ``id`` is used as the + primary key; an existing record with the same id is + overwritten. ``updated_at`` is refreshed to ``datetime.now()`` + before writing. + + Returns: + `TeamRecord`: + The stored record (with refreshed ``updated_at``). + """ + record.updated_at = datetime.now() + key = self._key( + self.key_config.team, + user_id=user_id, + team_id=record.id, + ) + index_key = self._key(self.key_config.team_index, user_id=user_id) + await self._set_with_ttl(key, record.model_dump_json()) + await self._client.sadd(index_key, record.id) + return record + + async def get_team( + self, + user_id: str, + team_id: str, + ) -> TeamRecord | None: + """Fetch a single team record by id. + + Args: + user_id (`str`): + The owner user id. + team_id (`str`): + The team id to look up. + + Returns: + `TeamRecord | None`: + The record, or ``None`` if no record exists at the + ``(user_id, team_id)`` key (e.g. expired or never created). + """ + key = self._key( + self.key_config.team, + user_id=user_id, + team_id=team_id, + ) + raw = await self._client.get(key) + if not raw: + return None + return TeamRecord.model_validate_json(raw) + + async def list_teams(self, user_id: str) -> list[TeamRecord]: + """Return all team records belonging to the given user. + + Reads the per-user team index (a Redis Set of team ids) and fetches + each record individually. Records whose keys have expired or been + deleted externally are silently skipped. + + Args: + user_id (`str`): + The owner user id whose teams to list. + + Returns: + `list[TeamRecord]`: + All team records for the user, in arbitrary order (the + index is a Set). + """ + index_key = self._key(self.key_config.team_index, user_id=user_id) + ids = await self._client.smembers(index_key) + records: list[TeamRecord] = [] + for team_id in ids: + raw = await self._client.get( + self._key( + self.key_config.team, + user_id=user_id, + team_id=team_id, + ), + ) + if raw: + records.append(TeamRecord.model_validate_json(raw)) + return records + + async def set_session_team_id( + self, + user_id: str, + session_id: str, + team_id: str | None, + ) -> None: + """Set or clear ``team_id`` on an existing session record. + + Bypasses :meth:`upsert_session` because that method does not + allow writing ``team_id`` (which is a relation column the + application normally only mutates via team operations). + Idempotent: a no-op if the session does not exist or already + holds the given value. + + Args: + user_id (`str`): + The owner user id. + session_id (`str`): + The session whose ``team_id`` should be updated. + team_id (`str | None`): + The new value. ``None`` detaches the session from any + team (used by :meth:`delete_team` and by the team + service when a session leaves a team). + """ + key = self._key( + self.key_config.session, + user_id=user_id, + session_id=session_id, + ) + raw = await self._client.get(key) + if not raw: + return + record = SessionRecord.model_validate_json(raw) + if record.team_id == team_id: + return + record.team_id = team_id + record.updated_at = datetime.now() + await self._set_with_ttl(key, record.model_dump_json()) + + async def delete_team(self, user_id: str, team_id: str) -> bool: + """Delete a team record and cascade-delete all of its workers. + + Cascade order (mirrors what SQL's ``ON DELETE CASCADE`` would do + for the same set of foreign keys): + + 1. For each ``member_id`` in :attr:`TeamData.member_ids`, call + :meth:`delete_agent`. Each call cascades the worker's single + session via the existing agent-cascade logic. + 2. Clear ``team_id`` on the leader session (referenced by + :attr:`TeamRecord.session_id`) — semantically equivalent to + ``ON DELETE SET NULL`` for that direction of the relationship. + Idempotent if the session has already been deleted (no-op). + 3. Delete the :class:`TeamRecord` key and remove it from the + per-user team index. + + The cascade is best-effort: Redis has no cross-key transaction, + so a process crash mid-cascade may leave residue. Each step is + idempotent so retries are safe. + + Args: + user_id (`str`): + The owner user id. + team_id (`str`): + The id of the team to delete. + + Returns: + `bool`: + ``True`` if the team record existed and was deleted, + ``False`` if no record was found at the + ``(user_id, team_id)`` key. + """ + team = await self.get_team(user_id, team_id) + if team is None: + # Make sure the index is also clean if the record vanished + # for any reason. + index_key = self._key(self.key_config.team_index, user_id=user_id) + await self._client.srem(index_key, team_id) + return False + + # Cascade: delete each worker agent (which cascades its session) + for member_id in team.data.member_ids: + await self.delete_agent(user_id, member_id) + + # Clear team_id on the leader session (idempotent) + await self.set_session_team_id(user_id, team.session_id, None) + + # Delete the TeamRecord key + index entry + key = self._key( + self.key_config.team, + user_id=user_id, + team_id=team_id, + ) + existed = await self._client.delete(key) + index_key = self._key(self.key_config.team_index, user_id=user_id) + await self._client.srem(index_key, team_id) + return bool(existed) + + # ------------------------------------------------------------------ + # Knowledge base persistence + # ------------------------------------------------------------------ + + async def upsert_knowledge_base( + self, + user_id: str, + record: KnowledgeBaseRecord, + ) -> KnowledgeBaseRecord: + """Persist a knowledge base record and register it in the user index. + + If a record with the same ``id`` already exists it is overwritten + and ``updated_at`` is refreshed; ``created_at`` is preserved. + + Args: + user_id (`str`): + The owner user id. + record (`KnowledgeBaseRecord`): + The fully-populated record to store. + + Returns: + `KnowledgeBaseRecord`: + The stored record (with ``updated_at`` refreshed). + """ + if record.user_id != user_id: + raise ValueError( + "record.user_id does not match the given user_id.", + ) + + key = self._key( + self.key_config.knowledge_base, + user_id=user_id, + knowledge_base_id=record.id, + ) + existing_raw = await self._client.get(key) + if existing_raw: + existing = KnowledgeBaseRecord.model_validate_json(existing_raw) + record.created_at = existing.created_at + record.updated_at = datetime.now() + + index_key = self._key( + self.key_config.knowledge_base_index, + user_id=user_id, + ) + await self._set_with_ttl(key, record.model_dump_json()) + await self._client.sadd(index_key, record.id) + return record + + async def get_knowledge_base( + self, + user_id: str, + knowledge_base_id: str, + ) -> KnowledgeBaseRecord | None: + """Fetch a single knowledge base record by id. + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The knowledge base id. + + Returns: + `KnowledgeBaseRecord | None`: + The record, or ``None`` if not found. + """ + key = self._key( + self.key_config.knowledge_base, + user_id=user_id, + knowledge_base_id=knowledge_base_id, + ) + raw = await self._client.get(key) + return KnowledgeBaseRecord.model_validate_json(raw) if raw else None + + async def list_knowledge_bases( + self, + user_id: str, + ) -> list[KnowledgeBaseRecord]: + """List all knowledge base records belonging to the given user. + + Reads the per-user knowledge base index Set to obtain all ids, + then fetches each record individually. Records whose keys have + expired or been deleted externally are silently skipped. + + Args: + user_id (`str`): + The owner user id. + + Returns: + `list[KnowledgeBaseRecord]`: + All knowledge base records for the user. + """ + index_key = self._key( + self.key_config.knowledge_base_index, + user_id=user_id, + ) + ids = await self._client.smembers(index_key) + records: list[KnowledgeBaseRecord] = [] + for kb_id in ids: + raw = await self._client.get( + self._key( + self.key_config.knowledge_base, + user_id=user_id, + knowledge_base_id=kb_id, + ), + ) + if raw: + records.append( + KnowledgeBaseRecord.model_validate_json(raw), + ) + return records + + async def delete_knowledge_base( + self, + user_id: str, + knowledge_base_id: str, + ) -> bool: + """Delete a knowledge base record and remove it from the index. + + Cascades the per-knowledge-base document index: every + :class:`KnowledgeDocumentRecord` indexed under the KB is + removed from storage so no orphan documents survive the KB + deletion. Cleanup of the underlying vector store collection + and blob payloads remains the caller's responsibility (the + manager + the blob store, respectively). + + Args: + user_id (`str`): + The owner user id. + knowledge_base_id (`str`): + The id of the record to delete. + + Returns: + `bool`: + ``True`` if the record existed and was deleted, + ``False`` if not found. + """ + # Cascade: drop every document record so the sweeper does not + # later try to redispatch leases for a KB that no longer exists. + documents = await self.list_knowledge_documents( + user_id, + knowledge_base_id, + ) + for document in documents: + await self.delete_knowledge_document( + user_id, + knowledge_base_id, + document.id, + ) + + key = self._key( + self.key_config.knowledge_base, + user_id=user_id, + knowledge_base_id=knowledge_base_id, + ) + index_key = self._key( + self.key_config.knowledge_base_index, + user_id=user_id, + ) + deleted = await self._client.delete(key) + await self._client.srem(index_key, knowledge_base_id) + return deleted > 0 + + # ------------------------------------------------------------------ + # Knowledge document persistence + # ------------------------------------------------------------------ + + def _document_key( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> str: + """Format the Redis key for one document record.""" + return self._key( + self.key_config.knowledge_document, + user_id=user_id, + knowledge_base_id=knowledge_base_id, + document_id=document_id, + ) + + def _document_index_key( + self, + user_id: str, + knowledge_base_id: str, + ) -> str: + """Format the per-KB document index Set key.""" + return self._key( + self.key_config.knowledge_document_index, + user_id=user_id, + knowledge_base_id=knowledge_base_id, + ) + + def _document_global_token( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> str: + """Encode (user_id, kb_id, doc_id) for the global sweeper index.""" + return f"{user_id}:{knowledge_base_id}:{document_id}" + + async def upsert_knowledge_document( + self, + user_id: str, + record: KnowledgeDocumentRecord, + ) -> KnowledgeDocumentRecord: + """Persist a document record and update its indexes. + + If a record with the same ``id`` already exists it is + overwritten and ``updated_at`` refreshed; ``created_at`` is + preserved. + + Args: + user_id (`str`): + The owner user id. Must match ``record.user_id``. + record (`KnowledgeDocumentRecord`): + The fully-populated record to persist. + + Returns: + `KnowledgeDocumentRecord`: + The stored record (with ``updated_at`` refreshed). + """ + if record.user_id != user_id: + raise ValueError( + "record.user_id does not match the given user_id.", + ) + + key = self._document_key( + user_id, + record.knowledge_base_id, + record.id, + ) + existing_raw = await self._client.get(key) + if existing_raw: + existing = KnowledgeDocumentRecord.model_validate_json( + existing_raw, + ) + record.created_at = existing.created_at + record.updated_at = datetime.now() + + await self._set_with_ttl(key, record.model_dump_json()) + await self._client.sadd( + self._document_index_key(user_id, record.knowledge_base_id), + record.id, + ) + await self._client.sadd( + self.key_config.knowledge_document_global_index, + self._document_global_token( + user_id, + record.knowledge_base_id, + record.id, + ), + ) + return record + + async def get_knowledge_document( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> KnowledgeDocumentRecord | None: + """Fetch a single knowledge document record by id.""" + raw = await self._client.get( + self._document_key(user_id, knowledge_base_id, document_id), + ) + return ( + KnowledgeDocumentRecord.model_validate_json(raw) if raw else None + ) + + async def list_knowledge_documents( + self, + user_id: str, + knowledge_base_id: str, + ) -> list[KnowledgeDocumentRecord]: + """List all documents in a knowledge base. + + Reads the per-KB document index Set and fetches each record. + Records whose keys have expired or been deleted externally are + silently skipped. + """ + ids = await self._client.smembers( + self._document_index_key(user_id, knowledge_base_id), + ) + records: list[KnowledgeDocumentRecord] = [] + for document_id in ids: + raw = await self._client.get( + self._document_key(user_id, knowledge_base_id, document_id), + ) + if raw: + records.append( + KnowledgeDocumentRecord.model_validate_json(raw), + ) + return records + + async def delete_knowledge_document( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + ) -> bool: + """Delete a document record and remove it from the indexes.""" + key = self._document_key(user_id, knowledge_base_id, document_id) + deleted = await self._client.delete(key) + await self._client.srem( + self._document_index_key(user_id, knowledge_base_id), + document_id, + ) + await self._client.srem( + self.key_config.knowledge_document_global_index, + self._document_global_token( + user_id, + knowledge_base_id, + document_id, + ), + ) + return deleted > 0 + + async def update_knowledge_document_status( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + status: KnowledgeDocumentStatus, + error: str | None = None, + chunk_count: int | None = None, + ) -> None: + """Update only the status-related fields of a document record. + + Reads the record, mutates the status fields in memory, and + writes it back. Not atomic across multiple writers — relies + on the indexing worker holding the lease, which serialises + status transitions for a single document. + """ + key = self._document_key(user_id, knowledge_base_id, document_id) + raw = await self._client.get(key) + if not raw: + return + record = KnowledgeDocumentRecord.model_validate_json(raw) + record.data.status = status + if error is not None: + record.data.error = error + if chunk_count is not None: + record.data.chunk_count = chunk_count + record.updated_at = datetime.now() + await self._set_with_ttl(key, record.model_dump_json()) + + async def acquire_knowledge_document_lease( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + processing_node: str, + lease_ttl: timedelta, + now: datetime | None = None, + ) -> bool: + """Compare-and-swap acquisition of the processing lease. + + Implementation is a read-modify-write under a per-document + ``WATCH`` so two workers racing on the same document cannot + both win. Redis ``WATCH`` aborts the transaction if the key + changes between WATCH and EXEC; we retry a few times on + ``WatchError`` and otherwise give up (treating the contention + as "someone else already holds the lease"). + """ + now = now or datetime.now() + new_deadline = now + lease_ttl + + async with self._client.pipeline(transaction=True) as pipe: + for _ in range(3): + try: + key = self._document_key( + user_id, + knowledge_base_id, + document_id, + ) + await pipe.watch(key) + raw = await pipe.get(key) + if not raw: + await pipe.unwatch() + return False + record = KnowledgeDocumentRecord.model_validate_json(raw) + holder = record.processing_node + deadline = record.data.lease_expires_at + if ( + holder is not None + and deadline is not None + and deadline > now + ): + await pipe.unwatch() + return False + record.processing_node = processing_node + record.data.lease_expires_at = new_deadline + record.updated_at = now + pipe.multi() + pipe.set(key, record.model_dump_json()) + if self.key_ttl is not None: + pipe.expire(key, self.key_ttl) + await pipe.execute() + return True + except _watch_error(): + # Another writer touched the key between WATCH and + # EXEC; loop and re-read. + continue + return False + + async def renew_knowledge_document_lease( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + processing_node: str, + lease_ttl: timedelta, + now: datetime | None = None, + ) -> bool: + """Extend the lease this worker already holds.""" + now = now or datetime.now() + new_deadline = now + lease_ttl + + async with self._client.pipeline(transaction=True) as pipe: + for _ in range(3): + try: + key = self._document_key( + user_id, + knowledge_base_id, + document_id, + ) + await pipe.watch(key) + raw = await pipe.get(key) + if not raw: + await pipe.unwatch() + return False + record = KnowledgeDocumentRecord.model_validate_json(raw) + if record.processing_node != processing_node: + await pipe.unwatch() + return False + record.data.lease_expires_at = new_deadline + record.updated_at = now + pipe.multi() + pipe.set(key, record.model_dump_json()) + if self.key_ttl is not None: + pipe.expire(key, self.key_ttl) + await pipe.execute() + return True + except _watch_error(): + continue + return False + + async def release_knowledge_document_lease( + self, + user_id: str, + knowledge_base_id: str, + document_id: str, + processing_node: str, + ) -> None: + """Release the lease only if this worker still owns it.""" + async with self._client.pipeline(transaction=True) as pipe: + for _ in range(3): + try: + key = self._document_key( + user_id, + knowledge_base_id, + document_id, + ) + await pipe.watch(key) + raw = await pipe.get(key) + if not raw: + await pipe.unwatch() + return + record = KnowledgeDocumentRecord.model_validate_json(raw) + if record.processing_node != processing_node: + await pipe.unwatch() + return + record.processing_node = None + record.data.lease_expires_at = None + record.updated_at = datetime.now() + pipe.multi() + pipe.set(key, record.model_dump_json()) + if self.key_ttl is not None: + pipe.expire(key, self.key_ttl) + await pipe.execute() + return + except _watch_error(): + continue + + async def list_knowledge_documents_with_expired_lease( + self, + now: datetime | None = None, + ) -> list[KnowledgeDocumentRecord]: + """Return non-terminal documents whose lease has expired. + + Reads the global document index and filters in-memory — there + is no Redis-side filter primitive that knows about our nested + ``data.lease_expires_at`` field. For production deployments + with very large document counts a secondary index would pay + off, but for the v1 workload the global scan is acceptable + because the sweep runs on a slow cadence (minutes) and only + cares about the small subset of non-terminal documents. + """ + now = now or datetime.now() + terminal = {"ready", "error"} + tokens = await self._client.smembers( + self.key_config.knowledge_document_global_index, + ) + records: list[KnowledgeDocumentRecord] = [] + for token in tokens: + try: + user_id, kb_id, document_id = token.split(":", 2) + except ValueError: + continue + raw = await self._client.get( + self._document_key(user_id, kb_id, document_id), + ) + if not raw: + continue + record = KnowledgeDocumentRecord.model_validate_json(raw) + if record.data.status in terminal: + continue + if record.processing_node is None: + continue + if ( + record.data.lease_expires_at is not None + and record.data.lease_expires_at < now + ): + records.append(record) + return records + + async def list_knowledge_documents_pending_since( + self, + threshold: datetime, + ) -> list[KnowledgeDocumentRecord]: + """Return documents stuck in ``pending`` since before ``threshold``.""" + tokens = await self._client.smembers( + self.key_config.knowledge_document_global_index, + ) + records: list[KnowledgeDocumentRecord] = [] + for token in tokens: + try: + user_id, kb_id, document_id = token.split(":", 2) + except ValueError: + continue + raw = await self._client.get( + self._document_key(user_id, kb_id, document_id), + ) + if not raw: + continue + record = KnowledgeDocumentRecord.model_validate_json(raw) + if record.data.status != "pending": + continue + if record.created_at < threshold: + records.append(record) + return records diff --git a/src/agentscope/app/storage/_utils.py b/src/agentscope/app/storage/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..923c70a2f2204ba8a0ebeea1ba85e1d5fdfef94d --- /dev/null +++ b/src/agentscope/app/storage/_utils.py @@ -0,0 +1,29 @@ +# -*- coding: utf-8 -*- +"""The utils for storage.""" + +from pydantic import BaseModel, SecretStr + + +def _dump_with_secrets(model: BaseModel) -> dict: + """Dump the BaseModel instance with SecretStr fields. Used for + storage. + + Args: + model (`BaseModel`): + The model instance to dump. + + Returns: + `dict`: + The dumped JSON with secrets included. + """ + # Use mode='json' so that Pydantic converts non-JSON-native types + # (e.g. datetime, UUID) to their JSON-compatible representations. + # SecretStr fields will be masked at this step. + result = model.model_dump(mode="json") + + for field_name, _ in model.__class__.model_fields.items(): + value = getattr(model, field_name) + if isinstance(value, SecretStr): + result[field_name] = value.get_secret_value() + + return result diff --git a/src/agentscope/app/workspace_manager/__init__.py b/src/agentscope/app/workspace_manager/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d723beb682596c0ce1a14e872ea48446a5e0002 --- /dev/null +++ b/src/agentscope/app/workspace_manager/__init__.py @@ -0,0 +1,15 @@ +# -*- coding: utf-8 -*- +"""The workspace manager classes, responsible for managing the resources +and their lifecycles, and filesystem isolation.""" + +from ._base import WorkspaceManagerBase +from ._local_workspace_manager import LocalWorkspaceManager +from ._docker_workspace_manager import DockerWorkspaceManager +from ._e2b_workspace_manager import E2BWorkspaceManager + +__all__ = [ + "WorkspaceManagerBase", + "LocalWorkspaceManager", + "DockerWorkspaceManager", + "E2BWorkspaceManager", +] diff --git a/src/agentscope/app/workspace_manager/_base.py b/src/agentscope/app/workspace_manager/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..9d77949ed68e169f43a4b775cb265670a9ca3847 --- /dev/null +++ b/src/agentscope/app/workspace_manager/_base.py @@ -0,0 +1,73 @@ +# -*- coding: utf-8 -*- +"""Workspace manager implementations.""" + +from abc import ABC, abstractmethod +from typing import Self + +from ...workspace import WorkspaceBase + + +class WorkspaceManagerBase(ABC): + """Abstract base for workspace managers. + + Subclasses are expected to be used as async context managers — entering + the context activates any background machinery the subclass needs (e.g. + a TTL sweeper task) and exiting it tears that machinery down and closes + every cached workspace via :meth:`close_all`. + + The default ``__aenter__`` / ``__aexit__`` cover the common case where a + subclass has no background machinery: enter is a no-op, exit just calls + :meth:`close_all`. Subclasses that own background tasks should override + both. + """ + + @abstractmethod + async def get_workspace( + self, + user_id: str, + agent_id: str, + session_id: str, + workspace_id: str, + ) -> WorkspaceBase: + """Return an initialized workspace. + + Args: + user_id (`str`): + The user id. + agent_id (`str`): + The agent id. + session_id (`str`): + The session id. + workspace_id (`str`): + The workspace id (reconnection credential). + """ + + @abstractmethod + async def create_workspace( + self, + user_id: str, + agent_id: str, + session_id: str, + ) -> WorkspaceBase: + """Create a new workspace and return it.""" + + @abstractmethod + async def close(self, workspace_id: str) -> None: + """Close and evict a single workspace from the cache.""" + + @abstractmethod + async def close_all(self) -> None: + """Close every cached workspace. + + Pure "close all currently tracked workspaces" semantics — does not + imply the manager itself is being torn down. Use ``async with`` (or + :meth:`__aexit__` directly) for full manager shutdown. + """ + + async def __aenter__(self) -> Self: + """Enter the manager's lifetime. Default is a no-op.""" + return self + + async def __aexit__(self, *exc: object) -> None: + """Exit the manager's lifetime — closes all cached workspaces.""" + await self.close_all() diff --git a/src/agentscope/app/workspace_manager/_docker_workspace_manager.py b/src/agentscope/app/workspace_manager/_docker_workspace_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..77e256bcedfd4f831bf068d36a2ee38786bdd857 --- /dev/null +++ b/src/agentscope/app/workspace_manager/_docker_workspace_manager.py @@ -0,0 +1,371 @@ +# -*- coding: utf-8 -*- +"""DockerWorkspaceManager — lifecycle manager for :class:`DockerWorkspace`. + +Mirrors :class:`LocalWorkspaceManager` 1:1 in its public surface +(``get_workspace`` / ``create_workspace`` / ``close`` / ``close_all``) +so that callers — notably :class:`agentscope.app._service.ChatService` — +do not branch on backend. + +Differences from the local manager (allowed to surface only via the +constructor): + +* Workdir layout is two levels — ``//`` — + and is bind-mounted to ``/workspace`` inside each container, so the + agent always sees a flat ``/workspace`` regardless of host layout. +* ``workspace_id`` is forwarded into :class:`DockerWorkspace` so the + container name (``as_ws_``) is stable across process + restarts. A cache miss after a restart deterministically re-attaches + to the same container slot via ``containers.create_or_replace``. +* Idle workspaces are evicted by a dedicated background sweeper task + started in :meth:`__aenter__` and cancelled in :meth:`__aexit__` — + not lazily on each :meth:`get_workspace` call. This keeps idle + resource consumption bounded even when no traffic is arriving. +* ``close_all`` shuts containers down in parallel + (:func:`asyncio.gather`) — Docker ``kill + delete`` is slow enough + that linear teardown on shutdown is noticeable. +""" + +import asyncio +import os +import time +from typing import Self + +from agentscope._logging import logger +from agentscope.mcp import MCPClient +from agentscope.workspace._docker import DockerWorkspace +from agentscope.workspace._docker._make_dockerfile import ( + DEFAULT_BASE_IMAGE, + DEFAULT_GATEWAY_PORT, +) +from ._base import WorkspaceManagerBase + +DEFAULT_SWEEP_INTERVAL = 300.0 + + +class DockerWorkspaceManager(WorkspaceManagerBase): + """Manages :class:`DockerWorkspace` instances with TTL-based caching. + + The manager owns a single set of image-build parameters + (``base_image`` / ``node_version`` / ``extra_pip``) shared by every + workspace it produces; the resulting image is content-hashed so + rebuilds are skipped on cache hits. + + Use the manager as an ``async with`` context manager: entering it + starts the TTL sweeper task, exiting it stops the sweeper and then + closes every cached workspace via :meth:`close_all`. + """ + + def __init__( + self, + basedir: str, + *, + base_image: str = DEFAULT_BASE_IMAGE, + node_version: str = "20", + extra_pip: list[str] | None = None, + gateway_port: int = DEFAULT_GATEWAY_PORT, + env: dict[str, str] | None = None, + default_mcps: list[MCPClient] | None = None, + skill_paths: list[str] | None = None, + ttl: float = 3600.0, + sweep_interval: float = DEFAULT_SWEEP_INTERVAL, + ) -> None: + """Initialize the docker workspace manager. + + Args: + basedir (`str`): + Host root under which per-user/per-agent workdir are + created (``//``). Each + workdir is bind-mounted to ``/workspace`` inside its + container. + base_image (`str`, defaults to `DEFAULT_BASE_IMAGE`): + Base Docker image; must provide ``python3``. + node_version (`str`, defaults to `"20"`): + Major Node.js version (e.g. ``"20"``) to bake into + the image. + extra_pip (`list[str] | None`, optional): + Extra Python packages to install into the gateway + venv at image-build time. + gateway_port (`int`, defaults to `DEFAULT_GATEWAY_PORT`): + TCP port the in-container gateway listens on (always + exposed to a randomly assigned host port). + env (`dict[str, str] | None`, optional): + Environment variables to set inside every workspace's + container. + default_mcps (`list[MCPClient] | None`, optional): + MCP clients seeded into brand-new workspaces. Ignored + on subsequent restarts of a workdir that already + persists ``.mcp``. + skill_paths (`list[str] | None`, optional): + Skill directories seeded into brand-new workspaces. + ttl (`float`, defaults to `3600.0`): + Seconds before an idle cached workspace is evicted + and its container torn down. + sweep_interval (`float`, defaults to `DEFAULT_SWEEP_INTERVAL`): + How often (seconds) the background sweeper wakes up + to look for idle workspaces. Defaults to 5 minutes. + """ + self._basedir = os.path.abspath(basedir) + self._base_image = base_image + self._node_version = node_version + self._extra_pip = list(extra_pip or []) + self._gateway_port = gateway_port + self._env = dict(env or {}) + self._default_mcps = list(default_mcps or []) + self._skill_paths = list(skill_paths or []) + self._ttl = ttl + self._sweep_interval = sweep_interval + + # workspace_id → (workspace, last_access_monotonic) + self._cache: dict[str, tuple[DockerWorkspace, float]] = {} + self._lock = asyncio.Lock() + self._sweep_task: asyncio.Task | None = None + + # ── isolation helpers ───────────────────────────────────────── + + def _workdir_for(self, user_id: str, agent_id: str) -> str: + """Resolve the host workdir for ``(user_id, agent_id)``. + + Two-level layout — ``//`` — so + different users never share a bind-mount even when their + ``agent_id`` collides. + """ + return os.path.join(self._basedir, user_id, agent_id) + + # ── workspace construction ──────────────────────────────────── + + async def _build_and_start( + self, + *, + workspace_id: str, + user_id: str, + agent_id: str, + ) -> DockerWorkspace: + """Create a :class:`DockerWorkspace` for ``(user_id, agent_id)`` + and run its full ``initialize``. + + ``workspace_id`` is forwarded so the container name is + deterministic and the same id round-trips through the cache. + """ + workdir = self._workdir_for(user_id, agent_id) + os.makedirs(workdir, exist_ok=True) + ws = DockerWorkspace( + workspace_id=workspace_id, + workdir=workdir, + base_image=self._base_image, + node_version=self._node_version, + extra_pip=self._extra_pip, + gateway_port=self._gateway_port, + env=self._env, + default_mcps=self._default_mcps, + skill_paths=self._skill_paths, + ) + await ws.initialize() + return ws + + # ── public API ──────────────────────────────────────────────── + + async def get_workspace( + self, + user_id: str, + agent_id: str, + session_id: str, + workspace_id: str, + ) -> DockerWorkspace: + """Return an initialised workspace, building one on cache miss. + + On miss the manager calls ``DockerWorkspace(workspace_id=…)`` + with a deterministic workdir derived from ``(user_id, + agent_id)``. Image build, container creation and gateway + startup all happen inside the workspace's ``initialize``. + + Eviction of idle workspaces is *not* performed here — the + background sweeper started by :meth:`__aenter__` handles that. + + Args: + user_id (`str`): + Owning user identifier. + agent_id (`str`): + Agent identifier (controls the workdir). + session_id (`str`): + Session identifier (unused for isolation; sessions + share a workdir and partition under + ``sessions//``). + workspace_id (`str`): + Stable workspace identifier — used both as the cache + key and the container name suffix. + + Returns: + `DockerWorkspace`: + A live, initialised workspace. + """ + del session_id # accepted for interface parity; not used here + + async with self._lock: + cached = self._cache.get(workspace_id) + if cached is not None: + ws, _ = cached + self._cache[workspace_id] = (ws, time.monotonic()) + return ws + + # Cache miss: build under the lock to prevent two concurrent + # get_workspace(workspace_id=X) calls from creating two + # workspaces for the same id. + async with self._lock: + cached = self._cache.get(workspace_id) + if cached is not None: + ws, _ = cached + self._cache[workspace_id] = (ws, time.monotonic()) + return ws + + ws = await self._build_and_start( + workspace_id=workspace_id, + user_id=user_id, + agent_id=agent_id, + ) + self._cache[workspace_id] = (ws, time.monotonic()) + return ws + + async def create_workspace( + self, + user_id: str, + agent_id: str, + session_id: str, + ) -> DockerWorkspace: + """Build a brand-new workspace and track it. + + A fresh ``workspace_id`` is allocated by + :class:`DockerWorkspace` itself; the caller should persist + ``workspace.workspace_id`` for later :meth:`get_workspace` + calls. + + Args: + user_id (`str`): + Owning user identifier. + agent_id (`str`): + Agent identifier (controls the workdir). + session_id (`str`): + Session identifier (accepted for parity; not used + here). + + Returns: + `DockerWorkspace`: + The newly built workspace, already initialised. + """ + del session_id # accepted for interface parity; not used here + + workdir = self._workdir_for(user_id, agent_id) + os.makedirs(workdir, exist_ok=True) + ws = DockerWorkspace( + workdir=workdir, + base_image=self._base_image, + node_version=self._node_version, + extra_pip=self._extra_pip, + gateway_port=self._gateway_port, + env=self._env, + default_mcps=self._default_mcps, + skill_paths=self._skill_paths, + ) + await ws.initialize() + async with self._lock: + self._cache[ws.workspace_id] = (ws, time.monotonic()) + return ws + + async def close(self, workspace_id: str) -> None: + """Close and evict a single workspace from the cache. + + No-op when the workspace_id is not tracked. + + Args: + workspace_id (`str`): + The workspace to close. + """ + async with self._lock: + entry = self._cache.pop(workspace_id, None) + if entry is None: + return + ws, _ = entry + await self._safe_close(ws) + + async def close_all(self) -> None: + """Close every cached workspace in parallel. + + Docker ``kill + delete`` is slow per container; doing it + sequentially on app shutdown produces a noticeable stall, so + we fan the calls out with :func:`asyncio.gather`. + """ + async with self._lock: + entries = list(self._cache.values()) + self._cache.clear() + if not entries: + return + await asyncio.gather( + *(self._safe_close(ws) for ws, _ in entries), + return_exceptions=True, + ) + + # ── async context manager ───────────────────────────────────── + + async def __aenter__(self) -> Self: + """Start the TTL sweeper task.""" + if self._sweep_task is None: + self._sweep_task = asyncio.create_task(self._sweep_loop()) + return self + + async def __aexit__(self, *exc: object) -> None: + """Stop the TTL sweeper task, then close every cached workspace.""" + if self._sweep_task is not None: + self._sweep_task.cancel() + try: + await self._sweep_task + except (asyncio.CancelledError, Exception): + pass + self._sweep_task = None + await self.close_all() + + # ── background sweeper ─────────────────────────────────────── + + async def _sweep_loop(self) -> None: + """Periodically evict idle workspaces. + + Runs forever until cancelled. Each tick pops every cache entry + whose last-access is older than ``ttl`` and closes it outside + the lock; exceptions during close are logged and swallowed so + one bad container does not poison the sweeper. + """ + while True: + try: + await asyncio.sleep(self._sweep_interval) + except asyncio.CancelledError: + return + try: + await self._sweep_once() + except Exception: + logger.exception("Docker workspace sweeper tick failed") + + async def _sweep_once(self) -> None: + """One sweeper tick: evict expired entries and close them.""" + now = time.monotonic() + async with self._lock: + expired_ids = [ + wid + for wid, (_, ts) in self._cache.items() + if now - ts > self._ttl + ] + evicted = [self._cache.pop(wid)[0] for wid in expired_ids] + if not evicted: + return + await asyncio.gather( + *(self._safe_close(ws) for ws in evicted), + return_exceptions=True, + ) + + @staticmethod + async def _safe_close(ws: DockerWorkspace) -> None: + """Close a workspace, logging any failure instead of raising.""" + try: + await ws.close() + except Exception: + logger.exception( + "Failed to close DockerWorkspace %s", + ws.workspace_id, + ) diff --git a/src/agentscope/app/workspace_manager/_e2b_workspace_manager.py b/src/agentscope/app/workspace_manager/_e2b_workspace_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..3586b29a22b4bf32b0adf677c4e2c28e438e76ea --- /dev/null +++ b/src/agentscope/app/workspace_manager/_e2b_workspace_manager.py @@ -0,0 +1,387 @@ +# -*- coding: utf-8 -*- +"""E2BWorkspaceManager — lifecycle manager for :class:`E2BWorkspace`. + +Mirrors :class:`DockerWorkspaceManager` 1:1 in its public surface +(``get_workspace`` / ``create_workspace`` / ``close`` / ``close_all``) +so callers — notably :class:`agentscope.app._service.ChatService` — +do not branch on backend. + +Differences from the Docker manager: + +* No ``basedir`` / ``_workdir_for`` — E2B sandboxes carry their own + filesystem state across pause/resume, so there is nothing to + bind-mount and nothing to lay out on the host. +* No image build parameters (``base_image`` / ``node_version``); E2B + attaches to a pre-built template plus a runtime bootstrap. +* Reattachment uses E2B sandbox metadata. The ``workspace_id`` is + written into the sandbox's metadata at create time and looked up via + ``AsyncSandbox.list(query=...)`` inside + :meth:`E2BWorkspace.initialize`. The manager itself is metadata-blind + — it just forwards ``workspace_id`` and lets the workspace handle the + reattach. +* ``user_id`` / ``agent_id`` are surfaced as extra sandbox metadata + (``agentscope.user.id`` / ``agentscope.agent.id``) so users can + filter their own sandboxes in the E2B dashboard. They do **not** + participate in cache key resolution; the cache is keyed strictly on + ``workspace_id`` (same as Docker). +* Idle workspaces are evicted by a dedicated background sweeper task + started in :meth:`__aenter__` and cancelled in :meth:`__aexit__` — + not lazily on each :meth:`get_workspace` call. +* ``close_all`` fans calls out with :func:`asyncio.gather` because + ``sandbox.pause()`` is a remote round-trip per sandbox; sequentialising + it on app shutdown produces a noticeable stall. +""" + +import asyncio +import time +from typing import Self + +from agentscope._logging import logger +from agentscope.mcp import MCPClient +from agentscope.workspace import E2BWorkspace +from agentscope.workspace._e2b._bootstrap import ( + DEFAULT_GATEWAY_PORT, + DEFAULT_TEMPLATE, + DEFAULT_TIMEOUT, +) +from ._base import WorkspaceManagerBase + +DEFAULT_SWEEP_INTERVAL = 300.0 + + +class E2BWorkspaceManager(WorkspaceManagerBase): + """Manages :class:`E2BWorkspace` instances with TTL-based caching. + + Use the manager as an ``async with`` context manager: entering it + starts the TTL sweeper task, exiting it stops the sweeper and then + closes every cached workspace via :meth:`close_all`. + """ + + def __init__( + self, + *, + template: str = DEFAULT_TEMPLATE, + api_key: str = "", + domain: str = "", + timeout_seconds: int = DEFAULT_TIMEOUT, + gateway_port: int = DEFAULT_GATEWAY_PORT, + env: dict[str, str] | None = None, + sandbox_metadata: dict[str, str] | None = None, + extra_pip: list[str] | None = None, + default_mcps: list[MCPClient] | None = None, + skill_paths: list[str] | None = None, + ttl: float = 3600.0, + sweep_interval: float = DEFAULT_SWEEP_INTERVAL, + ) -> None: + """Initialize the E2B workspace manager. + + Args: + template (`str`, defaults to `DEFAULT_TEMPLATE`): + E2B template id passed to every workspace this + manager produces. Defaults to ``"base"``. + api_key (`str`, defaults to `""`): + E2B API key. ``""`` falls back to the ``E2B_API_KEY`` + env var on the SDK side. + domain (`str`, defaults to `""`): + Optional custom E2B domain (self-hosted etc.). + timeout_seconds (`int`, defaults to `DEFAULT_TIMEOUT`): + Sandbox keep-alive timeout passed to + ``AsyncSandbox.create`` / ``AsyncSandbox.connect``. + gateway_port (`int`, defaults to `DEFAULT_GATEWAY_PORT`): + TCP port the in-sandbox gateway listens on. + env (`dict[str, str] | None`, optional): + Environment variables baked into the sandbox at + create time. + sandbox_metadata (`dict[str, str] | None`, optional): + Extra metadata merged with the per-workspace + ``agentscope.workspace.id`` / ``agentscope.user.id`` / + ``agentscope.agent.id`` keys. Useful for downstream + E2B dashboard filtering. + extra_pip (`list[str] | None`, optional): + Extra Python packages to install into the gateway + venv during bootstrap. + default_mcps (`list[MCPClient] | None`, optional): + MCP clients seeded into brand-new workspaces. Ignored + on subsequent reattachments — the sandbox's persisted + ``.mcp`` file wins. + skill_paths (`list[str] | None`, optional): + Skill directories seeded into brand-new workspaces. + ttl (`float`, defaults to `3600.0`): + Seconds before an idle cached workspace is evicted + and its sandbox paused. + sweep_interval (`float`, defaults to `DEFAULT_SWEEP_INTERVAL`): + How often (seconds) the background sweeper wakes up + to look for idle workspaces. Defaults to 5 minutes. + """ + self._template = template + self._api_key = api_key + self._domain = domain + self._timeout_seconds = timeout_seconds + self._gateway_port = gateway_port + self._env = dict(env or {}) + self._sandbox_metadata = dict(sandbox_metadata or {}) + self._extra_pip = list(extra_pip or []) + self._default_mcps = list(default_mcps or []) + self._skill_paths = list(skill_paths or []) + self._ttl = ttl + self._sweep_interval = sweep_interval + + # workspace_id → (workspace, last_access_monotonic) + self._cache: dict[str, tuple[E2BWorkspace, float]] = {} + self._lock = asyncio.Lock() + self._sweep_task: asyncio.Task | None = None + + # ── metadata helper ─────────────────────────────────────────── + + def _metadata_for( + self, + user_id: str, + agent_id: str, + ) -> dict[str, str]: + """Build the extra sandbox metadata for ``(user_id, agent_id)``. + + ``E2BWorkspace`` always sets ``agentscope.workspace.id`` itself; + we add the user/agent keys here so they show up alongside it + in the E2B dashboard's metadata filter UI. + """ + return { + "agentscope.user.id": user_id, + "agentscope.agent.id": agent_id, + **self._sandbox_metadata, + } + + # ── workspace construction ──────────────────────────────────── + + async def _build_and_start( + self, + *, + workspace_id: str | None, + user_id: str, + agent_id: str, + ) -> E2BWorkspace: + """Construct an :class:`E2BWorkspace` and run its full ``initialize``. + + ``workspace_id=None`` lets :class:`WorkspaceBase` allocate a + fresh UUID — used by :meth:`create_workspace`. Otherwise the + provided id is forwarded so reattachment by metadata works on + the second call. + """ + ws = E2BWorkspace( + workspace_id=workspace_id, + template=self._template, + api_key=self._api_key, + domain=self._domain, + timeout_seconds=self._timeout_seconds, + gateway_port=self._gateway_port, + env=self._env, + sandbox_metadata=self._metadata_for(user_id, agent_id), + extra_pip=self._extra_pip, + default_mcps=self._default_mcps, + skill_paths=self._skill_paths, + ) + await ws.initialize() + return ws + + # ── public API ──────────────────────────────────────────────── + + async def get_workspace( + self, + user_id: str, + agent_id: str, + session_id: str, + workspace_id: str, + ) -> E2BWorkspace: + """Return an initialised workspace, reattaching on cache miss. + + On miss the manager calls ``E2BWorkspace(workspace_id=…)`` and + relies on its ``initialize`` to find any existing sandbox via + ``AsyncSandbox.list(query=SandboxQuery(metadata=...))`` and + ``connect`` to it (auto-resuming if paused) — or to ``create`` + a fresh sandbox otherwise. + + Eviction of idle workspaces is *not* performed here — the + background sweeper started by :meth:`__aenter__` handles that. + + Args: + user_id (`str`): + Owning user identifier (forwarded as sandbox metadata + only — not part of the cache key). + agent_id (`str`): + Agent identifier (forwarded as sandbox metadata only + — not part of the cache key). + session_id (`str`): + Session identifier (unused; sandboxes are + per-workspace, sessions partition under + ``sessions//``). + workspace_id (`str`): + Stable workspace identifier — the cache key and the + value stored in the sandbox's + ``agentscope.workspace.id`` metadata. + + Returns: + `E2BWorkspace`: + A live, initialised workspace. + """ + del session_id # accepted for interface parity; not used here + + async with self._lock: + cached = self._cache.get(workspace_id) + if cached is not None: + ws, _ = cached + self._cache[workspace_id] = (ws, time.monotonic()) + return ws + + # Cache miss: build under the lock to prevent two concurrent + # get_workspace(workspace_id=X) calls from creating two + # workspaces (and thus two sandboxes) for the same id. + async with self._lock: + cached = self._cache.get(workspace_id) + if cached is not None: + ws, _ = cached + self._cache[workspace_id] = (ws, time.monotonic()) + return ws + + ws = await self._build_and_start( + workspace_id=workspace_id, + user_id=user_id, + agent_id=agent_id, + ) + self._cache[workspace_id] = (ws, time.monotonic()) + return ws + + async def create_workspace( + self, + user_id: str, + agent_id: str, + session_id: str, + ) -> E2BWorkspace: + """Build a brand-new workspace and track it. + + A fresh ``workspace_id`` is allocated by + :class:`WorkspaceBase`; the caller should persist + ``workspace.workspace_id`` for later :meth:`get_workspace` + calls. + + Args: + user_id (`str`): + Owning user identifier (forwarded as sandbox metadata). + agent_id (`str`): + Agent identifier (forwarded as sandbox metadata). + session_id (`str`): + Session identifier (accepted for parity; not used + here). + + Returns: + `E2BWorkspace`: + The newly built workspace, already initialised. + """ + del session_id # accepted for interface parity; not used here + + ws = await self._build_and_start( + workspace_id=None, + user_id=user_id, + agent_id=agent_id, + ) + async with self._lock: + self._cache[ws.workspace_id] = (ws, time.monotonic()) + return ws + + async def close(self, workspace_id: str) -> None: + """Close (= pause the sandbox) and evict a single workspace. + + No-op when the workspace_id is not tracked. + + Args: + workspace_id (`str`): + The workspace to close. + """ + async with self._lock: + entry = self._cache.pop(workspace_id, None) + if entry is None: + return + ws, _ = entry + await self._safe_close(ws) + + async def close_all(self) -> None: + """Close every cached workspace in parallel. + + ``sandbox.pause()`` is a remote round-trip per sandbox; doing + it sequentially on app shutdown produces a noticeable stall, + so we fan the calls out with :func:`asyncio.gather`. + """ + async with self._lock: + entries = list(self._cache.values()) + self._cache.clear() + if not entries: + return + await asyncio.gather( + *(self._safe_close(ws) for ws, _ in entries), + return_exceptions=True, + ) + + # ── async context manager ───────────────────────────────────── + + async def __aenter__(self) -> Self: + """Start the TTL sweeper task.""" + if self._sweep_task is None: + self._sweep_task = asyncio.create_task(self._sweep_loop()) + return self + + async def __aexit__(self, *exc: object) -> None: + """Stop the TTL sweeper task, then close every cached workspace.""" + if self._sweep_task is not None: + self._sweep_task.cancel() + try: + await self._sweep_task + except (asyncio.CancelledError, Exception): + pass + self._sweep_task = None + await self.close_all() + + # ── background sweeper ─────────────────────────────────────── + + async def _sweep_loop(self) -> None: + """Periodically pause idle workspaces. + + Runs forever until cancelled. Each tick pops every cache entry + whose last-access is older than ``ttl`` and closes it outside + the lock; exceptions during close are logged and swallowed so + one bad sandbox does not poison the sweeper. + """ + while True: + try: + await asyncio.sleep(self._sweep_interval) + except asyncio.CancelledError: + return + try: + await self._sweep_once() + except Exception: + logger.exception("E2B workspace sweeper tick failed") + + async def _sweep_once(self) -> None: + """One sweeper tick: evict expired entries and close them.""" + now = time.monotonic() + async with self._lock: + expired_ids = [ + wid + for wid, (_, ts) in self._cache.items() + if now - ts > self._ttl + ] + evicted = [self._cache.pop(wid)[0] for wid in expired_ids] + if not evicted: + return + await asyncio.gather( + *(self._safe_close(ws) for ws in evicted), + return_exceptions=True, + ) + + @staticmethod + async def _safe_close(ws: E2BWorkspace) -> None: + """Close a workspace, logging any failure instead of raising.""" + try: + await ws.close() + except Exception: + logger.exception( + "Failed to close E2BWorkspace %s", + ws.workspace_id, + ) diff --git a/src/agentscope/app/workspace_manager/_local_workspace_manager.py b/src/agentscope/app/workspace_manager/_local_workspace_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..e437d13c75e62e4f920b72b4ca49e029902f0a9f --- /dev/null +++ b/src/agentscope/app/workspace_manager/_local_workspace_manager.py @@ -0,0 +1,184 @@ +# -*- coding: utf-8 -*- +"""The local workspace manager.""" + +import asyncio +import os +import time + +from ..._logging import logger +from ...workspace import LocalWorkspace +from ._base import WorkspaceManagerBase + + +class LocalWorkspaceManager(WorkspaceManagerBase): + """Manages LocalWorkspace instances with TTL-based lazy lifecycle. + + Workspaces are keyed by ``workspace_id`` in the cache. On cache miss + the manager reconstructs the workspace from ``basedir/agent_id`` — the + workdir is deterministic for local workspaces so no storage lookup is + needed. + """ + + def __init__( + self, + basedir: str, + default_mcps: list | None = None, + skill_paths: list[str] | None = None, + ttl: float = 3600.0, + ) -> None: + """Initialize the local workspace manager. + + Args: + basedir (`str`): + Root directory under which per-agent workdir are + created. + default_mcps (`list | None`, optional): + MCP clients seeded into brand-new workspaces. + skill_paths (`list[str] | None`, optional): + Skill directories seeded into brand-new workspaces. + ttl (`float`, defaults to `3600.0`): + Seconds before an idle cached workspace is evicted. + """ + self._basedir = os.path.abspath(basedir) + self._default_mcps = default_mcps or [] + self._skill_paths = skill_paths or [] + self._ttl = ttl + # workspace_id → (workspace, last_access_monotonic) + self._cache: dict[str, tuple[LocalWorkspace, float]] = {} + self._lock = asyncio.Lock() + + def _pop_expired(self, now: float) -> list[LocalWorkspace]: + """Pop every cache entry whose last-access exceeds ``ttl``. + + Caller is responsible for closing the returned workspaces + *outside* the manager lock so a slow ``close()`` does not stall + unrelated ``get_workspace`` callers. + """ + expired_ids = [ + wid for wid, (_, ts) in self._cache.items() if now - ts > self._ttl + ] + return [self._cache.pop(wid)[0] for wid in expired_ids] + + async def get_workspace( + self, + user_id: str, + agent_id: str, + session_id: str, + workspace_id: str, + ) -> LocalWorkspace: + """Return an initialized workspace, reconstructing from + disk on cache miss. + + Mirrors the Docker / E2B managers' double-check pattern: a + first lock acquisition handles the cache-hit fast path and + collects expired entries; expired entries are then closed in + parallel *outside* the lock; on a miss a second acquisition + runs ``initialize()`` while holding the lock so two concurrent + cache misses for the same ``workspace_id`` cannot create two + workspaces. + """ + del user_id # accepted for interface parity; not used here + + # Phase 1: cache hit + collect expired. + async with self._lock: + now = time.monotonic() + expired = self._pop_expired(now) + cached = self._cache.get(workspace_id) + if cached is not None: + ws, _ = cached + self._cache[workspace_id] = (ws, now) + hit: LocalWorkspace | None = ws + else: + hit = None + + # Phase 2: close expired entries outside the lock, in parallel, + # so a slow stdio MCP shutdown does not block unrelated callers. + if expired: + await asyncio.gather( + *(self._safe_close(ws) for ws in expired), + return_exceptions=True, + ) + + if hit is not None: + return hit + + # Phase 3: build under the lock to prevent two concurrent + # get_workspace(workspace_id=X) calls from creating two + # workspaces for the same id. + async with self._lock: + cached = self._cache.get(workspace_id) + if cached is not None: + ws, _ = cached + self._cache[workspace_id] = (ws, time.monotonic()) + return ws + + # Workdir is deterministic for local workspaces — no storage needed + workdir = os.path.join(self._basedir, agent_id) + ws = LocalWorkspace( + workspace_id=workspace_id, + workdir=workdir, + default_mcps=self._default_mcps, + skill_paths=self._skill_paths, + ) + await ws.initialize() + self._cache[workspace_id] = (ws, time.monotonic()) + return ws + + async def create_workspace( + self, + user_id: str, + agent_id: str, + session_id: str, + ) -> LocalWorkspace: + """Create a new workspace for the given agent and return it.""" + del user_id, session_id # accepted for interface parity + + workdir = os.path.join(self._basedir, agent_id) + os.makedirs(workdir, exist_ok=True) + ws = LocalWorkspace( + workdir=workdir, + default_mcps=self._default_mcps, + skill_paths=self._skill_paths, + ) + await ws.initialize() + async with self._lock: + self._cache[ws.workspace_id] = (ws, time.monotonic()) + return ws + + async def close(self, workspace_id: str) -> None: + """Close and evict a single workspace from the cache.""" + async with self._lock: + entry = self._cache.pop(workspace_id, None) + if entry is None: + return + ws, _ = entry + await self._safe_close(ws) + + async def close_all(self) -> None: + """Close every cached workspace in parallel. + + Stdio MCP shutdown can be slow per workspace; doing it + sequentially on app shutdown produces a noticeable stall, so + we fan the calls out with :func:`asyncio.gather` (mirrors the + Docker / E2B managers). + """ + async with self._lock: + entries = list(self._cache.values()) + self._cache.clear() + if not entries: + return + await asyncio.gather( + *(self._safe_close(ws) for ws, _ in entries), + return_exceptions=True, + ) + + @staticmethod + async def _safe_close(ws: LocalWorkspace) -> None: + """Close a workspace, logging any failure instead of raising.""" + try: + await ws.close() + except Exception: + logger.exception( + "Failed to close LocalWorkspace %s", + ws.workspace_id, + ) diff --git a/src/agentscope/credential/__init__.py b/src/agentscope/credential/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f468d888d792b0e8b0c6cd4a1c2082840dd7b932 --- /dev/null +++ b/src/agentscope/credential/__init__.py @@ -0,0 +1,27 @@ +# -*- coding: utf-8 -*- +"""The credential module.""" + +from ._base import CredentialBase +from ._anthropic import AnthropicCredential +from ._dashscope import DashScopeCredential +from ._deepseek import DeepSeekCredential +from ._gemini import GeminiCredential +from ._moonshot import MoonshotCredential +from ._ollama import OllamaCredential +from ._openai import OpenAICredential +from ._xai import XAICredential +from ._factory import CredentialFactory + + +__all__ = [ + "CredentialBase", + "AnthropicCredential", + "DashScopeCredential", + "DeepSeekCredential", + "GeminiCredential", + "MoonshotCredential", + "OllamaCredential", + "OpenAICredential", + "XAICredential", + "CredentialFactory", +] diff --git a/src/agentscope/credential/_anthropic.py b/src/agentscope/credential/_anthropic.py new file mode 100644 index 0000000000000000000000000000000000000000..7a2848f7b437f6e355c51a91ee83d15381efdaec --- /dev/null +++ b/src/agentscope/credential/_anthropic.py @@ -0,0 +1,39 @@ +# -*- coding: utf-8 -*- +"""The Anthropic credential.""" +from typing import Literal, Type, TYPE_CHECKING + +from pydantic import Field, SecretStr, ConfigDict + +from ._base import CredentialBase + +if TYPE_CHECKING: + from ..model import ChatModelBase + + +class AnthropicCredential(CredentialBase): + """The Anthropic credential model.""" + + model_config = ConfigDict( + title="Anthropic API", + ) + + type: Literal["anthropic_credential"] = "anthropic_credential" + """The credential type.""" + + api_key: SecretStr = Field( + description="The Anthropic API key", + ) + """The API key.""" + + base_url: str | None = Field( + description="The base URL for the Anthropic API.", + default=None, + ) + """The base URL for the Anthropic API.""" + + @classmethod + def get_chat_model_class(cls) -> Type["ChatModelBase"]: + """Return the AnthropicChatModel class.""" + from ..model import AnthropicChatModel + + return AnthropicChatModel diff --git a/src/agentscope/credential/_base.py b/src/agentscope/credential/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..8d9fb411a5111a70472385095975a9c389982993 --- /dev/null +++ b/src/agentscope/credential/_base.py @@ -0,0 +1,96 @@ +# -*- coding: utf-8 -*- +"""The credential base class.""" +from typing import TYPE_CHECKING, Type + +from pydantic import BaseModel, Field + +from .._utils._common import _generate_id + +if TYPE_CHECKING: + from ..embedding import EmbeddingModelBase + from ..model import ChatModelBase, ModelCard + from ..tts import TTSModelBase + from ..tts._tts_model_card import TTSModelCard + + +class CredentialBase(BaseModel): + """The credential base class.""" + + id: str = Field( + default_factory=_generate_id, + description="The credential id", + ) + + name: str = Field( + default="", + description="User-facing display name for this credential.", + ) + + @classmethod + def get_chat_model_class(cls) -> Type["ChatModelBase"]: + """Return the :class:`ChatModelBase` subclass that consumes this + credential. Subclasses must override this method to return the + corresponding chat model class. + + Returns: + `Type[ChatModelBase]`: + The chat model class that uses this credential. + """ + raise NotImplementedError( + f"{cls.__name__} must implement ``get_chat_model_class``.", + ) + + @classmethod + def get_tts_model_classes(cls) -> list[Type["TTSModelBase"]]: + """Return the TTS model classes supported by this credential. + + Subclasses that support TTS should override this to return one or + more :class:`TTSModelBase` subclasses. The default returns an empty + list (provider does not support TTS). + + Returns: + `list[Type[TTSModelBase]]`: + The TTS model classes, or an empty list. + """ + return [] + + @classmethod + def list_tts_models(cls) -> list["TTSModelCard"]: + """List the candidate TTS models available under this credential. + + Returns: + `list[TTSModelCard]`: + A list of TTS model cards, or empty if TTS is not supported. + """ + cards: list["TTSModelCard"] = [] + for tts_cls in cls.get_tts_model_classes(): + cards.extend(tts_cls.list_models()) + return cards + + @classmethod + def list_models(cls) -> list["ModelCard"]: + """List the candidate chat models that are available under this + credential. The default implementation delegates to the + :meth:`ChatModelBase.list_models` of the class returned by + :meth:`get_chat_model_class`. + + Returns: + `list[ModelCard]`: + A list of candidate models described by their model cards. + """ + return cls.get_chat_model_class().list_models() + + @classmethod + def get_embedding_model_class(cls) -> Type["EmbeddingModelBase"] | None: + """Return the :class:`EmbeddingModelBase` subclass that consumes + this credential, or ``None`` if this provider does not support + embedding models. + + Subclasses that have a matching embedding implementation should + override this method. The default returns ``None``. + + Returns: + `Type[EmbeddingModelBase] | None`: + The embedding model class, or ``None``. + """ + return None diff --git a/src/agentscope/credential/_dashscope.py b/src/agentscope/credential/_dashscope.py new file mode 100644 index 0000000000000000000000000000000000000000..3521e3d8bf2cb37c833021d4837a7c24d6d31d6d --- /dev/null +++ b/src/agentscope/credential/_dashscope.py @@ -0,0 +1,67 @@ +# -*- coding: utf-8 -*- +"""The DashScope credential.""" +from typing import Literal, Type, TYPE_CHECKING + +from pydantic import ConfigDict, Field, SecretStr + +from ._base import CredentialBase + +if TYPE_CHECKING: + from ..embedding import EmbeddingModelBase + from ..model import ChatModelBase + from ..tts import TTSModelBase + +_DASHSCOPE_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1" + + +class DashScopeCredential(CredentialBase): + """The credential for DashScope API.""" + + model_config = ConfigDict( + title="DashScope API", + ) + + type: Literal["dashscope_credential"] = "dashscope_credential" + """The type of the credential.""" + + api_key: SecretStr = Field( + description="The DashScope API key.", + title="API Key", + ) + + base_url: str = Field( + default=_DASHSCOPE_BASE_URL, + title="API Base URL", + description=( + "The base URL for the DashScope OpenAI-compatible API endpoint." + ), + ) + + @classmethod + def get_chat_model_class(cls) -> Type["ChatModelBase"]: + """Return the DashScopeChatModel class.""" + from ..model import DashScopeChatModel + + return DashScopeChatModel + + @classmethod + def get_tts_model_classes(cls) -> list[Type["TTSModelBase"]]: + """Return the DashScope TTS model classes.""" + from ..tts import ( + DashScopeCosyVoiceRealtimeTTSModel, + DashScopeRealtimeTTSModel, + DashScopeTTSModel, + ) + + return [ + DashScopeTTSModel, + DashScopeRealtimeTTSModel, + DashScopeCosyVoiceRealtimeTTSModel, + ] + + @classmethod + def get_embedding_model_class(cls) -> Type["EmbeddingModelBase"]: + """Return the DashScopeEmbeddingModel class.""" + from ..embedding import DashScopeEmbeddingModel + + return DashScopeEmbeddingModel diff --git a/src/agentscope/credential/_deepseek.py b/src/agentscope/credential/_deepseek.py new file mode 100644 index 0000000000000000000000000000000000000000..35511fecc8656afea405d666621676a8361c2ad4 --- /dev/null +++ b/src/agentscope/credential/_deepseek.py @@ -0,0 +1,41 @@ +# -*- coding: utf-8 -*- +"""The DeepSeek credential.""" +from typing import Literal, Type, TYPE_CHECKING + +from pydantic import ConfigDict, Field, SecretStr + +from ._base import CredentialBase + +if TYPE_CHECKING: + from ..model import ChatModelBase + +_DEEPSEEK_BASE_URL = "https://api.deepseek.com" + + +class DeepSeekCredential(CredentialBase): + """The DeepSeek credential model.""" + + model_config = ConfigDict( + title="DeepSeek API", + ) + + type: Literal["deepseek_credential"] = "deepseek_credential" + """The credential type.""" + + api_key: SecretStr = Field( + description="The DeepSeek API key.", + ) + """The API key.""" + + base_url: str = Field( + default=_DEEPSEEK_BASE_URL, + description="The base URL for the DeepSeek API.", + ) + """The base URL for the DeepSeek API.""" + + @classmethod + def get_chat_model_class(cls) -> Type["ChatModelBase"]: + """Return the DeepSeekChatModel class.""" + from ..model import DeepSeekChatModel + + return DeepSeekChatModel diff --git a/src/agentscope/credential/_factory.py b/src/agentscope/credential/_factory.py new file mode 100644 index 0000000000000000000000000000000000000000..bb16f93902d9db1f13ac7eec3d15f01971598db1 --- /dev/null +++ b/src/agentscope/credential/_factory.py @@ -0,0 +1,116 @@ +# -*- coding: utf-8 -*- +"""The credential factory class.""" +from typing import Annotated, Type, Union, get_args, get_type_hints + +from pydantic import TypeAdapter, Field + +from ._anthropic import AnthropicCredential +from ._dashscope import DashScopeCredential +from ._deepseek import DeepSeekCredential +from ._gemini import GeminiCredential +from ._moonshot import MoonshotCredential +from ._ollama import OllamaCredential +from ._openai import OpenAICredential +from ._xai import XAICredential +from ._base import CredentialBase + + +class CredentialFactory: + """Registry and deserializer for :class:`CredentialBase` subclasses. + + Built-in credential types are pre-registered. Call + :meth:`register_credential` to add custom types before starting the app. + + Usage:: + + # Deserialize from storage + credential = CredentialFactory.from_dict(record.data) + + # Register a custom type + CredentialFactory.register_credential(MyCredential) + + # List schemas for the frontend form + schemas = CredentialFactory.list_schemas() + """ + + _classes: list[Type[CredentialBase]] = [ + AnthropicCredential, + DashScopeCredential, + DeepSeekCredential, + GeminiCredential, + MoonshotCredential, + OllamaCredential, + OpenAICredential, + XAICredential, + ] + _adapter: TypeAdapter[CredentialBase] | None = None + + @classmethod + def _get_adapter(cls) -> TypeAdapter[CredentialBase]: + if cls._adapter is None: + union = Annotated[ # type: ignore[valid-type] + Union[tuple(cls._classes)], + Field(discriminator="type"), + ] + cls._adapter = TypeAdapter(union) + return cls._adapter + + @classmethod + def register_credential(cls, credential_cls: Type[CredentialBase]) -> None: + """Register a custom :class:`CredentialBase` subclass. + + The class must define a ``type`` field with a unique ``Literal`` + default so Pydantic can use it as a discriminator. + + Args: + credential_cls: The subclass to register. + """ + if credential_cls in cls._classes: + return + cls._classes.append(credential_cls) + cls._adapter = None # invalidate so it's rebuilt on next use + + @classmethod + def from_dict(cls, data: dict) -> CredentialBase: + """Deserialize a credential dict (from storage) to a typed instance. + + Args: + data: Raw dict containing a ``"type"`` key. + + Returns: + A typed :class:`CredentialBase` subclass instance. + """ + return cls._get_adapter().validate_python(data) + + @classmethod + def get_credential_class( + cls, + provider: str, + ) -> Type[CredentialBase] | None: + """Return the credential class for the given provider type, or None. + + Args: + provider: The ``type`` discriminator value (e.g. ``"openai"``). + + Returns: + The matching :class:`CredentialBase` subclass, or ``None`` if not + found. + """ + for c in cls._classes: + hints = get_type_hints(c) + type_hint = hints.get("type") + if type_hint is None: + continue + args = get_args(type_hint) + + if args and args[0] == provider: + return c + return None + + @classmethod + def list_schemas(cls) -> list[dict]: + """Return JSON schemas for all registered credential types. + + Used by the frontend to render credential forms dynamically. + """ + return [c.model_json_schema() for c in cls._classes] diff --git a/src/agentscope/credential/_gemini.py b/src/agentscope/credential/_gemini.py new file mode 100644 index 0000000000000000000000000000000000000000..f65804ddd964554517bc4950a2cefb77b788df09 --- /dev/null +++ b/src/agentscope/credential/_gemini.py @@ -0,0 +1,41 @@ +# -*- coding: utf-8 -*- +"""The Google Gemini credential.""" +from typing import Literal, Type, TYPE_CHECKING + +from pydantic import ConfigDict, Field, SecretStr + +from ._base import CredentialBase + +if TYPE_CHECKING: + from ..embedding import EmbeddingModelBase + from ..model import ChatModelBase + + +class GeminiCredential(CredentialBase): + """The Google Gemini credential model.""" + + model_config = ConfigDict( + title="Gemini API", + ) + + type: Literal["gemini_credential"] = "gemini_credential" + """The credential type.""" + + api_key: SecretStr = Field( + description="The Google Gemini API key.", + ) + """The API key.""" + + @classmethod + def get_chat_model_class(cls) -> Type["ChatModelBase"]: + """Return the GeminiChatModel class.""" + from ..model import GeminiChatModel + + return GeminiChatModel + + @classmethod + def get_embedding_model_class(cls) -> Type["EmbeddingModelBase"]: + """Return the GeminiEmbeddingModel class.""" + from ..embedding import GeminiEmbeddingModel + + return GeminiEmbeddingModel diff --git a/src/agentscope/credential/_kimi.py b/src/agentscope/credential/_kimi.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/agentscope/credential/_moonshot.py b/src/agentscope/credential/_moonshot.py new file mode 100644 index 0000000000000000000000000000000000000000..c5fd374bae4673e93c334287a6905426ffe2417a --- /dev/null +++ b/src/agentscope/credential/_moonshot.py @@ -0,0 +1,41 @@ +# -*- coding: utf-8 -*- +"""The Moonshot AI credential.""" +from typing import Literal, Type, TYPE_CHECKING + +from pydantic import ConfigDict, Field, SecretStr + +from ._base import CredentialBase + +if TYPE_CHECKING: + from ..model import ChatModelBase + +_MOONSHOT_BASE_URL = "https://api.moonshot.cn/v1" + + +class MoonshotCredential(CredentialBase): + """The Moonshot AI credential model.""" + + model_config = ConfigDict( + title="Moonshot API", + ) + + type: Literal["moonshot_credential"] = "moonshot_credential" + """The credential type.""" + + api_key: SecretStr = Field( + description="The Moonshot AI API key.", + ) + """The API key.""" + + base_url: str = Field( + default=_MOONSHOT_BASE_URL, + description="The base URL for the Moonshot AI API.", + ) + """The base URL for the Moonshot AI API.""" + + @classmethod + def get_chat_model_class(cls) -> Type["ChatModelBase"]: + """Return the MoonshotChatModel class.""" + from ..model import MoonshotChatModel + + return MoonshotChatModel diff --git a/src/agentscope/credential/_ollama.py b/src/agentscope/credential/_ollama.py new file mode 100644 index 0000000000000000000000000000000000000000..2ce1d2e761808b4b2588b72aebad9ac3250c14f9 --- /dev/null +++ b/src/agentscope/credential/_ollama.py @@ -0,0 +1,45 @@ +# -*- coding: utf-8 -*- +"""The Ollama credential.""" +from typing import Literal, Type, TYPE_CHECKING + +from pydantic import ConfigDict, Field + +from ._base import CredentialBase + +if TYPE_CHECKING: + from ..embedding import EmbeddingModelBase + from ..model import ChatModelBase + + +class OllamaCredential(CredentialBase): + """The Ollama credential model (connection settings).""" + + model_config = ConfigDict( + title="Ollama API", + ) + + type: Literal["ollama_credential"] = "ollama_credential" + """The credential type.""" + + host: str | None = Field( + default=None, + description=( + "The Ollama server host URL. " + "Defaults to http://localhost:11434 if not specified." + ), + ) + """The Ollama server host URL.""" + + @classmethod + def get_chat_model_class(cls) -> Type["ChatModelBase"]: + """Return the OllamaChatModel class.""" + from ..model import OllamaChatModel + + return OllamaChatModel + + @classmethod + def get_embedding_model_class(cls) -> Type["EmbeddingModelBase"]: + """Return the OllamaEmbeddingModel class.""" + from ..embedding import OllamaEmbeddingModel + + return OllamaEmbeddingModel diff --git a/src/agentscope/credential/_openai.py b/src/agentscope/credential/_openai.py new file mode 100644 index 0000000000000000000000000000000000000000..f0848e0fcea77750c7e0ea7ea6732a75f57c03b8 --- /dev/null +++ b/src/agentscope/credential/_openai.py @@ -0,0 +1,56 @@ +# -*- coding: utf-8 -*- +"""The OpenAI credential.""" +from typing import Literal, Type, TYPE_CHECKING + +from pydantic import ConfigDict, Field, SecretStr + +from ._base import CredentialBase + +if TYPE_CHECKING: + from ..embedding import EmbeddingModelBase + from ..model import ChatModelBase + + +class OpenAICredential(CredentialBase): + """The OpenAI credential model.""" + + model_config = ConfigDict( + title="OpenAI API", + ) + + type: Literal["openai_credential"] = "openai_credential" + """The credential type.""" + + api_key: SecretStr = Field( + description="The OpenAI API key.", + ) + """The API key.""" + + organization: str | None = Field( + default=None, + description="The OpenAI organization ID.", + ) + """The OpenAI organization ID.""" + + base_url: str | None = Field( + default=None, + description=( + "The base URL for the OpenAI API. " + "Can be used for OpenAI-compatible endpoints." + ), + ) + """Custom base URL for OpenAI-compatible endpoints.""" + + @classmethod + def get_chat_model_class(cls) -> Type["ChatModelBase"]: + """Return the OpenAIChatModel class.""" + from ..model import OpenAIChatModel + + return OpenAIChatModel + + @classmethod + def get_embedding_model_class(cls) -> Type["EmbeddingModelBase"]: + """Return the OpenAIEmbeddingModel class.""" + from ..embedding import OpenAIEmbeddingModel + + return OpenAIEmbeddingModel diff --git a/src/agentscope/credential/_xai.py b/src/agentscope/credential/_xai.py new file mode 100644 index 0000000000000000000000000000000000000000..f3001f6e73608aaefb8dee10ee6ca9674db375e2 --- /dev/null +++ b/src/agentscope/credential/_xai.py @@ -0,0 +1,43 @@ +# -*- coding: utf-8 -*- +"""The xAI credential.""" +from typing import Literal, Type, TYPE_CHECKING + +from pydantic import ConfigDict, Field, SecretStr + +from ._base import CredentialBase + +if TYPE_CHECKING: + from ..model import ChatModelBase + + +class XAICredential(CredentialBase): + """The xAI credential model.""" + + model_config = ConfigDict( + title="xAI API", + ) + + type: Literal["xai_credential"] = "xai_credential" + """The credential type.""" + + api_key: SecretStr = Field( + description="The xAI API key.", + ) + """The xAI API key.""" + + api_host: str = Field( + default="api.x.ai", + title="API Host", + description=( + "The xAI API host (without scheme). Override to point at a " + "compatible/self-hosted endpoint." + ), + ) + """The xAI API host.""" + + @classmethod + def get_chat_model_class(cls) -> Type["ChatModelBase"]: + """Return the XAIChatModel class.""" + from ..model import XAIChatModel + + return XAIChatModel diff --git a/src/agentscope/embedding/__init__.py b/src/agentscope/embedding/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e5910a9253d9e1de8479ef55a91c382d5e5ada58 --- /dev/null +++ b/src/agentscope/embedding/__init__.py @@ -0,0 +1,27 @@ +# -*- coding: utf-8 -*- +"""The embedding module in agentscope.""" + +from ._embedding_base import EmbeddingModelBase +from ._embedding_model_card import EmbeddingModelCard +from ._embedding_usage import EmbeddingUsage +from ._embedding_response import EmbeddingResponse +from ._dashscope import DashScopeEmbeddingModel +from ._openai import OpenAIEmbeddingModel +from ._gemini import GeminiEmbeddingModel +from ._ollama import OllamaEmbeddingModel +from ._cache_base import EmbeddingCacheBase +from ._file_cache import FileEmbeddingCache + + +__all__ = [ + "EmbeddingModelBase", + "EmbeddingModelCard", + "EmbeddingUsage", + "EmbeddingResponse", + "DashScopeEmbeddingModel", + "OpenAIEmbeddingModel", + "GeminiEmbeddingModel", + "OllamaEmbeddingModel", + "EmbeddingCacheBase", + "FileEmbeddingCache", +] diff --git a/src/agentscope/embedding/_cache_base.py b/src/agentscope/embedding/_cache_base.py new file mode 100644 index 0000000000000000000000000000000000000000..5d3fa1addcaf75546e3148c835d0f0eba26f11d1 --- /dev/null +++ b/src/agentscope/embedding/_cache_base.py @@ -0,0 +1,63 @@ +# -*- coding: utf-8 -*- +"""The embedding cache base class.""" +from abc import abstractmethod +from typing import List, Any + +from ..types import ( + JSONSerializableObject, + Embedding, +) + + +class EmbeddingCacheBase: + """Base class for embedding caches, which is responsible for storing and + retrieving embeddings.""" + + @abstractmethod + async def store( + self, + embeddings: List[Embedding], + identifier: JSONSerializableObject, + overwrite: bool = False, + **kwargs: Any, + ) -> None: + """Store the embeddings with the given identifier. + + Args: + embeddings (`List[Embedding]`): + The embeddings to store. + identifier (`JSONSerializableObject`): + The identifier to distinguish the embeddings. + overwrite (`bool`, defaults to `False`): + Whether to overwrite existing embeddings with the same + identifier. If `True`, existing embeddings will be replaced. + """ + + @abstractmethod + async def retrieve( + self, + identifier: JSONSerializableObject, + ) -> List[Embedding] | None: + """Retrieve the embeddings with the given identifier. If not + found, return `None`. + + Args: + identifier (`JSONSerializableObject`): + The identifier to retrieve the embeddings. + """ + + @abstractmethod + async def remove( + self, + identifier: JSONSerializableObject, + ) -> None: + """Remove the embeddings with the given identifier. + + Args: + identifier (`JSONSerializableObject`): + The identifier to remove the embeddings. + """ + + @abstractmethod + async def clear(self) -> None: + """Clear all cached embeddings.""" diff --git a/src/agentscope/embedding/_dashscope/__init__.py b/src/agentscope/embedding/_dashscope/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5cf3d4e2d4ea7c864fd8533ec3ec9977a895d3c9 --- /dev/null +++ b/src/agentscope/embedding/_dashscope/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- +"""The DashScope embedding API modules.""" + +from ._model import DashScopeEmbeddingModel + +__all__ = [ + "DashScopeEmbeddingModel", +] diff --git a/src/agentscope/embedding/_dashscope/_model.py b/src/agentscope/embedding/_dashscope/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..d0a8ca585dd93b24c83c237dfc31b6980fcb565d --- /dev/null +++ b/src/agentscope/embedding/_dashscope/_model.py @@ -0,0 +1,529 @@ +# -*- coding: utf-8 -*- +"""The DashScope embedding model. + +Handles both text-only and multimodal models under a single class. +Text models (``text-embedding-v3``, ``text-embedding-v4``) accept +``list[str | TextBlock]``. Multimodal models (``qwen*-vl-embedding``, +``multimodal-embedding-*``, ``tongyi-embedding-vision-*``) +additionally accept :class:`~agentscope.message.DataBlock`. +The model name determines which DashScope API endpoint is used. + +Text payloads may be passed either as bare ``str`` or as +:class:`~agentscope.message.TextBlock` — the latter is unpacked to its +``.text`` field on entry so the rest of the pipeline only deals with +``str`` and ``DataBlock``. +""" +from __future__ import annotations + +import asyncio +from datetime import datetime +from dataclasses import dataclass +from typing import Any + +from .._cache_base import EmbeddingCacheBase +from .._embedding_response import EmbeddingResponse +from .._embedding_usage import EmbeddingUsage +from .._embedding_base import EmbeddingModelBase +from ..._logging import logger +from ...credential import CredentialBase +from ...message import DataBlock, Base64Source, TextBlock, URLSource + +#: Model name prefixes that route to the multimodal API. +_MULTIMODAL_PREFIXES = ( + "multimodal-embedding-", + "tongyi-embedding-vision-", + "qwen3-vl-embedding", + "qwen2.5-vl-embedding", +) + + +@dataclass +class _MultimodalLimits: + """Per-request constraints for a multimodal embedding model.""" + + max_elements: int = 20 + """Maximum total content elements per API call.""" + + max_images: int = 5 + """Maximum image elements per API call.""" + + max_videos: int = 1 + """Maximum video elements per API call.""" + + +#: Known per-model multimodal constraints (from DashScope docs). +_MODEL_LIMITS: dict[str, _MultimodalLimits] = { + "qwen3-vl-embedding": _MultimodalLimits( + max_elements=20, + max_images=5, + max_videos=1, + ), + "qwen2.5-vl-embedding": _MultimodalLimits( + max_elements=20, + max_images=5, + max_videos=1, + ), + "tongyi-embedding-vision-plus": _MultimodalLimits( + max_elements=20, + max_images=64, + max_videos=8, + ), + "tongyi-embedding-vision-flash": _MultimodalLimits( + max_elements=20, + max_images=64, + max_videos=8, + ), + "multimodal-embedding-v1": _MultimodalLimits( + max_elements=20, + max_images=1, + max_videos=1, + ), +} + +#: Fallback for unknown multimodal models — safest constraints. +_DEFAULT_LIMITS = _MultimodalLimits( + max_elements=20, + max_images=1, + max_videos=1, +) + + +class DashScopeEmbeddingModel(EmbeddingModelBase[str | TextBlock | DataBlock]): + """Unified DashScope embedding model. + + Routes to the text or multimodal DashScope API based on the model + name. + + - **Text mode** (``text-embedding-*``): uses the base class's + simple batch splitting + concurrent retry. + - **Multimodal mode** (``qwen*-vl-*``, ``multimodal-*``, + ``tongyi-embedding-vision-*``): overrides ``__call__`` to + perform content-aware batching that respects per-model limits + on total elements, images, and videos per request. + """ + + #: Text-mode batch size (from DashScope docs: 10 for v3/v4, 25 for + #: v1/v2). Multimodal models use :data:`_MODEL_LIMITS` instead. + _TEXT_BATCH_SIZE: int = 10 + + def __init__( + self, + credential: CredentialBase, + model: str, + dimensions: int | None, + parameters: "DashScopeEmbeddingModel.Parameters | None" = None, + embedding_cache: EmbeddingCacheBase | None = None, + context_size: int = 8192, + max_retries: int = 3, + retry_delay: float = 1.0, + ) -> None: + """Initialize the DashScope embedding model. + + Args: + credential (`CredentialBase`): + A :class:`~agentscope.credential.DashScopeCredential` + instance providing the API key. + model (`str`): + The embedding model name (e.g. + ``"text-embedding-v4"`` or + ``"qwen3-vl-embedding"``). + dimensions (`int | None`): + The output embedding vector dimensions. Required at + the contract level — see :class:`EmbeddingModelBase` + for the rationale. ``None`` is accepted only for + backward compatibility with legacy configs that + persisted ``dimensions`` inside ``parameters``. + parameters (`DashScopeEmbeddingModel.Parameters | None`, \ + defaults to ``None``): + Provider-specific non-dimensional parameters. Currently + empty for DashScope. + embedding_cache (`EmbeddingCacheBase | None`, defaults to \ + ``None``): + Optional embedding cache. + context_size (`int`, defaults to ``8192``): + Maximum input tokens per text. + max_retries (`int`, defaults to ``3``): + Number of retries on transient failures. + retry_delay (`float`, defaults to ``1.0``): + Seconds between retry attempts. + """ + self._is_multimodal: bool = model.startswith(_MULTIMODAL_PREFIXES) + + super().__init__( + credential=credential, + model=model, + dimensions=dimensions, + parameters=parameters, + context_size=context_size, + batch_size=self._TEXT_BATCH_SIZE, + max_retries=max_retries, + retry_delay=retry_delay, + ) + self.api_key: str = credential.api_key.get_secret_value() + self.embedding_cache: EmbeddingCacheBase | None = embedding_cache + + # Resolve multimodal constraints. + if self._is_multimodal: + self._limits = _MODEL_LIMITS.get(model, _DEFAULT_LIMITS) + + @classmethod + def _get_retryable_exceptions(cls) -> tuple[type[Exception], ...]: + """Return retryable exception types. + + DashScope SDK does not expose typed exception classes. We + retry on ``RuntimeError``, which is raised when the API + returns a non-200 status code. + """ + return (RuntimeError,) + + # ------------------------------------------------------------------ + # __call__ — override for multimodal content-aware batching + # ------------------------------------------------------------------ + + async def __call__( + self, + inputs: list[str | TextBlock | DataBlock], + **kwargs: Any, + ) -> EmbeddingResponse: + """Embed inputs with batching and retry. + + For text models, delegates to the base class (simple + ``batch_size`` splitting). For multimodal models, performs + content-aware batching that respects per-model limits on + total elements, images, and videos per request. + + Args: + inputs (`list[str | TextBlock | DataBlock]`): + The input data to embed. ``TextBlock`` items are + unpacked to their ``.text`` field on entry, so the + remainder of the pipeline only sees ``str`` and + ``DataBlock``. + **kwargs: + Forwarded to the DashScope API. + + Returns: + `EmbeddingResponse`: Merged response for all inputs. + """ + normalized: list[str | DataBlock] = [ + item.text if isinstance(item, TextBlock) else item + for item in inputs + ] + + if not self._is_multimodal: + # Text mode — use base class batching. + return await super().__call__(normalized, **kwargs) + + # Multimodal mode — content-aware batching. + batches = self._split_multimodal_batches(normalized) + + if len(batches) > 1: + logger.info( + "Embedding %d multimodal inputs in %d batches for " + "model %s (limits: elements=%d, images=%d, videos=%d).", + len(normalized), + len(batches), + self.model, + self._limits.max_elements, + self._limits.max_images, + self._limits.max_videos, + ) + + results: list[EmbeddingResponse] = await asyncio.gather( + *(self._call_with_retry(batch, **kwargs) for batch in batches), + ) + + return self._merge_responses(results) + + def _split_multimodal_batches( + self, + inputs: list[str | DataBlock], + ) -> list[list[str | DataBlock]]: + """Split inputs into batches that satisfy multimodal limits. + + Greedy algorithm: keep adding items to the current batch + until adding the next item would violate any constraint, + then start a new batch. + + Args: + inputs (`list[str | DataBlock]`): + All inputs to split. + + Returns: + `list[list[str | DataBlock]]`: List of batches. + """ + limits = self._limits + batches: list[list[str | DataBlock]] = [] + current_batch: list[str | DataBlock] = [] + n_elements = 0 + n_images = 0 + n_videos = 0 + + for item in inputs: + # Determine what this item contributes. + is_image = False + is_video = False + if isinstance(item, DataBlock): + media_type = item.source.media_type + is_image = media_type.startswith("image/") + is_video = media_type.startswith("video/") + + # Check if adding this item would exceed any limit. + would_exceed = ( + n_elements + 1 > limits.max_elements + or (is_image and n_images + 1 > limits.max_images) + or (is_video and n_videos + 1 > limits.max_videos) + ) + + if would_exceed and current_batch: + batches.append(current_batch) + current_batch = [] + n_elements = 0 + n_images = 0 + n_videos = 0 + + current_batch.append(item) + n_elements += 1 + if is_image: + n_images += 1 + if is_video: + n_videos += 1 + + if current_batch: + batches.append(current_batch) + + return batches + + # ------------------------------------------------------------------ + # _call_api — single batch dispatch + # ------------------------------------------------------------------ + + async def _call_api( + self, + inputs: list[str | DataBlock], + **kwargs: Any, + ) -> EmbeddingResponse: + """Route to the text or multimodal DashScope API. + + Args: + inputs (`list[str | DataBlock]`): + A single batch. For text models every element must be + ``str``; for multimodal models elements may also be + :class:`~agentscope.message.DataBlock`. + **kwargs: + Forwarded to the DashScope API. + + Returns: + `EmbeddingResponse`: Embedding vectors and usage info. + """ + if self._is_multimodal: + return await self._call_multimodal(inputs, **kwargs) + return await self._call_text(inputs, **kwargs) + + # ------------------------------------------------------------------ + # Text API + # ------------------------------------------------------------------ + + async def _call_text( + self, + inputs: list[str | DataBlock], + **kwargs: Any, + ) -> EmbeddingResponse: + """Call the DashScope text embedding API for a single batch. + + Args: + inputs (`list[str | DataBlock]`): + Must all be ``str``; raises ``ValueError`` otherwise. + **kwargs: + Forwarded to the API. + + Returns: + `EmbeddingResponse`: Embedding vectors and usage info. + """ + texts: list[str] = [] + for item in inputs: + if not isinstance(item, str): + raise ValueError( + f"Text embedding model {self.model!r} only accepts " + f"str inputs, got {type(item).__name__}.", + ) + texts.append(item) + + api_kwargs: dict[str, Any] = { + "input": texts, + "model": self.model, + "dimension": self.dimensions, + **kwargs, + } + + if self.embedding_cache: + cached = await self.embedding_cache.retrieve( + identifier=api_kwargs, + ) + if cached: + return EmbeddingResponse( + embeddings=cached, + usage=EmbeddingUsage(tokens=0, time=0), + source="cache", + ) + + import dashscope + + start_time = datetime.now() + response = dashscope.embeddings.TextEmbedding.call( + api_key=self.api_key, + **api_kwargs, + ) + time = (datetime.now() - start_time).total_seconds() + + if response.status_code != 200: + raise RuntimeError( + f"DashScope text embedding API error: {response}", + ) + + embeddings = [ + entry["embedding"] for entry in response.output["embeddings"] + ] + if self.embedding_cache: + await self.embedding_cache.store( + identifier=api_kwargs, + embeddings=embeddings, + ) + return EmbeddingResponse( + embeddings=embeddings, + usage=EmbeddingUsage( + tokens=response.usage["total_tokens"], + time=time, + ), + ) + + # ------------------------------------------------------------------ + # Multimodal API + # ------------------------------------------------------------------ + + async def _call_multimodal( + self, + inputs: list[str | DataBlock], + **kwargs: Any, + ) -> EmbeddingResponse: + """Call the DashScope multimodal embedding API for a single batch. + + Args: + inputs (`list[str | DataBlock]`): + ``str`` for text, ``DataBlock`` for images / videos. + **kwargs: + Forwarded to the API. + + Returns: + `EmbeddingResponse`: Embedding vectors and usage info. + """ + formatted: list[dict[str, str]] = [] + for item in inputs: + if isinstance(item, str): + formatted.append({"text": item}) + elif isinstance(item, DataBlock): + formatted.append(self._format_data_block(item)) + else: + raise ValueError( + f"Invalid input: {item!r}. Expected str or DataBlock.", + ) + + api_kwargs: dict[str, Any] = { + "input": formatted, + "model": self.model, + "api_key": self.api_key, + **kwargs, + } + + # Exclude api_key from cache identifier to avoid persisting secrets + # and to keep cache valid across key rotations. + cache_identifier = { + k: v for k, v in api_kwargs.items() if k != "api_key" + } + + if self.embedding_cache: + cached = await self.embedding_cache.retrieve( + identifier=cache_identifier, + ) + if cached: + return EmbeddingResponse( + embeddings=cached, + usage=EmbeddingUsage(tokens=0, time=0), + source="cache", + ) + + import dashscope + + start_time = datetime.now() + res = dashscope.MultiModalEmbedding.call(**api_kwargs) + time = (datetime.now() - start_time).total_seconds() + + if res.status_code != 200: + raise RuntimeError( + f"DashScope multimodal embedding API error: {res}", + ) + + embeddings = [entry["embedding"] for entry in res.output["embeddings"]] + if self.embedding_cache: + await self.embedding_cache.store( + identifier=cache_identifier, + embeddings=embeddings, + ) + return EmbeddingResponse( + embeddings=embeddings, + usage=EmbeddingUsage( + tokens=res.usage.get("image_tokens", 0) + + res.usage.get("input_tokens", 0), + time=time, + ), + source="api", + ) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + @staticmethod + def _format_data_block(block: DataBlock) -> dict[str, str]: + """Convert a :class:`~agentscope.message.DataBlock` to the dict + format expected by the DashScope multimodal embedding API. + + The ``DataBlock.source.media_type`` determines whether the + block is treated as an image or video. + + Args: + block (`DataBlock`): + A data block with a ``Base64Source`` or ``URLSource``. + + Returns: + `dict[str, str]`: E.g. + ``{"image": "data:image/png;base64,..."}`` or + ``{"video": "https://..."}``. + + Raises: + `ValueError`: If the media type is unsupported or a video + block uses a non-URL source. + """ + + source = block.source + media_type = source.media_type + + if media_type.startswith("video/"): + if not isinstance(source, URLSource): + raise ValueError( + "Multimodal embedding API only supports URL input " + f"for video data, got {type(source).__name__}.", + ) + return {"video": str(source.url)} + + if media_type.startswith("image/"): + if isinstance(source, Base64Source): + return { + "image": f"data:{media_type};" f"base64,{source.data}", + } + if isinstance(source, URLSource): + return {"image": str(source.url)} + + raise ValueError( + f"Unsupported media type {media_type!r} in DataBlock. " + f"Expected image/* or video/*.", + ) diff --git a/src/agentscope/embedding/_dashscope/_models/multimodal-embedding-v1.yaml b/src/agentscope/embedding/_dashscope/_models/multimodal-embedding-v1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0c0246bbb4f7ec8473ff622f69fee7b52d92f8fe --- /dev/null +++ b/src/agentscope/embedding/_dashscope/_models/multimodal-embedding-v1.yaml @@ -0,0 +1,16 @@ +name: multimodal-embedding-v1 +label: Multimodal Embedding v1 +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/bmp + +output_types: + - application/x-embedding + +context_size: 512 + +dimensions: 1024 diff --git a/src/agentscope/embedding/_dashscope/_models/qwen2.5-vl-embedding.yaml b/src/agentscope/embedding/_dashscope/_models/qwen2.5-vl-embedding.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c57b4a934b60931e20645fb8b5c46bc5e49f42db --- /dev/null +++ b/src/agentscope/embedding/_dashscope/_models/qwen2.5-vl-embedding.yaml @@ -0,0 +1,20 @@ +name: qwen2.5-vl-embedding +label: Qwen2.5 VL Embedding +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/webp + - image/bmp + - image/tiff + - video/mp4 + +output_types: + - application/x-embedding + +context_size: 32000 + +dimensions: 1024 +supported_dimensions: [2048, 1024, 768, 512] diff --git a/src/agentscope/embedding/_dashscope/_models/qwen3-vl-embedding.yaml b/src/agentscope/embedding/_dashscope/_models/qwen3-vl-embedding.yaml new file mode 100644 index 0000000000000000000000000000000000000000..10c30879e88cdbd2f5a11791a04949bce74065ef --- /dev/null +++ b/src/agentscope/embedding/_dashscope/_models/qwen3-vl-embedding.yaml @@ -0,0 +1,20 @@ +name: qwen3-vl-embedding +label: Qwen3 VL Embedding +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/webp + - image/bmp + - image/tiff + - video/mp4 + +output_types: + - application/x-embedding + +context_size: 32000 + +dimensions: 2560 +supported_dimensions: [2560, 2048, 1536, 1024, 768, 512, 256] diff --git a/src/agentscope/embedding/_dashscope/_models/text-embedding-v3.yaml b/src/agentscope/embedding/_dashscope/_models/text-embedding-v3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..84c3155a51ef4ec49b7626a69c1ba1f2b4e6a9e8 --- /dev/null +++ b/src/agentscope/embedding/_dashscope/_models/text-embedding-v3.yaml @@ -0,0 +1,14 @@ +name: text-embedding-v3 +label: Text Embedding v3 +status: active + +input_types: + - text/plain + +output_types: + - application/x-embedding + +context_size: 8192 + +dimensions: 1024 +supported_dimensions: [1024, 768, 512, 256, 128, 64] diff --git a/src/agentscope/embedding/_dashscope/_models/text-embedding-v4.yaml b/src/agentscope/embedding/_dashscope/_models/text-embedding-v4.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cd686084c34636a8c8f7ffe2e123933b83acc419 --- /dev/null +++ b/src/agentscope/embedding/_dashscope/_models/text-embedding-v4.yaml @@ -0,0 +1,14 @@ +name: text-embedding-v4 +label: Text Embedding v4 +status: active + +input_types: + - text/plain + +output_types: + - application/x-embedding + +context_size: 8192 + +dimensions: 1024 +supported_dimensions: [2048, 1536, 1024, 768, 512, 256, 128, 64] diff --git a/src/agentscope/embedding/_dashscope/_models/tongyi-embedding-vision-flash.yaml b/src/agentscope/embedding/_dashscope/_models/tongyi-embedding-vision-flash.yaml new file mode 100644 index 0000000000000000000000000000000000000000..51b03599efb154b74476274f12b75f42972cc545 --- /dev/null +++ b/src/agentscope/embedding/_dashscope/_models/tongyi-embedding-vision-flash.yaml @@ -0,0 +1,21 @@ +name: tongyi-embedding-vision-flash +label: Tongyi Embedding Vision Flash +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/webp + - image/bmp + - image/tiff + - video/mp4 + - video/mpeg + +output_types: + - application/x-embedding + +context_size: 1024 + +dimensions: 768 +supported_dimensions: [768, 512, 256, 128, 64] diff --git a/src/agentscope/embedding/_dashscope/_models/tongyi-embedding-vision-plus.yaml b/src/agentscope/embedding/_dashscope/_models/tongyi-embedding-vision-plus.yaml new file mode 100644 index 0000000000000000000000000000000000000000..979b88bc9608e8a5190d0ae94870c1513e232eb7 --- /dev/null +++ b/src/agentscope/embedding/_dashscope/_models/tongyi-embedding-vision-plus.yaml @@ -0,0 +1,21 @@ +name: tongyi-embedding-vision-plus +label: Tongyi Embedding Vision Plus +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/webp + - image/bmp + - image/tiff + - video/mp4 + - video/mpeg + +output_types: + - application/x-embedding + +context_size: 1024 + +dimensions: 1152 +supported_dimensions: [1152, 1024, 512, 256, 128, 64] diff --git a/src/agentscope/embedding/_embedding_base.py b/src/agentscope/embedding/_embedding_base.py new file mode 100644 index 0000000000000000000000000000000000000000..ef07bdb420c7b5cce9da6da622ddfd9d607a4e40 --- /dev/null +++ b/src/agentscope/embedding/_embedding_base.py @@ -0,0 +1,449 @@ +# -*- coding: utf-8 -*- +"""The embedding model base class.""" +from __future__ import annotations + +import asyncio +import inspect +from abc import abstractmethod +from pathlib import Path +from typing import Any, Generic, TypeVar, Type, Union + +from pydantic import BaseModel, ConfigDict + +from ._embedding_model_card import EmbeddingModelCard +from ._embedding_response import EmbeddingResponse +from ._embedding_usage import EmbeddingUsage +from .._logging import logger +from ..credential import CredentialBase +from ..message import DataBlock, TextBlock + +#: Type variable for embedding input elements. +#: +#: Bound to the union of all element shapes the framework supports; +#: each concrete subclass narrows it to its accepted input via +#: ``class Foo(EmbeddingModelBase[str | TextBlock]): ...`` so callers +#: get accurate IDE completion and type-checking on ``__call__``. +#: +#: ``__call__`` accepts ``TextBlock`` and unpacks it to ``.text`` +#: before invoking :meth:`_call_api`; therefore ``_call_api``'s input +#: type is intentionally decoupled from :data:`InputT` (it is typed +#: ``list[Any]`` on the base and narrowed by each subclass). +InputT = TypeVar("InputT", bound=Union[str, TextBlock, DataBlock]) + + +class EmbeddingModelBase(Generic[InputT]): + """Base class for embedding models. + + Generic over :data:`InputT` so that text-only subclasses + (``EmbeddingModelBase[str]``) and multimodal subclasses + (``EmbeddingModelBase[str | TextBlock | DataBlock]``) expose the + correct ``inputs`` type to the IDE. + + Follows the same pattern as :class:`~agentscope.model.ChatModelBase`: + + - ``__call__`` splits inputs into batches of size + :attr:`batch_size`, calls :meth:`_call_api` for each batch + **concurrently** via :func:`asyncio.gather`, and merges the + results. Each batch call is wrapped with retry logic. + - Subclasses only implement :meth:`_call_api` for a **single + batch** — no batching or retry code needed. + - Each subclass may override :meth:`_get_retryable_exceptions` to + declare provider-specific retriable errors. + """ + + class Parameters(BaseModel): + """Provider-specific tunables for embedding models. + + Intentionally empty in the base — ``dimensions`` is a contract + property of an embedding model (its output vector size), not a + tunable knob, so it lives directly on the instance via the + required :paramref:`__init__.dimensions` argument. Subclasses + extend this class to expose **non-dimensional** knobs (e.g. + Gemini's ``task_type``, Dashscope's ``text_type``). + + ``extra="allow"`` is set so old persisted configs (where + ``dimensions`` lived in ``parameters``) keep deserialising + without raising; :meth:`EmbeddingModelBase.__init__` extracts + it back out for backward compatibility. + """ + + model_config = ConfigDict(extra="allow") + + credential: CredentialBase + """The API credential.""" + + model: str + """The embedding model name.""" + + dimensions: int + """The output embedding vector dimensions. + + Set directly from the :paramref:`__init__.dimensions` argument — + a required, first-class field rather than something derived from + :attr:`parameters`. + """ + + context_size: int + """Maximum input length (in tokens) per single input item.""" + + batch_size: int + """Maximum number of input items per API call.""" + + max_retries: int + """The maximum number of retries for the underlying API.""" + + retry_delay: float + """Seconds to sleep between retry attempts.""" + + supports_multimodal: bool = False + """Whether this model instance accepts :class:`DataBlock` inputs in + addition to text. Text-only models keep the default ``False``; + multimodal subclasses must set it to ``True`` (per instance when + routing depends on the model name).""" + + def __init__( + self, + credential: CredentialBase, + model: str, + dimensions: int | None, + parameters: BaseModel | None, + context_size: int, + batch_size: int, + max_retries: int, + retry_delay: float, + ) -> None: + """Initialize the embedding model base class. + + Args: + credential (`CredentialBase`): + The API credential used for authentication. + model (`str`): + The name of the embedding model. + dimensions (`int | None`): + The output embedding vector dimensions for this + instance. Required and first-class — see the class + docstring for the rationale of keeping ``dimensions`` + outside :class:`Parameters`. For backward compatibility + with older configs that stored ``dimensions`` inside + :class:`Parameters`, ``None`` is accepted at the + signature level and is back-filled from + ``parameters.dimensions`` when present. + parameters (`BaseModel | None`): + Provider-specific non-dimensional parameters. When + ``None``, the default ``Parameters()`` is used. + context_size (`int`): + Maximum input length (in tokens) per single input item. + batch_size (`int`): + Maximum number of input items per API call. When + ``__call__`` receives more items, it splits them into + batches and calls :meth:`_call_api` concurrently. + max_retries (`int`): + The maximum number of retries for each batch API call. + Only exceptions listed in + :meth:`_get_retryable_exceptions` count against this + budget. + retry_delay (`float`): + Seconds to sleep between retry attempts. + """ + resolved_parameters = parameters or self.Parameters() + # Backward-compat: older session/KB configs persisted + # ``dimensions`` inside ``parameters``. Promote it to the + # constructor argument when the caller did not pass one + # explicitly, then strip it from the parameters object so it + # never reaches provider-specific request payloads. + param_dump = resolved_parameters.model_dump() + legacy_dimensions = param_dump.pop("dimensions", None) + if dimensions is None: + if legacy_dimensions is None: + raise ValueError( + "dimensions is required: pass it explicitly to " + "EmbeddingModelBase.__init__ or include it in the " + "legacy `parameters` mapping.", + ) + dimensions = int(legacy_dimensions) + resolved_parameters = type(resolved_parameters)(**param_dump) + elif legacy_dimensions is not None: + # Both routes set it — explicit constructor wins, strip the + # legacy mirror so it can't drift. + resolved_parameters = type(resolved_parameters)(**param_dump) + + if dimensions <= 0: + raise ValueError( + f"dimensions must be a positive integer, got {dimensions}.", + ) + + self.credential = credential + self.model = model + self.dimensions = dimensions + self.parameters = resolved_parameters + self.context_size = context_size + self.batch_size = batch_size + self.max_retries = max_retries + self.retry_delay = retry_delay + + @classmethod + def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]: + """Return exception types that should trigger a retry. + + Defaults to an empty tuple (no retries). Subclasses can + override to declare provider-specific retryable exceptions. + """ + return () + + # ------------------------------------------------------------------ + # Public API — batching + concurrent retry + # ------------------------------------------------------------------ + + async def __call__( + self, + inputs: list[InputT], + **kwargs: Any, + ) -> EmbeddingResponse: + """Embed a list of inputs with automatic batching and retry. + + The inputs are split into chunks of :attr:`batch_size`. All + chunks are dispatched **concurrently** via + :func:`asyncio.gather`. Each chunk is individually retried up + to ``max_retries`` times on retryable errors. Results are + merged into a single :class:`EmbeddingResponse` preserving the + original input order. + + Args: + inputs (`list[InputT]`): + The input data to embed. For text-only models this is + ``list[str]``; for multimodal models it is + ``list[str | TextBlock | DataBlock]``. Any + :class:`TextBlock` items are transparently unpacked to + their ``.text`` field on entry, so subclasses' batching + and ``_call_api`` only have to handle ``str`` (and + ``DataBlock`` for multimodal variants). + **kwargs: + Additional keyword arguments forwarded to + :meth:`_call_api`. + + Returns: + `EmbeddingResponse`: + A merged response containing embeddings for all inputs. + """ + if not inputs: + return EmbeddingResponse( + embeddings=[], + usage=EmbeddingUsage(tokens=0, time=0), + ) + + normalized: list[Any] = [ + item.text if isinstance(item, TextBlock) else item + for item in inputs + ] + + # Split into batches. + batches = [ + normalized[i : i + self.batch_size] + for i in range(0, len(normalized), self.batch_size) + ] + + if len(batches) > 1: + logger.info( + "Embedding %d inputs in %d batches (batch_size=%d) " + "for model %s.", + len(normalized), + len(batches), + self.batch_size, + self.model, + ) + + # Dispatch all batches concurrently, each with retry. + results: list[EmbeddingResponse] = await asyncio.gather( + *(self._call_with_retry(batch, **kwargs) for batch in batches), + ) + + return self._merge_responses(results) + + # ------------------------------------------------------------------ + # Internal — merge multiple batch responses + # ------------------------------------------------------------------ + + @staticmethod + def _merge_responses( + responses: list[EmbeddingResponse], + ) -> EmbeddingResponse: + """Merge multiple batch :class:`EmbeddingResponse` objects into + one, preserving input order. + + Args: + responses (`list[EmbeddingResponse]`): + Batch responses to merge. + + Returns: + `EmbeddingResponse`: The merged response. + """ + if len(responses) == 1: + return responses[0] + + all_embeddings: list = [] + total_tokens = 0 + total_time = 0.0 + + for resp in responses: + all_embeddings.extend(resp.embeddings) + if resp.usage: + total_time += resp.usage.time + if resp.usage.tokens: + total_tokens += resp.usage.tokens + + return EmbeddingResponse( + embeddings=all_embeddings, + usage=EmbeddingUsage( + tokens=total_tokens, + time=total_time, + ), + source="api", + ) + + # ------------------------------------------------------------------ + # Internal — retry wrapper for a single batch + # ------------------------------------------------------------------ + + async def _call_with_retry( + self, + inputs: list[Any], + **kwargs: Any, + ) -> EmbeddingResponse: + """Call :meth:`_call_api` with retry logic for a single batch. + + Args: + inputs (`list[Any]`): + A single batch of inputs (size ≤ ``batch_size``), already + normalised by :meth:`__call__` (any :class:`TextBlock` + items unpacked to their ``.text``). Typed as + ``list[Any]`` because the concrete element shape depends + on the subclass — see :meth:`_call_api`. + **kwargs: + Forwarded to :meth:`_call_api`. + """ + retryable = tuple(self._get_retryable_exceptions()) + last_error: Exception | None = None + + for attempt in range(self.max_retries + 1): + try: + return await self._call_api(inputs, **kwargs) + except Exception as e: + if not isinstance(e, retryable): + raise + last_error = e + if attempt < self.max_retries: + logger.warning( + "Batch attempt %d failed for embedding model " + "%s: %s. Retrying in %.1fs...", + attempt + 1, + self.model, + str(e), + self.retry_delay, + ) + await asyncio.sleep(self.retry_delay) + else: + logger.warning( + "All %d attempt(s) failed for a batch of " + "embedding model %s.", + self.max_retries + 1, + self.model, + ) + + if last_error is not None: + raise last_error + raise RuntimeError( + f"Failed to call embedding model {self.model} after " + f"{self.max_retries + 1} retries.", + ) + + # ------------------------------------------------------------------ + # Abstract — subclasses implement this for a single batch + # ------------------------------------------------------------------ + + @abstractmethod + async def _call_api( + self, + inputs: list[Any], + **kwargs: Any, + ) -> EmbeddingResponse: + """Call the underlying embedding API for a **single batch**. + + Subclasses must implement this method. The batch splitting, + concurrency, and retry logic are handled by :meth:`__call__` + — this method only needs to handle one API call. + + .. note:: + The parameter is typed ``list[Any]`` rather than + ``list[InputT]`` because :meth:`__call__` unpacks + :class:`TextBlock` items to their ``.text`` field *before* + dispatching to this method. The element shape this method + actually receives is therefore subclass-specific + (:data:`InputT` minus :class:`TextBlock`). Subclasses + should override with their concrete narrower type, e.g. + ``list[str]`` for text-only models or + ``list[str | DataBlock]`` for multimodal ones. + + Args: + inputs (`list[Any]`): + A batch of inputs (guaranteed ``len(inputs) <= + self.batch_size``). + **kwargs: + Additional keyword arguments. + + Returns: + `EmbeddingResponse`: + The embedding response for this batch. + """ + + # ------------------------------------------------------------------ + # Model card discovery + # ------------------------------------------------------------------ + + @classmethod + def list_models( + cls, + custom_yaml_dir: str | None = None, + ) -> list[EmbeddingModelCard]: + """List candidate embedding models from YAML files. + + Each concrete subclass should live in its own provider + subdirectory (e.g. ``embedding/_openai/_model.py``) with a + sibling ``_models/`` directory containing YAML files — identical + to the layout used by :class:`~agentscope.model.ChatModelBase`. + + Args: + custom_yaml_dir (`str | None`): + Override the YAML directory. + + Returns: + `list[EmbeddingModelCard]`: + A list of embedding model cards. + """ + if custom_yaml_dir is None: + subclass_file = Path(inspect.getfile(cls)) + yaml_dir = subclass_file.parent / "_models" + else: + yaml_dir = Path(custom_yaml_dir) + + if not yaml_dir.is_dir(): + return [] + + yaml_files = list(yaml_dir.glob("*.yaml")) + + model_cards = [] + for yaml_file in yaml_files: + try: + card = EmbeddingModelCard.from_yaml( + yaml_path=str(yaml_file), + parameter_class=cls.Parameters, + ) + model_cards.append(card) + except Exception as e: + logger.warning( + "Failed to load embedding model card %s: %s", + yaml_file, + str(e), + ) + continue + + return model_cards diff --git a/src/agentscope/embedding/_embedding_model_card.py b/src/agentscope/embedding/_embedding_model_card.py new file mode 100644 index 0000000000000000000000000000000000000000..b280662deb6529ef1eb61e04f9cca86fea848876 --- /dev/null +++ b/src/agentscope/embedding/_embedding_model_card.py @@ -0,0 +1,179 @@ +# -*- coding: utf-8 -*- +"""The embedding model card class.""" +from __future__ import annotations + +import copy +from typing import Literal, Self, Type + +import yaml +from pydantic import BaseModel, Field + + +class EmbeddingModelCard(BaseModel): + """A card describing an embedding model's capabilities. + + Mirrors :class:`~agentscope.model.ModelCard` but tailored for + embedding models. Uses ``input_types`` / ``output_types`` to + describe model capabilities, and ``parameter_schema`` (built from + the embedding class's ``Parameters`` + YAML ``parameter_overrides``) + to tell the frontend which knobs the user can adjust. + + The output type ``application/x-embedding`` indicates that the + model produces dense vector embeddings. + """ + + type: Literal["embedding_model"] = "embedding_model" + """The card type, always ``"embedding_model"``.""" + + name: str = Field(description="The model name used in API calls.") + """The model name (e.g. ``"text-embedding-3-small"``).""" + + label: str = Field(description="Human-readable label for the frontend.") + """Display label (e.g. ``"Text Embedding 3 Small"``).""" + + status: Literal["active", "deprecated", "sunset"] = Field( + default="active", + description="The model lifecycle status.", + ) + """The model status.""" + + input_types: list[str] = Field( + default=["text/plain"], + description="Supported input media types.", + ) + """Supported input types (e.g. ``["text/plain"]``, + ``["text/plain", "image/jpeg", "image/png"]``).""" + + output_types: list[str] = Field( + default=["application/x-embedding"], + description="Supported output media types.", + ) + """Output types. ``application/x-embedding`` for vector output.""" + + dimensions: int = Field( + ..., + description="Default output vector dimensions for this model.", + gt=0, + ) + """The default output dimensions for this model. + + First-class top-level field — kept outside of + :attr:`parameter_schema` so that callers can rely on a strongly + typed ``int`` rather than the soft ``parameter_schema['properties'] + ['dimensions']['default']`` lookup. + """ + + supported_dimensions: list[int] | None = Field( + default=None, + description=( + "If set, the only dimensions this model can produce. " + "``None`` means dimensions are fixed at " + ":attr:`dimensions` and cannot be overridden." + ), + ) + """Optional set of allowed output dimensions. + + Set for Matryoshka-style models (e.g. OpenAI's + ``text-embedding-3-*``) that can be truncated to a smaller size. + ``None`` indicates a fixed-dimension model. + """ + + context_size: int | None = Field( + default=None, + description="Maximum input length (in tokens) per request.", + gt=0, + ) + """Maximum input context size, if known.""" + + parameter_schema: dict = Field( + default_factory=dict, + description=( + "JSON Schema for user-configurable parameters " + "(built from the Parameters class + YAML overrides)." + ), + ) + """The parameter schema sent to the frontend for form rendering. + Empty ``properties`` means nothing to configure (e.g. fixed + dimensions).""" + + parameter_overrides: dict[str, dict] = Field( + default_factory=dict, + description="Raw parameter overrides from the YAML file.", + ) + """The raw parameter overrides, preserved for reference.""" + + @classmethod + def from_yaml( + cls, + yaml_path: str, + parameter_class: Type[BaseModel], + ) -> Self: + """Load an embedding model card from a YAML file. + + Merges the base ``parameter_class`` JSON Schema with + ``parameter_overrides`` from the YAML — identical to the + approach used by :meth:`~agentscope.model.ModelCard.from_yaml`. + + Args: + yaml_path (`str`): + Path to the YAML file. + parameter_class (`Type[BaseModel]`): + The ``Parameters`` class from the embedding model subclass. + + Returns: + `EmbeddingModelCard`: The loaded model card. + """ + with open(yaml_path, "r", encoding="utf-8") as f: + config = yaml.safe_load(f) + + if "dimensions" not in config: + raise ValueError( + f"Embedding model card {yaml_path!r} is missing the " + f"required top-level 'dimensions' field.", + ) + + # Build parameter schema from the Parameters class + base_schema = parameter_class.model_json_schema() + properties = copy.deepcopy(base_schema.get("properties", {})) + + # Apply parameter_overrides (same logic as ModelCard.from_yaml) + overrides = config.get("parameter_overrides", {}) + for param_name, override in overrides.items(): + if override is None: + # null means remove + properties.pop(param_name, None) + continue + + if isinstance(override, dict): + if override.get("hidden"): + properties.pop(param_name, None) + continue + + # Simple dict merge + if param_name in properties: + properties[param_name] = { + **properties[param_name], + **override, + } + + final_schema = { + "type": "object", + "properties": properties, + "required": base_schema.get("required", []), + } + + return cls( + name=config["name"], + label=config["label"], + status=config.get("status", "active"), + input_types=config.get("input_types", ["text/plain"]), + output_types=config.get( + "output_types", + ["application/x-embedding"], + ), + dimensions=config["dimensions"], + supported_dimensions=config.get("supported_dimensions"), + context_size=config.get("context_size"), + parameter_schema=final_schema, + parameter_overrides=overrides, + ) diff --git a/src/agentscope/embedding/_embedding_response.py b/src/agentscope/embedding/_embedding_response.py new file mode 100644 index 0000000000000000000000000000000000000000..8b855f09c414f7c380fbd78862699b941b26bd78 --- /dev/null +++ b/src/agentscope/embedding/_embedding_response.py @@ -0,0 +1,32 @@ +# -*- coding: utf-8 -*- +"""The embedding response class.""" +from dataclasses import dataclass, field +from typing import Literal, List + +from ._embedding_usage import EmbeddingUsage +from .._utils._common import _get_timestamp +from .._utils._mixin import DictMixin +from ..types import Embedding + + +@dataclass +class EmbeddingResponse(DictMixin): + """The embedding response class.""" + + embeddings: List[Embedding] + """The embedding data""" + + id: str = field(default_factory=lambda: _get_timestamp(True)) + """The identity of the embedding response""" + + created_at: str = field(default_factory=_get_timestamp) + """The timestamp of the embedding response creation""" + + type: Literal["embedding"] = field(default_factory=lambda: "embedding") + """The type of the response, must be `embedding`.""" + + usage: EmbeddingUsage | None = field(default_factory=lambda: None) + """The usage of the embedding model API invocation, if available.""" + + source: Literal["cache", "api"] = field(default_factory=lambda: "api") + """If the response comes from the cache or the API.""" diff --git a/src/agentscope/embedding/_embedding_usage.py b/src/agentscope/embedding/_embedding_usage.py new file mode 100644 index 0000000000000000000000000000000000000000..7da5e56f287563f184aa980b49f578d3967bcd20 --- /dev/null +++ b/src/agentscope/embedding/_embedding_usage.py @@ -0,0 +1,20 @@ +# -*- coding: utf-8 -*- +"""The embedding usage class in agentscope.""" +from dataclasses import dataclass, field +from typing import Literal + +from .._utils._mixin import DictMixin + + +@dataclass +class EmbeddingUsage(DictMixin): + """The usage of an embedding model API invocation.""" + + time: float + """The time used in seconds.""" + + tokens: int | None = field(default_factory=lambda: None) + """The number of tokens used, if available.""" + + type: Literal["embedding"] = field(default_factory=lambda: "embedding") + """The type of the usage, must be `embedding`.""" diff --git a/src/agentscope/embedding/_file_cache.py b/src/agentscope/embedding/_file_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..68901560a6c3f499d0f292b50531d03917514bd9 --- /dev/null +++ b/src/agentscope/embedding/_file_cache.py @@ -0,0 +1,187 @@ +# -*- coding: utf-8 -*- +"""A file embedding cache implementation for storing and retrieving +embeddings in binary files.""" +import hashlib +import json +import os +from typing import Any, List + +import numpy as np + +from ._cache_base import EmbeddingCacheBase +from .._logging import logger +from ..types import ( + Embedding, + JSONSerializableObject, +) + + +class FileEmbeddingCache(EmbeddingCacheBase): + """The embedding cache class that stores each embeddings vector in + binary files.""" + + def __init__( + self, + cache_dir: str = "./.cache/embeddings", + max_file_number: int | None = None, + max_cache_size: int | None = None, + ) -> None: + """Initialize the file embedding cache class. + + Args: + cache_dir (`str`, defaults to `"./.cache/embeddings"`): + The directory to store the embedding files. + max_file_number (`int | None`, defaults to `None`): + The maximum number of files to keep in the cache directory. If + exceeded, the oldest files will be removed. + max_cache_size (`int | None`, defaults to `None`): + The maximum size of the cache directory in MB. If exceeded, + the oldest files will be removed until the size is within the + limit. + """ + self._cache_dir = os.path.abspath(cache_dir) + self.max_file_number = max_file_number + self.max_cache_size = max_cache_size + + @property + def cache_dir(self) -> str: + """The cache directory where the embedding files are stored.""" + if not os.path.exists(self._cache_dir): + os.makedirs(self._cache_dir, exist_ok=True) + return self._cache_dir + + async def store( + self, + embeddings: List[Embedding], + identifier: JSONSerializableObject, + overwrite: bool = False, + **kwargs: Any, + ) -> None: + """Store the embeddings with the given identifier. + + Args: + embeddings (`List[Embedding]`): + The embeddings to store. + identifier (`JSONSerializableObject`): + The identifier to distinguish the embeddings, which will be + used to generate a hashable filename, so it should be + JSON serializable (e.g. a string, number, list, dict). + overwrite (`bool`, defaults to `False`): + Whether to overwrite existing embeddings with the same + identifier. If `True`, existing embeddings will be replaced. + """ + filename = self._get_filename(identifier) + path_file = os.path.join(self.cache_dir, filename) + + if os.path.exists(path_file): + if not os.path.isfile(path_file): + raise RuntimeError( + f"Path {path_file} exists but is not a file.", + ) + + if overwrite: + np.save(path_file, embeddings) + await self._maintain_cache_dir() + else: + np.save(path_file, embeddings) + await self._maintain_cache_dir() + + async def retrieve( + self, + identifier: JSONSerializableObject, + ) -> List[Embedding] | None: + """Retrieve the embeddings with the given identifier. If not found, + return `None`. + + Args: + identifier (`JSONSerializableObject`): + The identifier to retrieve the embeddings, which will be + used to generate a hashable filename, so it should be + JSON serializable (e.g. a string, number, list, dict). + """ + filename = self._get_filename(identifier) + path_file = os.path.join(self.cache_dir, filename) + + if os.path.exists(path_file): + return np.load(os.path.join(self.cache_dir, filename)).tolist() + return None + + async def remove(self, identifier: JSONSerializableObject) -> None: + """Remove the embeddings with the given identifier. + + Args: + identifier (`JSONSerializableObject`): + The identifiers to remove the embeddings, which will be + used to generate a hashable filename, so it should be + JSON serializable (e.g. a string, number, list, dict). + """ + filename = self._get_filename(identifier) + path_file = os.path.join(self.cache_dir, filename) + + if os.path.exists(path_file): + os.remove(path_file) + else: + raise FileNotFoundError(f"File {path_file} does not exist.") + + async def clear(self) -> None: + """Clear the cache directory by removing all files.""" + for filename in os.listdir(self.cache_dir): + if filename.endswith(".npy"): + os.remove(os.path.join(self.cache_dir, filename)) + + def _get_cache_size(self) -> float: + """Get the current size of the cache directory in MB.""" + total_size = 0 + for filename in os.listdir(self.cache_dir): + if filename.endswith(".npy"): + path_file = os.path.join(self.cache_dir, filename) + if os.path.isfile(path_file): + total_size += os.path.getsize(path_file) + return total_size / (1024.0 * 1024.0) + + @staticmethod + def _get_filename(identifier: JSONSerializableObject) -> str: + """Generate a filename based on the identifier.""" + json_str = json.dumps(identifier, ensure_ascii=False) + return hashlib.sha256(json_str.encode("utf-8")).hexdigest() + ".npy" + + async def _maintain_cache_dir(self) -> None: + """Maintain the cache directory by removing old files if the number of + files exceeds the maximum limit or if the cache size exceeds the + maximum size.""" + files = [ + (_.name, _.stat().st_mtime) + for _ in os.scandir(self.cache_dir) + if _.is_file() and _.name.endswith(".npy") + ] + files.sort(key=lambda x: x[1]) + + if self.max_file_number and len(files) > self.max_file_number: + for file_name, _ in files[: 0 - self.max_file_number]: + os.remove(os.path.join(self.cache_dir, file_name)) + logger.info( + "Remove cached embedding file %s for limited number " + "of files (%d).", + file_name, + self.max_file_number, + ) + files = files[0 - self.max_file_number :] + + if ( + self.max_cache_size is not None + and self._get_cache_size() > self.max_cache_size + ): + removed_files = [] + for filename, _ in files: + os.remove(os.path.join(self.cache_dir, filename)) + removed_files.append(filename) + if self._get_cache_size() <= self.max_cache_size: + break + + if removed_files: + logger.info( + "Remove %d cached embedding file(s) for limited " + "cache size (%d MB).", + len(removed_files), + self.max_cache_size, + ) diff --git a/src/agentscope/embedding/_gemini/__init__.py b/src/agentscope/embedding/_gemini/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a097d7040ba9a96e95bf05bed58c7a82f79fb0e0 --- /dev/null +++ b/src/agentscope/embedding/_gemini/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- +"""The Gemini embedding API modules.""" + +from ._model import GeminiEmbeddingModel + +__all__ = [ + "GeminiEmbeddingModel", +] diff --git a/src/agentscope/embedding/_gemini/_model.py b/src/agentscope/embedding/_gemini/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..b382de96a6fbaf305b6f29cbfb16580c3495d0ea --- /dev/null +++ b/src/agentscope/embedding/_gemini/_model.py @@ -0,0 +1,485 @@ +# -*- coding: utf-8 -*- +"""The Google Gemini embedding model. + +Handles both text-only and multimodal models under a single class. +``gemini-embedding-001`` accepts ``list[str | TextBlock]``. +``gemini-embedding-2`` additionally accepts +:class:`~agentscope.message.DataBlock` (images, video, audio, PDF). +The model name determines the API call style. + +Text payloads may be passed either as bare ``str`` or as +:class:`~agentscope.message.TextBlock` — the latter is unpacked to its +``.text`` field on entry so the rest of the pipeline only deals with +``str`` and ``DataBlock``. +""" +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from datetime import datetime +from typing import Any + +from .._cache_base import EmbeddingCacheBase +from .._embedding_response import EmbeddingResponse +from .._embedding_usage import EmbeddingUsage +from .._embedding_base import EmbeddingModelBase +from ..._logging import logger +from ...credential import CredentialBase +from ...message import DataBlock, TextBlock + +#: Model name prefixes that use the multimodal API path. +_MULTIMODAL_PREFIXES = ("gemini-embedding-2",) + + +@dataclass +class _MultimodalLimits: + """Per-request constraints for Gemini multimodal embedding.""" + + max_elements: int = 20 + """Maximum total content elements per API call.""" + + max_images: int = 6 + """Maximum image elements per API call.""" + + max_videos: int = 1 + """Maximum video elements per API call.""" + + max_audios: int = 1 + """Maximum audio elements per API call.""" + + max_pdfs: int = 1 + """Maximum PDF documents per API call.""" + + +_MODEL_LIMITS: dict[str, _MultimodalLimits] = { + "gemini-embedding-2": _MultimodalLimits( + max_elements=20, + max_images=6, + max_videos=1, + max_audios=1, + max_pdfs=1, + ), +} + +_DEFAULT_LIMITS = _MultimodalLimits() + + +class GeminiEmbeddingModel(EmbeddingModelBase[str | TextBlock | DataBlock]): + """Unified Google Gemini embedding model. + + Routes to the text-only or multimodal Gemini API based on the + model name. + + - **Text mode** (``gemini-embedding-001``): uses the base class's + simple batch splitting + concurrent retry. The Gemini API + accepts a list of strings and returns individual embeddings. + - **Multimodal mode** (``gemini-embedding-2``): overrides + ``__call__`` with content-aware batching (respecting per-model + limits on images, videos, audios, PDFs). Each input is wrapped + in a ``Content`` object so the API returns separate embeddings. + + Key API differences from other providers: + + - Dimensions are controlled via ``output_dimensionality`` in the + ``config`` parameter (not a top-level ``dimensions`` field). + - ``gemini-embedding-001`` supports ``task_type`` in config; + ``gemini-embedding-2`` uses prompt prefixes instead. + """ + + #: Text-mode batch size. Gemini docs don't specify an explicit + #: limit; we use a conservative default. + _TEXT_BATCH_SIZE: int = 100 + + def __init__( + self, + credential: CredentialBase, + model: str, + dimensions: int | None, + parameters: "GeminiEmbeddingModel.Parameters | None" = None, + embedding_cache: EmbeddingCacheBase | None = None, + context_size: int = 8192, + max_retries: int = 3, + retry_delay: float = 1.0, + ) -> None: + """Initialize the Gemini embedding model. + + Args: + credential (`CredentialBase`): + A :class:`~agentscope.credential.GeminiCredential` + instance providing the API key. + model (`str`): + The embedding model name (e.g. + ``"gemini-embedding-001"`` or + ``"gemini-embedding-2"``). + dimensions (`int | None`): + The output embedding vector dimensions. Required at + the contract level — see :class:`EmbeddingModelBase` + for the rationale. ``None`` is accepted only for + backward compatibility with legacy configs that + persisted ``dimensions`` inside ``parameters``. + parameters (`GeminiEmbeddingModel.Parameters | None`, \ + defaults to ``None``): + Provider-specific non-dimensional parameters. Currently + empty for Gemini. + embedding_cache (`EmbeddingCacheBase | None`, defaults to \ + ``None``): + Optional embedding cache. + context_size (`int`, defaults to ``8192``): + Maximum input tokens. 2048 for ``gemini-embedding-001``, + 8192 for ``gemini-embedding-2``. + max_retries (`int`, defaults to ``3``): + Number of retries on transient failures. + retry_delay (`float`, defaults to ``1.0``): + Seconds between retry attempts. + """ + from google import genai + + self._is_multimodal: bool = model.startswith(_MULTIMODAL_PREFIXES) + + super().__init__( + credential=credential, + model=model, + dimensions=dimensions, + parameters=parameters, + context_size=context_size, + batch_size=self._TEXT_BATCH_SIZE, + max_retries=max_retries, + retry_delay=retry_delay, + ) + self.supports_multimodal = self._is_multimodal + + self.client: genai.Client = genai.Client( + api_key=credential.api_key.get_secret_value(), + ) + self.embedding_cache: EmbeddingCacheBase | None = embedding_cache + + if self._is_multimodal: + self._limits = _MODEL_LIMITS.get(model, _DEFAULT_LIMITS) + + # ------------------------------------------------------------------ + # __call__ — override for multimodal content-aware batching + # ------------------------------------------------------------------ + + async def __call__( + self, + inputs: list[str | TextBlock | DataBlock], + **kwargs: Any, + ) -> EmbeddingResponse: + """Embed inputs with batching and retry. + + For text models, delegates to the base class. For multimodal + models, performs content-aware batching that respects per-model + limits on images, videos, audios, and PDFs. + + Args: + inputs (`list[str | TextBlock | DataBlock]`): + The input data to embed. ``TextBlock`` items are + unpacked to their ``.text`` field on entry, so the + remainder of the pipeline only sees ``str`` and + ``DataBlock``. + **kwargs: + Forwarded to the Gemini API config. + + Returns: + `EmbeddingResponse`: Merged response for all inputs. + """ + normalized: list[str | DataBlock] = [ + item.text if isinstance(item, TextBlock) else item + for item in inputs + ] + + if not self._is_multimodal: + return await super().__call__(normalized, **kwargs) + + batches = self._split_multimodal_batches(normalized) + + if len(batches) > 1: + logger.info( + "Embedding %d multimodal inputs in %d batches " + "for model %s.", + len(normalized), + len(batches), + self.model, + ) + + results: list[EmbeddingResponse] = await asyncio.gather( + *(self._call_with_retry(batch, **kwargs) for batch in batches), + ) + + return self._merge_responses(results) + + def _split_multimodal_batches( + self, + inputs: list[str | DataBlock], + ) -> list[list[str | DataBlock]]: + """Split inputs into batches respecting Gemini multimodal limits. + + Greedy: keep adding items until any constraint would be + violated, then start a new batch. + + Args: + inputs (`list[str | DataBlock]`): + All inputs to split. + + Returns: + `list[list[str | DataBlock]]`: List of batches. + """ + limits = self._limits + batches: list[list[str | DataBlock]] = [] + current: list[str | DataBlock] = [] + n_elem = 0 + n_img = 0 + n_vid = 0 + n_aud = 0 + n_pdf = 0 + + for item in inputs: + is_img = is_vid = is_aud = is_pdf = False + if isinstance(item, DataBlock): + mt = item.source.media_type + is_img = mt.startswith("image/") + is_vid = mt.startswith("video/") + is_aud = mt.startswith("audio/") + is_pdf = mt == "application/pdf" + + would_exceed = ( + n_elem + 1 > limits.max_elements + or (is_img and n_img + 1 > limits.max_images) + or (is_vid and n_vid + 1 > limits.max_videos) + or (is_aud and n_aud + 1 > limits.max_audios) + or (is_pdf and n_pdf + 1 > limits.max_pdfs) + ) + + if would_exceed and current: + batches.append(current) + current = [] + n_elem = n_img = n_vid = n_aud = n_pdf = 0 + + current.append(item) + n_elem += 1 + n_img += is_img + n_vid += is_vid + n_aud += is_aud + n_pdf += is_pdf + + if current: + batches.append(current) + + return batches + + # ------------------------------------------------------------------ + # _call_api — single batch + # ------------------------------------------------------------------ + + async def _call_api( + self, + inputs: list[str | DataBlock], + **kwargs: Any, + ) -> EmbeddingResponse: + """Route to text or multimodal Gemini API for a single batch. + + Args: + inputs (`list[str | DataBlock]`): + A single batch of inputs. + **kwargs: + Extra keyword arguments merged into the Gemini + ``EmbedContentConfig``. + + Returns: + `EmbeddingResponse`: Embedding vectors and usage info. + """ + if self._is_multimodal: + return await self._call_multimodal(inputs, **kwargs) + return await self._call_text(inputs, **kwargs) + + # ------------------------------------------------------------------ + # Text API (gemini-embedding-001) + # ------------------------------------------------------------------ + + async def _call_text( + self, + inputs: list[str | DataBlock], + **kwargs: Any, + ) -> EmbeddingResponse: + """Call the Gemini text embedding API for a single batch. + + Passes the list of strings directly to ``embed_content``, + which returns one embedding per string. + + Args: + inputs (`list[str | DataBlock]`): + Must all be ``str``; raises ``ValueError`` otherwise. + **kwargs: + Merged into ``EmbedContentConfig`` (e.g. + ``task_type``). + + Returns: + `EmbeddingResponse`: Embedding vectors and usage info. + """ + from google.genai import types + + texts: list[str] = [] + for item in inputs: + if not isinstance(item, str): + raise ValueError( + f"Text embedding model {self.model!r} only accepts " + f"str inputs, got {type(item).__name__}.", + ) + texts.append(item) + + config = types.EmbedContentConfig( + output_dimensionality=self.dimensions, + **kwargs, + ) + + cache_key = { + "model": self.model, + "contents": texts, + "output_dimensionality": self.dimensions, + **kwargs, + } + + if self.embedding_cache: + cached = await self.embedding_cache.retrieve( + identifier=cache_key, + ) + if cached: + return EmbeddingResponse( + embeddings=cached, + usage=EmbeddingUsage(tokens=0, time=0), + source="cache", + ) + + start_time = datetime.now() + response = self.client.models.embed_content( + model=self.model, + contents=texts, + config=config, + ) + time = (datetime.now() - start_time).total_seconds() + + embeddings = [item.values for item in response.embeddings] + + if self.embedding_cache: + await self.embedding_cache.store( + identifier=cache_key, + embeddings=embeddings, + ) + + return EmbeddingResponse( + embeddings=embeddings, + usage=EmbeddingUsage(time=time), + ) + + # ------------------------------------------------------------------ + # Multimodal API (gemini-embedding-2) + # ------------------------------------------------------------------ + + async def _call_multimodal( + self, + inputs: list[str | DataBlock], + **kwargs: Any, + ) -> EmbeddingResponse: + """Call the Gemini multimodal embedding API for a single batch. + + Each input is wrapped in a separate ``Content`` object so the + API returns one embedding per input (not one aggregated + embedding). + + Args: + inputs (`list[str | DataBlock]`): + ``str`` for text, ``DataBlock`` for images / video / + audio / PDF. + **kwargs: + Merged into ``EmbedContentConfig``. + + Returns: + `EmbeddingResponse`: Embedding vectors and usage info. + """ + from google.genai import types + + contents: list[types.Content] = [] + for item in inputs: + if isinstance(item, str): + contents.append( + types.Content( + parts=[types.Part.from_text(text=item)], + ), + ) + elif isinstance(item, DataBlock): + contents.append( + types.Content( + parts=[self._data_block_to_part(item)], + ), + ) + else: + raise ValueError( + f"Invalid input: {item!r}. Expected str or DataBlock.", + ) + + config = types.EmbedContentConfig( + output_dimensionality=self.dimensions, + **kwargs, + ) + + start_time = datetime.now() + response = self.client.models.embed_content( + model=self.model, + contents=contents, + config=config, + ) + time = (datetime.now() - start_time).total_seconds() + + embeddings = [item.values for item in response.embeddings] + + return EmbeddingResponse( + embeddings=embeddings, + usage=EmbeddingUsage(time=time), + ) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + @staticmethod + def _data_block_to_part(block: DataBlock) -> Any: + """Convert a :class:`~agentscope.message.DataBlock` to a Gemini + ``Part`` object. + + Args: + block (`DataBlock`): + A data block with ``Base64Source`` or ``URLSource``. + + Returns: + A ``google.genai.types.Part`` instance. + + Raises: + `ValueError`: If the source type is unsupported. + """ + from google.genai import types + from ...message import Base64Source, URLSource + + source = block.source + + if isinstance(source, Base64Source): + import base64 + + return types.Part.from_bytes( + data=base64.b64decode(source.data), + mime_type=source.media_type, + ) + + if isinstance(source, URLSource): + # Gemini SDK doesn't have a direct from_url for + # embed_content; download or use File API. + # For now, raise — callers should use Base64Source. + raise ValueError( + "Gemini embedding API requires inline data " + "(Base64Source). URLSource is not directly supported " + f"for embedding. Got URL: {source.url}", + ) + + raise ValueError( + f"Unsupported source type {type(source).__name__} " + f"in DataBlock.", + ) diff --git a/src/agentscope/embedding/_gemini/_models/gemini-embedding-001.yaml b/src/agentscope/embedding/_gemini/_models/gemini-embedding-001.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c70d67b4a42e43940b060fc30f458d9f503d073b --- /dev/null +++ b/src/agentscope/embedding/_gemini/_models/gemini-embedding-001.yaml @@ -0,0 +1,14 @@ +name: gemini-embedding-001 +label: Gemini Embedding 001 +status: active + +input_types: + - text/plain + +output_types: + - application/x-embedding + +context_size: 2048 + +dimensions: 3072 +supported_dimensions: [3072, 1536, 768, 512, 256, 128] diff --git a/src/agentscope/embedding/_gemini/_models/gemini-embedding-2.yaml b/src/agentscope/embedding/_gemini/_models/gemini-embedding-2.yaml new file mode 100644 index 0000000000000000000000000000000000000000..105335acb3d13003e6ef46c65b8c6fd9cad3bc7c --- /dev/null +++ b/src/agentscope/embedding/_gemini/_models/gemini-embedding-2.yaml @@ -0,0 +1,20 @@ +name: gemini-embedding-2 +label: Gemini Embedding 2 +status: active + +input_types: + - text/plain + - image/png + - image/jpeg + - video/mp4 + - audio/mpeg + - audio/wav + - application/pdf + +output_types: + - application/x-embedding + +context_size: 8192 + +dimensions: 3072 +supported_dimensions: [3072, 1536, 768] diff --git a/src/agentscope/embedding/_ollama/__init__.py b/src/agentscope/embedding/_ollama/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5930ba99bce2211ecd3fb34386e1bfad31ae2870 --- /dev/null +++ b/src/agentscope/embedding/_ollama/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- +"""The Ollama embedding API modules.""" + +from ._model import OllamaEmbeddingModel + +__all__ = [ + "OllamaEmbeddingModel", +] diff --git a/src/agentscope/embedding/_ollama/_model.py b/src/agentscope/embedding/_ollama/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..6264d42268a006565717b5bc5e2fd6e8e2129f8c --- /dev/null +++ b/src/agentscope/embedding/_ollama/_model.py @@ -0,0 +1,132 @@ +# -*- coding: utf-8 -*- +"""The Ollama embedding model.""" + +from datetime import datetime +from typing import Any + +from .._embedding_response import EmbeddingResponse +from .._embedding_usage import EmbeddingUsage +from .._cache_base import EmbeddingCacheBase +from .._embedding_base import EmbeddingModelBase +from ...credential import CredentialBase +from ...message import TextBlock + + +class OllamaEmbeddingModel(EmbeddingModelBase[str | TextBlock]): + """Ollama text embedding model. + + Wraps locally-hosted embedding models served by Ollama (e.g. + ``nomic-embed-text``, ``mxbai-embed-large``). Inherits batching + and retry logic from :class:`EmbeddingModelBase`. ``TextBlock`` + items in the input list are unpacked to their ``.text`` field by + the base class before ``_call_api`` runs, so this subclass only + has to handle plain ``str``. + """ + + _TEXT_BATCH_SIZE: int = 512 + + def __init__( + self, + credential: CredentialBase, + model: str, + dimensions: int | None, + parameters: "OllamaEmbeddingModel.Parameters | None" = None, + embedding_cache: EmbeddingCacheBase | None = None, + context_size: int = 8192, + max_retries: int = 3, + retry_delay: float = 1.0, + ) -> None: + """Initialize the Ollama embedding model. + + Args: + credential (`CredentialBase`): + An :class:`~agentscope.credential.OllamaCredential` + instance providing the host URL. + model (`str`): + The embedding model name (e.g. + ``"nomic-embed-text"``). + dimensions (`int | None`): + The output embedding vector dimensions. Required at + the contract level — see :class:`EmbeddingModelBase` + for the rationale. ``None`` is accepted only for + backward compatibility with legacy configs that + persisted ``dimensions`` inside ``parameters``. + parameters (`OllamaEmbeddingModel.Parameters | None`, \ + defaults to ``None``): + Provider-specific non-dimensional parameters. Currently + empty for Ollama. + embedding_cache (`EmbeddingCacheBase | None`, defaults to \ + ``None``): + Optional embedding cache. + context_size (`int`, defaults to ``8192``): + Maximum input tokens per text. + max_retries (`int`, defaults to ``3``): + Number of retries on transient failures. + retry_delay (`float`, defaults to ``1.0``): + Seconds between retry attempts. + """ + super().__init__( + credential=credential, + model=model, + dimensions=dimensions, + parameters=parameters, + context_size=context_size, + batch_size=self._TEXT_BATCH_SIZE, + max_retries=max_retries, + retry_delay=retry_delay, + ) + self.host: str | None = getattr(credential, "host", None) + self.embedding_cache: EmbeddingCacheBase | None = embedding_cache + + async def _call_api( + self, + inputs: list[str], + **kwargs: Any, + ) -> EmbeddingResponse: + """Call the Ollama embedding API for a single batch. + + Args: + inputs (`list[str]`): + A batch of texts to embed. + **kwargs: + Extra keyword arguments forwarded to the Ollama API. + + Returns: + `EmbeddingResponse`: Embedding vectors and usage info. + """ + api_kwargs: dict[str, Any] = { + "input": inputs, + "model": self.model, + "dimensions": self.dimensions, + **kwargs, + } + + if self.embedding_cache: + cached = await self.embedding_cache.retrieve( + identifier=api_kwargs, + ) + if cached: + return EmbeddingResponse( + embeddings=cached, + usage=EmbeddingUsage(tokens=0, time=0), + source="cache", + ) + + import ollama + + client = ollama.AsyncClient(host=self.host) + + start_time = datetime.now() + response = await client.embed(**api_kwargs) + time = (datetime.now() - start_time).total_seconds() + + if self.embedding_cache: + await self.embedding_cache.store( + identifier=api_kwargs, + embeddings=response.embeddings, + ) + + return EmbeddingResponse( + embeddings=response.embeddings, + usage=EmbeddingUsage(time=time), + ) diff --git a/src/agentscope/embedding/_openai/__init__.py b/src/agentscope/embedding/_openai/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..26fdc2237aee0f968777b6abc0fa9c021e5eaf31 --- /dev/null +++ b/src/agentscope/embedding/_openai/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- +"""The OpenAI embedding API modules.""" + +from ._model import OpenAIEmbeddingModel + +__all__ = [ + "OpenAIEmbeddingModel", +] diff --git a/src/agentscope/embedding/_openai/_model.py b/src/agentscope/embedding/_openai/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..4b389ebaaa48edf101e5f19deeced49d6a955c25 --- /dev/null +++ b/src/agentscope/embedding/_openai/_model.py @@ -0,0 +1,179 @@ +# -*- coding: utf-8 -*- +"""The OpenAI embedding model.""" +from __future__ import annotations + +from datetime import datetime +from typing import Any, Type + +from .._embedding_response import EmbeddingResponse +from .._embedding_usage import EmbeddingUsage +from .._cache_base import EmbeddingCacheBase +from .._embedding_base import EmbeddingModelBase +from ...credential import CredentialBase +from ...message import TextBlock + + +class OpenAIEmbeddingModel(EmbeddingModelBase[str | TextBlock]): + """OpenAI text embedding model. + + Supports ``text-embedding-3-small``, ``text-embedding-3-large``, + and other OpenAI-compatible embedding models. Inherits batching + and retry logic from :class:`EmbeddingModelBase`. ``TextBlock`` + items in the input list are unpacked to their ``.text`` field by + the base class before ``_call_api`` runs, so this subclass only + has to handle plain ``str``. + """ + + #: OpenAI does not document an explicit per-request item limit; + #: the constraint is on total tokens. We use a conservative + #: default that works well in practice. + _TEXT_BATCH_SIZE: int = 2048 + + def __init__( + self, + credential: CredentialBase, + model: str, + dimensions: int | None, + parameters: "OpenAIEmbeddingModel.Parameters | None" = None, + pass_dimensions: bool = True, + embedding_cache: EmbeddingCacheBase | None = None, + context_size: int = 8191, + max_retries: int = 3, + retry_delay: float = 1.0, + ) -> None: + """Initialize the OpenAI embedding model. + + Args: + credential (`CredentialBase`): + An :class:`~agentscope.credential.OpenAICredential` + instance providing the API key and optional base URL / + organization. + model (`str`): + The embedding model name (e.g. + ``"text-embedding-3-small"``). + dimensions (`int | None`): + The output embedding vector dimensions. Required at + the contract level — see :class:`EmbeddingModelBase` + for the rationale. ``None`` is accepted only for + backward compatibility with legacy configs that + persisted ``dimensions`` inside ``parameters``. + parameters (`OpenAIEmbeddingModel.Parameters | None`, \ + defaults to ``None``): + Provider-specific non-dimensional parameters. Currently + empty for OpenAI. + pass_dimensions (`bool`, defaults to `True`): + Whether to pass the ``dimensions`` parameter to the API. + Some OpenAI-compatible providers do not support it. + embedding_cache (`EmbeddingCacheBase | None`, defaults to \ + ``None``): + Optional embedding cache. + context_size (`int`, defaults to ``8191``): + Maximum input tokens per text. + max_retries (`int`, defaults to ``3``): + Number of retries on transient failures. + retry_delay (`float`, defaults to ``1.0``): + Seconds between retry attempts. + """ + import openai + + super().__init__( + credential=credential, + model=model, + dimensions=dimensions, + parameters=parameters, + context_size=context_size, + batch_size=self._TEXT_BATCH_SIZE, + max_retries=max_retries, + retry_delay=retry_delay, + ) + + client_kwargs: dict[str, Any] = {} + if getattr(credential, "base_url", None) is not None: + client_kwargs["base_url"] = credential.base_url + if getattr(credential, "organization", None) is not None: + client_kwargs["organization"] = credential.organization + + self.client: openai.AsyncClient = openai.AsyncClient( + api_key=credential.api_key.get_secret_value(), + **client_kwargs, + ) + self.pass_dimensions = pass_dimensions + self.embedding_cache: EmbeddingCacheBase | None = embedding_cache + + @classmethod + def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]: + """Return OpenAI exceptions that warrant a retry.""" + import openai + + return ( + openai.APIConnectionError, + openai.APITimeoutError, + openai.RateLimitError, + openai.InternalServerError, + ) + + async def _call_api( + self, + inputs: list[str], + **kwargs: Any, + ) -> EmbeddingResponse: + """Call the OpenAI embedding API for a single batch. + + Args: + inputs (`list[str]`): + A batch of texts to embed. + **kwargs: + Extra keyword arguments forwarded to the OpenAI API. + + Returns: + `EmbeddingResponse`: Embedding vectors and usage info. + """ + api_kwargs: dict[str, Any] = { + "input": inputs, + "model": self.model, + "encoding_format": "float", + **kwargs, + } + if self.pass_dimensions: + api_kwargs["dimensions"] = self.dimensions + + if self.embedding_cache: + cached = await self.embedding_cache.retrieve( + identifier=api_kwargs, + ) + if cached: + return EmbeddingResponse( + embeddings=cached, + usage=EmbeddingUsage(tokens=0, time=0), + source="cache", + ) + + start_time = datetime.now() + response = await self.client.embeddings.create(**api_kwargs) + time = (datetime.now() - start_time).total_seconds() + + embeddings: list[Any] = [None] * len(inputs) + for pos, item in enumerate(response.data): + index = getattr(item, "index", pos) + if not isinstance(index, int): + index = pos + if 0 <= index < len(inputs): + embeddings[index] = item.embedding or getattr( + item, + "dense_embedding", + None, + ) + + if self.embedding_cache: + await self.embedding_cache.store( + identifier=api_kwargs, + embeddings=embeddings, + ) + + return EmbeddingResponse( + embeddings=embeddings, + usage=EmbeddingUsage( + tokens=response.usage.total_tokens, + time=time, + ), + ) diff --git a/src/agentscope/embedding/_openai/_models/text-embedding-3-large.yaml b/src/agentscope/embedding/_openai/_models/text-embedding-3-large.yaml new file mode 100644 index 0000000000000000000000000000000000000000..335525f3b45bef0429c7e8e1d1dd339ee1872c6d --- /dev/null +++ b/src/agentscope/embedding/_openai/_models/text-embedding-3-large.yaml @@ -0,0 +1,14 @@ +name: text-embedding-3-large +label: Text Embedding 3 Large +status: active + +input_types: + - text/plain + +output_types: + - application/x-embedding + +context_size: 8191 + +dimensions: 3072 +supported_dimensions: [3072, 1536, 1024, 768, 512, 256] diff --git a/src/agentscope/embedding/_openai/_models/text-embedding-3-small.yaml b/src/agentscope/embedding/_openai/_models/text-embedding-3-small.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dfaba1a594b9436df6cbf49022c4fbd8f380707b --- /dev/null +++ b/src/agentscope/embedding/_openai/_models/text-embedding-3-small.yaml @@ -0,0 +1,14 @@ +name: text-embedding-3-small +label: Text Embedding 3 Small +status: active + +input_types: + - text/plain + +output_types: + - application/x-embedding + +context_size: 8191 + +dimensions: 1536 +supported_dimensions: [1536, 1024, 768, 512, 256] diff --git a/src/agentscope/event/__init__.py b/src/agentscope/event/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ea5c6b46b60d6a7129b5cf299d02f649e1dec23b --- /dev/null +++ b/src/agentscope/event/__init__.py @@ -0,0 +1,71 @@ +# -*- coding: utf-8 -*- +"""The event module of agentscope.""" + +from ._event import ( + EventType, + EventBase, + ReplyStartEvent, + ReplyEndEvent, + ModelCallStartEvent, + ModelCallEndEvent, + TextBlockStartEvent, + TextBlockDeltaEvent, + TextBlockEndEvent, + DataBlockStartEvent, + DataBlockDeltaEvent, + DataBlockEndEvent, + ThinkingBlockStartEvent, + ThinkingBlockDeltaEvent, + ThinkingBlockEndEvent, + HintBlockEvent, + ToolCallStartEvent, + ToolCallDeltaEvent, + ToolCallEndEvent, + ToolResultStartEvent, + ToolResultTextDeltaEvent, + ToolResultDataDeltaEvent, + ToolResultEndEvent, + ExceedMaxItersEvent, + RequireUserConfirmEvent, + RequireExternalExecutionEvent, + UserConfirmResultEvent, + ExternalExecutionResultEvent, + CustomEvent, + AgentEvent, + ConfirmResult, +) + + +__all__ = [ + "EventType", + "EventBase", + "ReplyStartEvent", + "ReplyEndEvent", + "ModelCallStartEvent", + "ModelCallEndEvent", + "TextBlockStartEvent", + "TextBlockDeltaEvent", + "TextBlockEndEvent", + "DataBlockStartEvent", + "DataBlockDeltaEvent", + "DataBlockEndEvent", + "ThinkingBlockStartEvent", + "ThinkingBlockDeltaEvent", + "ThinkingBlockEndEvent", + "HintBlockEvent", + "ToolCallStartEvent", + "ToolCallDeltaEvent", + "ToolCallEndEvent", + "ToolResultStartEvent", + "ToolResultTextDeltaEvent", + "ToolResultDataDeltaEvent", + "ToolResultEndEvent", + "ExceedMaxItersEvent", + "RequireUserConfirmEvent", + "RequireExternalExecutionEvent", + "UserConfirmResultEvent", + "ExternalExecutionResultEvent", + "CustomEvent", + "AgentEvent", + "ConfirmResult", +] diff --git a/src/agentscope/event/_event.py b/src/agentscope/event/_event.py new file mode 100644 index 0000000000000000000000000000000000000000..773c7097eff3aa506f097badb25995ba58c4952c --- /dev/null +++ b/src/agentscope/event/_event.py @@ -0,0 +1,506 @@ +# -*- coding: utf-8 -*- +"""Event types for agent execution.""" +from datetime import datetime +from enum import StrEnum +from typing import Any, Dict, Literal, List, TypeAlias + +from pydantic import BaseModel, Field, ConfigDict + +from .._utils._common import _generate_id +from ..message import ( + DataBlock, + TextBlock, + ToolCallBlock, + ToolResultBlock, + ToolResultState, +) +from ..permission import PermissionRule + + +class EventType(StrEnum): + """Event type enumeration.""" + + REPLY_START = "REPLY_START" + REPLY_END = "REPLY_END" + + MODEL_CALL_START = "MODEL_CALL_START" + MODEL_CALL_END = "MODEL_CALL_END" + + TEXT_BLOCK_START = "TEXT_BLOCK_START" + TEXT_BLOCK_DELTA = "TEXT_BLOCK_DELTA" + TEXT_BLOCK_END = "TEXT_BLOCK_END" + + DATA_BLOCK_START = "DATA_BLOCK_START" + DATA_BLOCK_DELTA = "DATA_BLOCK_DELTA" + DATA_BLOCK_END = "DATA_BLOCK_END" + + THINKING_BLOCK_START = "THINKING_BLOCK_START" + THINKING_BLOCK_DELTA = "THINKING_BLOCK_DELTA" + THINKING_BLOCK_END = "THINKING_BLOCK_END" + + HINT_BLOCK = "HINT_BLOCK" + + TOOL_CALL_START = "TOOL_CALL_START" + TOOL_CALL_DELTA = "TOOL_CALL_DELTA" + TOOL_CALL_END = "TOOL_CALL_END" + + TOOL_RESULT_START = "TOOL_RESULT_START" + TOOL_RESULT_TEXT_DELTA = "TOOL_RESULT_TEXT_DELTA" + TOOL_RESULT_DATA_DELTA = "TOOL_RESULT_DATA_DELTA" + TOOL_RESULT_END = "TOOL_RESULT_END" + + EXCEED_MAX_ITERS = "EXCEED_MAX_ITERS" + + REQUIRE_USER_CONFIRM = "REQUIRE_USER_CONFIRM" + REQUIRE_EXTERNAL_EXECUTION = "REQUIRE_EXTERNAL_EXECUTION" + + USER_CONFIRM_RESULT = "USER_CONFIRM_RESULT" + EXTERNAL_EXECUTION_RESULT = "EXTERNAL_EXECUTION_RESULT" + + CUSTOM = "CUSTOM" + + +class EventBase(BaseModel): + """Base event class.""" + + model_config = ConfigDict(use_enum_values=True) + + id: str = Field(default_factory=_generate_id) + """Unique event identifier.""" + created_at: str = Field(default_factory=lambda: datetime.now().isoformat()) + """ISO 8601 timestamp of when the event was created.""" + metadata: Dict[str, Any] = Field(default_factory=dict) + """Optional metadata attached to the event.""" + + +class ReplyStartEvent(EventBase): + """Reply start event.""" + + type: Literal[EventType.REPLY_START] = EventType.REPLY_START + """Event type.""" + session_id: str + """ID of the session this reply belongs to.""" + reply_id: str + """ID of the reply message produced by this reply.""" + name: str + """Name of the agent.""" + role: Literal["user", "assistant", "system"] = "assistant" + """Role of the agent.""" + + +class ReplyEndEvent(EventBase): + """Reply end event.""" + + type: Literal[EventType.REPLY_END] = EventType.REPLY_END + """Event type.""" + session_id: str + """ID of the session this reply belongs to.""" + reply_id: str + """ID of the reply message produced by this reply.""" + + +class ModelCallStartEvent(EventBase): + """Model call start event.""" + + type: Literal[EventType.MODEL_CALL_START] = EventType.MODEL_CALL_START + """Event type.""" + reply_id: str + """ID of the reply message this model call belongs to.""" + model_name: str + """Name of the model being called.""" + + +class ModelCallEndEvent(EventBase): + """Model call end event.""" + + type: Literal[EventType.MODEL_CALL_END] = EventType.MODEL_CALL_END + """Event type.""" + reply_id: str + """ID of the reply message this model call belongs to.""" + input_tokens: int + """Number of input tokens consumed.""" + output_tokens: int + """Number of output tokens generated.""" + + +class TextBlockStartEvent(EventBase): + """Text block start event.""" + + type: Literal[EventType.TEXT_BLOCK_START] = EventType.TEXT_BLOCK_START + """Event type.""" + reply_id: str + """ID of the reply message this block belongs to.""" + block_id: str + """Unique identifier of the text block.""" + + +class TextBlockDeltaEvent(EventBase): + """Text block delta event.""" + + type: Literal[EventType.TEXT_BLOCK_DELTA] = EventType.TEXT_BLOCK_DELTA + """Event type.""" + reply_id: str + """ID of the reply message this block belongs to.""" + block_id: str + """Unique identifier of the text block.""" + delta: str + """Incremental text content.""" + + +class TextBlockEndEvent(EventBase): + """Text block end event.""" + + type: Literal[EventType.TEXT_BLOCK_END] = EventType.TEXT_BLOCK_END + """Event type.""" + reply_id: str + """ID of the reply message this block belongs to.""" + block_id: str + """Unique identifier of the text block.""" + + +class DataBlockStartEvent(EventBase): + """Data block start event.""" + + type: Literal[EventType.DATA_BLOCK_START] = EventType.DATA_BLOCK_START + """Event type.""" + reply_id: str + """ID of the reply message this block belongs to.""" + block_id: str + """Unique identifier of the data block.""" + media_type: str + """MIME type of the data content (e.g. "image/png").""" + + +class DataBlockDeltaEvent(EventBase): + """Data block delta event.""" + + type: Literal[EventType.DATA_BLOCK_DELTA] = EventType.DATA_BLOCK_DELTA + """Event type.""" + reply_id: str + """ID of the reply message this block belongs to.""" + block_id: str + """Unique identifier of the data block.""" + data: str + """Incremental base64-encoded data.""" + media_type: str + """MIME type of the data content.""" + + +class DataBlockEndEvent(EventBase): + """Data block end event.""" + + type: Literal[EventType.DATA_BLOCK_END] = EventType.DATA_BLOCK_END + """Event type.""" + reply_id: str + """ID of the reply message this block belongs to.""" + block_id: str + """Unique identifier of the data block.""" + + +class ThinkingBlockStartEvent(EventBase): + """Thinking block start event.""" + + type: Literal[ + EventType.THINKING_BLOCK_START + ] = EventType.THINKING_BLOCK_START + """Event type.""" + reply_id: str + """ID of the reply message this block belongs to.""" + block_id: str + """Unique identifier of the thinking block.""" + + +class ThinkingBlockDeltaEvent(EventBase): + """Thinking block delta event.""" + + type: Literal[ + EventType.THINKING_BLOCK_DELTA + ] = EventType.THINKING_BLOCK_DELTA + """Event type.""" + reply_id: str + """ID of the reply message this block belongs to.""" + block_id: str + """Unique identifier of the thinking block.""" + delta: str + """Incremental thinking text content.""" + + +class ThinkingBlockEndEvent(EventBase): + """Thinking block end event.""" + + type: Literal[EventType.THINKING_BLOCK_END] = EventType.THINKING_BLOCK_END + """Event type.""" + reply_id: str + """ID of the reply message this block belongs to.""" + block_id: str + """Unique identifier of the thinking block.""" + + +class HintBlockEvent(EventBase): + """One-shot hint block event. + + Unlike text/thinking blocks, hint blocks are not streamed — the + full content is available at creation time (team messages, + background tool results, user interruptions, …). A single event + carries the complete :class:`~agentscope.message.HintBlock`. + + The ``hint`` field mirrors :attr:`HintBlock.hint` and may be a + plain string or a list of :class:`TextBlock` / :class:`DataBlock` + for multimodal content. + """ + + type: Literal[EventType.HINT_BLOCK] = EventType.HINT_BLOCK + """Event type.""" + reply_id: str + """ID of the reply message this block belongs to.""" + block_id: str + """Unique identifier of the hint block.""" + source: str | None = None + """Sender or origin of this hint (e.g. ``"alice"``, ``"system"``).""" + hint: str | List[TextBlock | DataBlock] + """Complete hint content — ``str`` or ``list[TextBlock | DataBlock]``.""" + + +class ToolCallStartEvent(EventBase): + """Tool call start event.""" + + type: Literal[EventType.TOOL_CALL_START] = EventType.TOOL_CALL_START + """Event type.""" + reply_id: str + """ID of the reply message this tool call belongs to.""" + tool_call_id: str + """Unique identifier of the tool call.""" + tool_call_name: str + """Name of the tool being called.""" + + +class ToolCallDeltaEvent(EventBase): + """Tool call delta event.""" + + type: Literal[EventType.TOOL_CALL_DELTA] = EventType.TOOL_CALL_DELTA + """Event type.""" + reply_id: str + """ID of the reply message this tool call belongs to.""" + tool_call_id: str + """Unique identifier of the tool call.""" + delta: str + """Incremental tool call arguments (JSON fragment).""" + + +class ToolCallEndEvent(EventBase): + """Tool call end event.""" + + type: Literal[EventType.TOOL_CALL_END] = EventType.TOOL_CALL_END + """Event type.""" + reply_id: str + """ID of the reply message this tool call belongs to.""" + tool_call_id: str + """Unique identifier of the tool call.""" + + +class ToolResultStartEvent(EventBase): + """Tool result start event.""" + + type: Literal[EventType.TOOL_RESULT_START] = EventType.TOOL_RESULT_START + """Event type.""" + reply_id: str + """ID of the reply message this tool result belongs to.""" + tool_call_id: str + """ID of the corresponding tool call.""" + tool_call_name: str + """Name of the tool that was called.""" + + +class ToolResultTextDeltaEvent(EventBase): + """Tool result text delta event.""" + + type: Literal[ + EventType.TOOL_RESULT_TEXT_DELTA + ] = EventType.TOOL_RESULT_TEXT_DELTA + """Event type.""" + reply_id: str + """ID of the reply message this tool result belongs to.""" + tool_call_id: str + """ID of the corresponding tool call.""" + delta: str + """Incremental text content of the tool result.""" + + +class ToolResultDataDeltaEvent(EventBase): + """Tool result data delta event.""" + + type: Literal[ + EventType.TOOL_RESULT_DATA_DELTA + ] = EventType.TOOL_RESULT_DATA_DELTA + """Event type.""" + reply_id: str + """ID of the reply message this tool result belongs to.""" + tool_call_id: str + """ID of the corresponding tool call.""" + block_id: str = Field(default_factory=_generate_id) + """Unique identifier of the data block created by this event.""" + media_type: str + """MIME type of the binary content.""" + data: str | None = None + """Base64-encoded binary data, mutually exclusive with `url`.""" + url: str | None = None + """URL pointing to the binary content, mutually exclusive with `data`.""" + + +class ToolResultEndEvent(EventBase): + """Tool result end event.""" + + model_config = ConfigDict(use_enum_values=True) + + type: Literal[EventType.TOOL_RESULT_END] = EventType.TOOL_RESULT_END + """Event type.""" + reply_id: str + """ID of the reply message this tool result belongs to.""" + tool_call_id: str + """ID of the corresponding tool call.""" + state: ToolResultState + """Final execution state of the tool call.""" + metadata: dict[str, Any] = Field(default_factory=dict) + """Optional metadata attached to the tool result event.""" + + +class ExceedMaxItersEvent(EventBase): + """Exceeded max iteration event.""" + + type: Literal[EventType.EXCEED_MAX_ITERS] = EventType.EXCEED_MAX_ITERS + """Event type.""" + reply_id: str + """ID of the reply message associated with this run.""" + name: str + """Name of the agent.""" + + +class RequireUserConfirmEvent(EventBase): + """Require user confirm event.""" + + type: Literal[ + EventType.REQUIRE_USER_CONFIRM + ] = EventType.REQUIRE_USER_CONFIRM + """Event type.""" + reply_id: str + """ID of the reply message associated with this run.""" + tool_calls: List[ToolCallBlock] + """Tool calls pending user confirmation.""" + + +class RequireExternalExecutionEvent(EventBase): + """Require external execution event.""" + + type: Literal[ + EventType.REQUIRE_EXTERNAL_EXECUTION + ] = EventType.REQUIRE_EXTERNAL_EXECUTION + """Event type.""" + reply_id: str + """ID of the reply message associated with this run.""" + tool_calls: List[ToolCallBlock] + """Tool calls to be executed externally.""" + + +class ConfirmResult(BaseModel): + """Confirm result for a tool call.""" + + confirmed: bool + """Whether the user confirmed the tool call.""" + tool_call: ToolCallBlock + """The tool call that was confirmed or rejected.""" + rules: list[PermissionRule] | None = None + """The allowed permission rules for this tool call. This field is only + applicable when ``confirmed`` is True. In case user modification is + needed, complete permission rules are used here instead of references to + the suggested rules in ``RequireUserConfirmEvent``.""" + + +class UserConfirmResultEvent(EventBase): + """User confirm result event.""" + + type: Literal[ + EventType.USER_CONFIRM_RESULT + ] = EventType.USER_CONFIRM_RESULT + """Event type.""" + reply_id: str + """ID of the reply message associated with this run.""" + confirm_results: list[ConfirmResult] + """Confirmation results for each pending tool call.""" + + +class ExternalExecutionResultEvent(EventBase): + """External execution result event.""" + + type: Literal[ + EventType.EXTERNAL_EXECUTION_RESULT + ] = EventType.EXTERNAL_EXECUTION_RESULT + """Event type.""" + reply_id: str + """ID of the reply message associated with this run.""" + execution_results: List[ToolResultBlock] + """Results returned by the external executor.""" + + +class CustomEvent(EventBase): + """Generic extensible event for signals that don't fit a specific + ``AgentEvent`` subtype. + + Used by service-layer middleware to notify front-end subscribers + about state changes (task progress, team membership, permission + updates, …) without polluting the core agent event enum with + application-specific types. + + Front-end implementations should handle unknown ``name`` values + gracefully — skip with no error. + + Attributes: + name (`str`): + Identifies the kind of notification. Well-known values: + + - ``"state_updated"`` — agent state (tasks / permission) + changed during a tool call. + - ``"team_updated"`` — team membership changed (member + added / team created or dissolved). + + value (`dict`): + Arbitrary JSON-serializable payload whose schema depends + on ``name``. May be empty. + """ + + type: Literal[EventType.CUSTOM] = EventType.CUSTOM + """Event type discriminator.""" + name: str + """Kind of notification — see class docstring for well-known values.""" + value: dict = Field(default_factory=dict) + """Arbitrary payload.""" + + +AgentEvent: TypeAlias = ( + ReplyStartEvent + | ReplyEndEvent + | ExceedMaxItersEvent + | RequireUserConfirmEvent + | RequireExternalExecutionEvent + | ModelCallStartEvent + | ModelCallEndEvent + | TextBlockStartEvent + | TextBlockDeltaEvent + | TextBlockEndEvent + | DataBlockStartEvent + | DataBlockDeltaEvent + | DataBlockEndEvent + | ThinkingBlockStartEvent + | ThinkingBlockDeltaEvent + | ThinkingBlockEndEvent + | HintBlockEvent + | ToolCallStartEvent + | ToolCallDeltaEvent + | ToolCallEndEvent + | ToolResultStartEvent + | ToolResultTextDeltaEvent + | ToolResultDataDeltaEvent + | ToolResultEndEvent + | UserConfirmResultEvent + | ExternalExecutionResultEvent + | CustomEvent +) diff --git a/src/agentscope/exception/__init__.py b/src/agentscope/exception/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..966f77cbb9e2d6096127085674f7313a7929d8af --- /dev/null +++ b/src/agentscope/exception/__init__.py @@ -0,0 +1,22 @@ +# -*- coding: utf-8 -*- +"""The exception module in agentscope.""" + +from ._base import ( + AgentOrientedException, + DeveloperOrientedException, +) +from ._tool import ( + ToolInterruptedError, + ToolNotFoundError, + ToolJSONDecodeError, + ToolGroupInactiveError, +) + +__all__ = [ + "AgentOrientedException", + "DeveloperOrientedException", + "ToolInterruptedError", + "ToolNotFoundError", + "ToolJSONDecodeError", + "ToolGroupInactiveError", +] diff --git a/src/agentscope/exception/_base.py b/src/agentscope/exception/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..f8ae89f737d8e33465a2aeeeff927b1f6a847e76 --- /dev/null +++ b/src/agentscope/exception/_base.py @@ -0,0 +1,27 @@ +# -*- coding: utf-8 -*- +"""The base exception class in agentscope.""" + + +class AgentOrientedException(Exception): + """The base class for all agent-oriented exceptions. These exceptions are + expect to the captured and exposed to the agent during runtime, so that + agents can handle the error appropriately during the runtime. + """ + + def __init__(self, message: str): + """Initialize the exception with a message.""" + super().__init__(message) + self.message = message + + def __str__(self) -> str: + """Return the string representation of the exception.""" + return f"{self.__class__.__name__}: {self.message}" + + +class DeveloperOrientedException(Exception): + """The exception should be raised to the developers.""" + + def __init__(self, message: str): + """Initialize the exception with a message.""" + super().__init__(message) + self.message = message diff --git a/src/agentscope/exception/_tool.py b/src/agentscope/exception/_tool.py new file mode 100644 index 0000000000000000000000000000000000000000..9a470b34330474bed4c40c3bedbb90012f300e48 --- /dev/null +++ b/src/agentscope/exception/_tool.py @@ -0,0 +1,20 @@ +# -*- coding: utf-8 -*- +"""The tool-related exceptions in agentscope.""" + +from ._base import AgentOrientedException + + +class ToolNotFoundError(AgentOrientedException): + """Exception raised when a tool was not found.""" + + +class ToolInterruptedError(AgentOrientedException): + """Exception raised when a tool calling was interrupted by the user.""" + + +class ToolJSONDecodeError(AgentOrientedException): + """Exception raised when tool arguments fail JSON decoding or repair.""" + + +class ToolGroupInactiveError(AgentOrientedException): + """Exception raised when a tool group is inactive.""" diff --git a/src/agentscope/formatter/__init__.py b/src/agentscope/formatter/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fa61f28da1b4bfac1e5398deede77e29bbc33034 --- /dev/null +++ b/src/agentscope/formatter/__init__.py @@ -0,0 +1,62 @@ +# -*- coding: utf-8 -*- +"""The formatter module in agentscope.""" + +from ._formatter_base import FormatterBase +from ._dashscope_formatter import ( + DashScopeChatFormatter, + DashScopeMultiAgentFormatter, +) +from ._anthropic_formatter import ( + AnthropicChatFormatter, + AnthropicMultiAgentFormatter, +) +from ._openai_formatter import ( + OpenAIChatFormatter, + OpenAIMultiAgentFormatter, +) +from ._gemini_formatter import ( + GeminiChatFormatter, + GeminiMultiAgentFormatter, +) +from ._ollama_formatter import ( + OllamaChatFormatter, + OllamaMultiAgentFormatter, +) +from ._deepseek_formatter import ( + DeepSeekChatFormatter, + DeepSeekMultiAgentFormatter, +) +from ._openai_response_formatter import ( + OpenAIResponseFormatter, + OpenAIResponseMultiAgentFormatter, +) +from ._moonshot_formatter import ( + MoonshotChatFormatter, + MoonshotMultiAgentFormatter, +) +from ._xai_formatter import ( + XAIChatFormatter, + XAIMultiAgentFormatter, +) + +__all__ = [ + "FormatterBase", + "DashScopeChatFormatter", + "DashScopeMultiAgentFormatter", + "OpenAIChatFormatter", + "OpenAIMultiAgentFormatter", + "AnthropicChatFormatter", + "AnthropicMultiAgentFormatter", + "GeminiChatFormatter", + "GeminiMultiAgentFormatter", + "OllamaChatFormatter", + "OllamaMultiAgentFormatter", + "DeepSeekChatFormatter", + "DeepSeekMultiAgentFormatter", + "OpenAIResponseFormatter", + "OpenAIResponseMultiAgentFormatter", + "MoonshotChatFormatter", + "MoonshotMultiAgentFormatter", + "XAIChatFormatter", + "XAIMultiAgentFormatter", +] diff --git a/src/agentscope/formatter/_anthropic_formatter.py b/src/agentscope/formatter/_anthropic_formatter.py new file mode 100644 index 0000000000000000000000000000000000000000..e94db1639dde0acdaafbc229dafc82a871e2d38a --- /dev/null +++ b/src/agentscope/formatter/_anthropic_formatter.py @@ -0,0 +1,507 @@ +# -*- coding: utf-8 -*- +"""The Anthropic formatter module.""" +import base64 +import fnmatch +import json +from abc import ABC +from typing import Any + +import requests +from pydantic import Field + +from ._formatter_base import FormatterBase +from .._logging import logger +from ..message import ( + Msg, + TextBlock, + ThinkingBlock, + HintBlock, + DataBlock, + ToolCallBlock, + ToolResultBlock, + URLSource, + Base64Source, +) + + +class _AnthropicFormatterBase(FormatterBase, ABC): + """Mixin for formatting Anthropic formatters to avoid duplication between + AnthropicChatFormatter and AnthropicMultiAgentFormatter.""" + + # pylint: disable=too-many-branches + async def _format_messages( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Format message objects into Anthropic API format. + + Args: + msgs (`list[Msg]`): + The list of message objects to format. + + Returns: + `list[dict[str, Any]]`: + The formatted messages as a list of dictionaries. + + .. note:: Anthropic suggests always passing all previous thinking + blocks back to the API in subsequent calls to maintain reasoning + continuity. For more details, please refer to + `Anthropic's documentation + `_. + """ + self.assert_list_of_msgs(msgs) + + messages: list[dict] = [] + for msg in msgs: # pylint: disable=too-many-nested-blocks + content_blocks: list = [] + has_tool_result = False + + for block in msg.get_content_blocks(): + if ( + has_tool_result + and content_blocks + and not isinstance( + block, + ToolResultBlock, + ) + ): + messages.append( + {"role": "user", "content": content_blocks}, + ) + content_blocks = [] + has_tool_result = False + + if isinstance(block, TextBlock): + content_blocks.append( + {"type": "text", "text": block.text}, + ) + + elif isinstance(block, ThinkingBlock): + # Anthropic rejects thinking blocks without a valid + # signature ("Invalid `signature` in `thinking` block"). + # ThinkingBlocks from other providers (OpenAI, DeepSeek, + # ...) carry no signature, so drop them instead of + # forwarding an empty one. + signature = getattr(block, "signature", None) + if signature: + content_blocks.append( + { + "type": "thinking", + "thinking": block.thinking, + "signature": signature, + }, + ) + else: + logger.debug( + "Dropping ThinkingBlock without signature; " + "Anthropic requires a valid signature.", + ) + + elif isinstance(block, HintBlock): + if content_blocks: + role = "user" if has_tool_result else msg.role + messages.append( + {"role": role, "content": content_blocks}, + ) + content_blocks = [] + has_tool_result = False + + if isinstance(block.hint, str): + messages.append( + { + "role": "user", + "content": [ + {"type": "text", "text": block.hint}, + ], + }, + ) + else: + hint_parts: list[dict] = [] + for sub in block.hint: + if isinstance(sub, TextBlock): + hint_parts.append( + {"type": "text", "text": sub.text}, + ) + elif isinstance(sub, DataBlock): + formatted_sub = ( + self._format_anthropic_data_block(sub) + ) + if formatted_sub: + hint_parts.append(formatted_sub) + if hint_parts: + messages.append( + {"role": "user", "content": hint_parts}, + ) + + elif isinstance(block, DataBlock): + formatted_block = self._format_anthropic_data_block(block) + if formatted_block: + content_blocks.append(formatted_block) + + elif isinstance(block, ToolCallBlock): + content_blocks.append( + { + "type": "tool_use", + "id": block.id, + "name": block.name, + # Anthropic API expects input as a dict, not a + # JSON string. + "input": json.loads(block.input or "{}"), + }, + ) + + elif isinstance(block, ToolResultBlock): + # Only flush when we have non-tool-result content + # (i.e. the preceding assistant turn). Once + # `has_tool_result` is True we are already accumulating + # tool_results into the current user message, so we must + # NOT flush on each additional ToolResultBlock — doing so + # would split parallel results into separate user messages + # which strict endpoints (e.g. DeepSeek) reject with 400. + if content_blocks and not has_tool_result: + role = "user" if has_tool_result else msg.role + messages.append( + {"role": role, "content": content_blocks}, + ) + content_blocks = [] + + tool_result_content: list[dict] = [] + output = block.output + if isinstance(output, str): + tool_result_content.append( + {"type": "text", "text": output}, + ) + else: + for out_block in output: + if isinstance(out_block, TextBlock): + tool_result_content.append( + {"type": "text", "text": out_block.text}, + ) + elif isinstance(out_block, DataBlock): + fmt_block = self._format_anthropic_data_block( + out_block, + ) + if fmt_block: + tool_result_content.append(fmt_block) + else: + source = out_block.source + main_type = source.media_type.split("/")[0] + if isinstance(source, URLSource): + fallback = ( + f"[{main_type} file returned, " + f"URL: {source.url}]" + ) + else: + fallback = ( + f"[{main_type} file returned, " + f"type: {source.media_type}]" + ) + tool_result_content.append( + {"type": "text", "text": fallback}, + ) + + content_blocks.append( + { + "type": "tool_result", + "tool_use_id": block.id, + "content": tool_result_content, + }, + ) + # Anthropic requires tool_result to be in a "user" message. + has_tool_result = True + + else: + logger.warning( + "Unsupported block type %s in the message, skipped.", + type(block), + ) + + if content_blocks: + # Anthropic requires `tool_result` blocks to be in a `user` + # message regardless of the containing Msg's role. + role = "user" if has_tool_result else msg.role + messages.append( + { + "role": role, + "content": content_blocks, + }, + ) + + return messages + + def _format_anthropic_data_block( + self, + block: DataBlock, + ) -> dict[str, Any] | None: + """Format a DataBlock into Anthropic API format. + + Args: + block (`DataBlock`): + The data block to format. + + Returns: + `dict[str, Any] | None`: + The formatted data block, or None if the media type is not + supported. + """ + source = block.source + media_type = source.media_type + + # Check if media type is supported + if not any( + fnmatch.fnmatch(media_type, pattern) + for pattern in self.supported_input_media_types + ): + logger.warning( + "Media type %s is not supported, skipped.", + media_type, + ) + return None + + # Anthropic only supports images + if not media_type.startswith("image/"): + logger.warning( + "Anthropic only supports image data, got %s, skipped.", + media_type, + ) + return None + + return self._format_image_source(source) + + @staticmethod + def _format_image_source( + source: URLSource | Base64Source, + ) -> dict[str, Any]: + """Format an image source into Anthropic API format. + + Args: + source (`URLSource | Base64Source`): + The image source to format. + + Returns: + `dict[str, Any]`: + The formatted image source. + """ + if isinstance(source, Base64Source): + return { + "type": "image", + "source": { + "type": "base64", + "media_type": source.media_type, + "data": source.data, + }, + } + elif isinstance(source, URLSource): + url = str(source.url) + if url.startswith("file://"): + # Local file - read and convert to base64 + file_path = url.removeprefix("file://") + with open(file_path, "rb") as f: + data = base64.b64encode(f.read()).decode("utf-8") + return { + "type": "image", + "source": { + "type": "base64", + "media_type": source.media_type, + "data": data, + }, + } + else: + # Remote URL - download and convert to base64 + response = requests.get(url, timeout=30) + response.raise_for_status() + data = base64.b64encode(response.content).decode("utf-8") + return { + "type": "image", + "source": { + "type": "base64", + "media_type": source.media_type, + "data": data, + }, + } + else: + raise ValueError(f"Unsupported source type: {type(source)}") + + +class AnthropicChatFormatter(_AnthropicFormatterBase): + """The Anthropic formatter class for chatbot scenario, where only a user + and an agent are involved. We use the `role` field to identify different + entities in the conversation. + """ + + input_types: list[str] = Field( + default_factory=lambda: ["text/plain", "image/*"], + description=( + "The supported input types. " + 'Defaults to ``["text/plain", "image/*"]``.' + ), + ) + + async def format( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Format message objects into Anthropic API format. + + Args: + msgs (`list[Msg]`): + The list of message objects to format. + + Returns: + `list[dict[str, Any]]`: + The formatted messages as a list of dictionaries. + + .. note:: Anthropic suggests always passing all previous thinking + blocks back to the API in subsequent calls to maintain reasoning + continuity. For more details, please refer to + `Anthropic's documentation + `_. + """ + return await self._format_messages(msgs) + + +class AnthropicMultiAgentFormatter(_AnthropicFormatterBase): + """Anthropic formatter for multi-agent conversations, where more than + a user and an agent are involved. + """ + + conversation_history_prompt: str = Field( + default=( + "# Conversation History\n" + "The content between tags contains " + "your conversation history\n" + ), + description="The prompt to use for the conversation history section.", + ) + + input_types: list[str] = Field( + default_factory=lambda: ["text/plain", "image/*"], + description=( + "The supported input types. " + 'Defaults to ``["text/plain", "image/*"]``.' + ), + ) + + async def format(self, msgs: list[Msg]) -> list[dict[str, Any]]: + """Format input messages into the structure required by the Anthropic + API for multi-agent conversations.""" + self.assert_list_of_msgs(msgs) + + formatted_msgs = [] + start_index = 0 + if len(msgs) > 0 and msgs[0].role == "system": + formatted_msgs.append( + await self._format_system_message(msgs[0]), + ) + start_index = 1 + + is_first_agent_message = True + async for typ, group in self._group_messages(msgs[start_index:]): + match typ: + case "tool_sequence": + formatted_msgs.extend( + await self._format_messages(group), + ) + case "agent_message": + formatted_msgs.extend( + await self._format_agent_message( + group, + is_first_agent_message, + ), + ) + is_first_agent_message = False + + return formatted_msgs + + async def _format_agent_message( + self, + msgs: list[Msg], + is_first: bool, + ) -> list[dict[str, Any]]: + """Format agent messages into conversation history.""" + conversation_blocks = [] + accumulated_text = [] + + for msg in msgs: + agent_name = msg.name or "Agent" + agent_text_parts = [] + + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + agent_text_parts.append(block.text) + elif isinstance(block, DataBlock): + formatted_block = self._format_anthropic_data_block(block) + if formatted_block: + if accumulated_text: + conversation_blocks.append( + { + "type": "text", + "text": "\n".join(accumulated_text), + }, + ) + accumulated_text = [] + conversation_blocks.append(formatted_block) + + if agent_text_parts: + agent_message = f"{agent_name}: {' '.join(agent_text_parts)}" + accumulated_text.append(agent_message) + + if accumulated_text: + conversation_blocks.append( + { + "type": "text", + "text": "\n".join(accumulated_text), + }, + ) + + if conversation_blocks and is_first: + if conversation_blocks[0].get("text"): + conversation_blocks[0]["text"] = ( + self.conversation_history_prompt + + "\n" + + conversation_blocks[0]["text"] + ) + else: + conversation_blocks.insert( + 0, + { + "type": "text", + "text": self.conversation_history_prompt + + "\n", + }, + ) + + if conversation_blocks[-1].get("text"): + conversation_blocks[-1]["text"] += "\n" + else: + conversation_blocks.append( + {"type": "text", "text": ""}, + ) + + if conversation_blocks: + return [ + { + "role": "user", + "content": conversation_blocks, + }, + ] + + return [] + + @staticmethod + async def _format_system_message(msg: Msg) -> dict[str, Any]: + """Format a system message.""" + text_parts = [] + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + text_parts.append(block.text) + + return { + "role": "system", + "content": [ + { + "type": "text", + "text": "\n".join(text_parts), + }, + ], + } diff --git a/src/agentscope/formatter/_dashscope_formatter.py b/src/agentscope/formatter/_dashscope_formatter.py new file mode 100644 index 0000000000000000000000000000000000000000..bd3b5be71bcc51235aa26390edc8bc43522c69b8 --- /dev/null +++ b/src/agentscope/formatter/_dashscope_formatter.py @@ -0,0 +1,563 @@ +# -*- coding: utf-8 -*- +"""The DashScope formatter module (OpenAI-compatible format).""" + +import base64 +from typing import Any +from fnmatch import fnmatch +from abc import ABC + +from pydantic import Field + +from ._formatter_base import FormatterBase +from .._logging import logger +from ..message import ( + Msg, + TextBlock, + ThinkingBlock, + ToolResultBlock, + URLSource, + DataBlock, + ToolCallBlock, + Base64Source, + HintBlock, +) + + +class _DashScopeFormatterBase(FormatterBase, ABC): + """Base class for DashScope formatters (OpenAI-compatible format), + providing shared data block formatting logic.""" + + input_types: list[str] = Field( + default_factory=lambda: [ + "text/plain", + "image/*", + "audio/*", + "video/*", + ], + description=( + "The supported input types, aligned with the model card's " + "``input_types`` field. Media types (non ``text/plain`` / " + "``application/x-thinking`` entries) are used to filter " + "``DataBlock``\\s; ``application/x-thinking`` enables passing " + "``reasoning_content`` back to the API." + ), + ) + + @property + def supported_input_media_types(self) -> list[str]: + """Derive supported media types from :attr:`input_types`, excluding + ``text/plain`` and ``application/x-thinking``.""" + return [ + t + for t in self.input_types + if t not in ("text/plain", "application/x-thinking") + ] + + @property + def supports_thinking_input(self) -> bool: + """Return ``True`` if ``application/x-thinking`` is listed in + :attr:`input_types`, meaning the model accepts ``reasoning_content`` + in the conversation history.""" + return "application/x-thinking" in self.input_types + + def _format_dashscope_data_block( + self, + block: DataBlock, + ) -> dict[str, Any] | None: + """Format a DataBlock into the OpenAI-compatible format for + DashScope API. + + Supports: + - Images: ``{"type": "image_url", "image_url": {"url": ...}}`` + - Videos: ``{"type": "video_url", "video_url": {"url": ...}}`` + - Audio: ``{"type": "input_audio", "input_audio": {...}}`` + + Args: + block (`DataBlock`): + The DataBlock to format. + + Returns: + `dict[str, Any] | None`: + A dictionary representing the formatted DataBlock, or ``None`` + if the media type is unsupported. + """ + if not any( + fnmatch(block.source.media_type, pattern) + for pattern in self.supported_input_media_types + ): + logger.warning( + "Unsupported media type %s for DashScope API. Supported " + "types: %s. This block will be skipped.", + block.source.media_type, + ", ".join(self.supported_input_media_types), + ) + return None + + main_type = block.source.media_type.split("/")[0] + + if main_type == "image": + return self._format_image_source(block.source) + + if main_type == "video": + return self._format_video_source(block.source) + + if main_type == "audio": + return self._format_audio_source(block.source) + + logger.warning( + "Unsupported main media type %s for DashScope API. " + "This block will be skipped.", + main_type, + ) + return None + + @staticmethod + def _format_image_source( + source: URLSource | Base64Source, + ) -> dict[str, Any]: + """Convert an image source to OpenAI-compatible ``image_url`` format. + + Local ``file://`` URLs are read from disk and converted to base64 + data URIs. Remote URLs are passed through unchanged. + """ + if isinstance(source, Base64Source): + url = f"data:{source.media_type};base64,{source.data}" + elif isinstance(source, URLSource): + url_str = str(source.url) + if url_str.startswith("file://"): + local_path = url_str.removeprefix("file://") + with open(local_path, "rb") as f: + encoded = base64.b64encode(f.read()).decode("utf-8") + url = f"data:{source.media_type};base64,{encoded}" + else: + url = url_str + else: + raise ValueError(f"Unsupported image source type: {type(source)}") + + return { + "type": "image_url", + "image_url": {"url": url}, + } + + @staticmethod + def _format_video_source( + source: URLSource | Base64Source, + ) -> dict[str, Any]: + """Convert a video source to DashScope's ``video_url`` format + (OpenAI-compatible extension). + + Local ``file://`` URLs are read from disk and converted to base64 + data URIs. Remote URLs are passed through unchanged. + """ + if isinstance(source, Base64Source): + url = f"data:{source.media_type};base64,{source.data}" + elif isinstance(source, URLSource): + url_str = str(source.url) + if url_str.startswith("file://"): + local_path = url_str.removeprefix("file://") + with open(local_path, "rb") as f: + encoded = base64.b64encode(f.read()).decode("utf-8") + url = f"data:{source.media_type};base64,{encoded}" + else: + url = url_str + else: + raise ValueError(f"Unsupported video source type: {type(source)}") + + return { + "type": "video_url", + "video_url": {"url": url}, + } + + @staticmethod + def _format_audio_source( + source: URLSource | Base64Source, + ) -> dict[str, Any]: + """Convert an audio source to DashScope ``input_audio`` format. + + DashScope's compatible API accepts URLs directly in the ``data`` + field (unlike standard OpenAI which requires base64). Local files + are still read and base64-encoded. + """ + if isinstance(source, Base64Source): + fmt = source.media_type.split("/")[-1] + return { + "type": "input_audio", + "input_audio": { + "data": source.data, + "format": fmt, + }, + } + + if isinstance(source, URLSource): + url_str = str(source.url) + fmt = source.media_type.split("/")[-1] + if url_str.startswith("file://"): + local_path = url_str.removeprefix("file://") + with open(local_path, "rb") as f: + data = base64.b64encode(f.read()).decode("utf-8") + return { + "type": "input_audio", + "input_audio": { + "data": data, + "format": fmt, + }, + } + else: + return { + "type": "input_audio", + "input_audio": { + "data": url_str, + "format": fmt, + }, + } + + raise ValueError(f"Unsupported audio source type: {type(source)}") + + +class DashScopeChatFormatter(_DashScopeFormatterBase): + """The DashScope formatter class for chatbot scenario (OpenAI-compatible + format), where only a user and an agent are involved. We use the ``role`` + field to identify different entities in the conversation. + + This formatter outputs messages in the OpenAI Chat Completions format, + with DashScope-specific extensions for video (``video_url``) and + thinking (``reasoning_content``). + """ + + # pylint: disable=too-many-branches + async def format( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Format message objects into DashScope OpenAI-compatible format. + + Args: + msgs (`list[Msg]`): + The list of message objects to format. + + Returns: + `list[dict[str, Any]]`: + The formatted messages as a list of dictionaries. + """ + self.assert_list_of_msgs(msgs) + + formatted_msgs: list[dict] = [] + i = 0 + while i < len(msgs): + msg = msgs[i] + content_blocks: list[dict] = [] + tool_calls = [] + thinking_parts: list[str] = [] + + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + content_blocks.append({"type": "text", "text": block.text}) + + elif isinstance(block, DataBlock): + formatted_block = self._format_dashscope_data_block( + block, + ) + if formatted_block: + content_blocks.append(formatted_block) + + elif isinstance(block, HintBlock): + if content_blocks or tool_calls or thinking_parts: + msg_openai: dict[str, Any] = { + "role": msg.role, + "content": content_blocks or None, + } + if tool_calls: + msg_openai["tool_calls"] = tool_calls + if thinking_parts: + msg_openai["reasoning_content"] = "\n".join( + thinking_parts, + ) + formatted_msgs.append(msg_openai) + content_blocks = [] + tool_calls = [] + thinking_parts = [] + + if isinstance(block.hint, str): + formatted_msgs.append( + { + "role": "user", + "content": [ + {"type": "text", "text": block.hint}, + ], + }, + ) + else: + hint_parts: list[dict] = [] + for sub in block.hint: + if isinstance(sub, TextBlock): + hint_parts.append( + {"type": "text", "text": sub.text}, + ) + elif isinstance(sub, DataBlock): + formatted_sub = ( + self._format_dashscope_data_block( + sub, + ) + ) + if formatted_sub: + hint_parts.append(formatted_sub) + if hint_parts: + formatted_msgs.append( + {"role": "user", "content": hint_parts}, + ) + + elif isinstance(block, ToolCallBlock): + tool_calls.append( + { + "id": block.id, + "type": "function", + "function": { + "name": block.name, + "arguments": block.input, + }, + }, + ) + + elif isinstance(block, ThinkingBlock): + if self.supports_thinking_input: + thinking_parts.append(block.thinking) + + elif isinstance(block, ToolResultBlock): + if content_blocks or tool_calls or thinking_parts: + msg_flush: dict[str, Any] = { + "role": msg.role, + "content": content_blocks or None, + } + if tool_calls: + msg_flush["tool_calls"] = tool_calls + if thinking_parts: + msg_flush["reasoning_content"] = "\n".join( + thinking_parts, + ) + formatted_msgs.append(msg_flush) + content_blocks = [] + tool_calls = [] + thinking_parts = [] + + ( + textual_output, + multimodal_data, + ) = self.convert_tool_result_to_string(block.output) + + formatted_msgs.append( + { + "role": "tool", + "tool_call_id": block.id, + "content": textual_output, + "name": block.name, + }, + ) + + if multimodal_data: + promo_content = [] + for item in multimodal_data: + if isinstance(item, TextBlock): + promo_content.append( + {"type": "text", "text": item.text}, + ) + elif isinstance(item, DataBlock): + fmt_item = self._format_dashscope_data_block( + item, + ) + if fmt_item is not None: + promo_content.append(fmt_item) + if promo_content: + formatted_msgs.append( + { + "role": "user", + "content": promo_content, + }, + ) + + else: + logger.warning( + "Unsupported block type %s in the message, skipped.", + type(block), + ) + + msg_dashscope: dict[str, Any] = { + "role": msg.role, + "content": content_blocks or None, + } + + if tool_calls: + msg_dashscope["tool_calls"] = tool_calls + + if thinking_parts: + msg_dashscope["reasoning_content"] = "\n".join(thinking_parts) + + if ( + msg_dashscope["content"] + or msg_dashscope.get("tool_calls") + or msg_dashscope.get("reasoning_content") + ): + formatted_msgs.append(msg_dashscope) + + i += 1 + + return formatted_msgs + + +class DashScopeMultiAgentFormatter(_DashScopeFormatterBase): + """DashScope formatter for multi-agent conversations (OpenAI-compatible + format), where more than a user and an agent are involved. + + .. note:: This formatter will combine previous messages (except tool + calls/results) into a history section in the first system message with + the conversation history prompt. + + .. note:: For tool calls/results, they will be presented as separate + messages as required by the API. Therefore, the tool calls/results + messages are expected to be placed at the end of the input messages. + + .. tip:: Telling the assistant's name in the system prompt is very + important in multi-agent conversations. So that LLM can know who it + is playing as. + """ + + conversation_history_prompt: str = Field( + description="The conversation history prompt.", + default=( + "# Conversation History\n" + "The content between tags contains " + "your conversation history\n" + ), + ) + + async def format(self, msgs: list[Msg]) -> list[dict]: + """Format input messages into the structure required by the DashScope + OpenAI-compatible API. + + To support multi-agent conversations, this formatter processes messages + as follows: + + - Prepends an instruction before the first conversation history + section. + - Combines conversation turns into a history section, where each entry + is formatted as ``{name}: {content}``. + - Wraps the conversation history with ```` and ```` + tags. + + Returns: + `list[dict[str, Any]]`: + A list of dictionaries formatted for the DashScope API. + """ + + formatted_msgs = [] + start_index = 0 + if len(msgs) > 0 and msgs[0].role == "system": + formatted_msgs.append( + await self._format_system_message(msgs[0]), + ) + start_index = 1 + + is_first_agent_message = True + async for typ, group in self._group_messages(msgs[start_index:]): + match typ: + case "tool_sequence": + formatted_msgs.extend( + await self._format_tool_sequence(group), + ) + case "agent_message": + formatted_msgs.extend( + await self._format_agent_message( + group, + is_first_agent_message, + ), + ) + is_first_agent_message = False + + return formatted_msgs + + async def _format_tool_sequence( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Given a sequence of tool call/result messages, format them into + the required format for the DashScope API.""" + return await DashScopeChatFormatter( + input_types=self.input_types, + ).format(msgs) + + async def _format_agent_message( + self, + msgs: list[Msg], + is_first: bool = True, + ) -> list[dict[str, Any]]: + """Given a sequence of messages without tool calls/results, format + them into a user message with conversation history tags.""" + if is_first: + conversation_history_prompt = self.conversation_history_prompt + else: + conversation_history_prompt = "" + + formatted_msgs: list[dict] = [] + conversation_blocks: list = [] + accumulated_text = [] + media_blocks: list[dict] = [] + + for msg in msgs: + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + accumulated_text.append(f"{msg.name}: {block.text}") + + elif isinstance(block, DataBlock): + formatted_block = self._format_dashscope_data_block( + block, + ) + if formatted_block is not None: + media_blocks.append(formatted_block) + + if accumulated_text: + conversation_blocks.append( + {"text": "\n".join(accumulated_text)}, + ) + + if conversation_blocks: + if conversation_blocks[0].get("text"): + conversation_blocks[0]["text"] = ( + conversation_history_prompt + + "\n" + + conversation_blocks[0]["text"] + ) + else: + conversation_blocks.insert( + 0, + {"text": conversation_history_prompt + "\n"}, + ) + + if conversation_blocks[-1].get("text"): + conversation_blocks[-1]["text"] += "\n" + else: + conversation_blocks.append({"text": ""}) + + conversation_blocks_text = "\n".join( + b.get("text", "") for b in conversation_blocks + ) + + content_list: list[dict[str, Any]] = [] + if conversation_blocks_text: + content_list.append( + {"type": "text", "text": conversation_blocks_text}, + ) + content_list.extend(media_blocks) + + if content_list: + formatted_msgs.append({"role": "user", "content": content_list}) + + return formatted_msgs + + @staticmethod + async def _format_system_message( + msg: Msg, + ) -> dict[str, Any]: + """Format system message for DashScope API.""" + return { + "role": "system", + "content": msg.get_text_content(), + } diff --git a/src/agentscope/formatter/_deepseek_formatter.py b/src/agentscope/formatter/_deepseek_formatter.py new file mode 100644 index 0000000000000000000000000000000000000000..ff9e4077d43ac17f001dc2ef5ba1f18af8a1b7a7 --- /dev/null +++ b/src/agentscope/formatter/_deepseek_formatter.py @@ -0,0 +1,313 @@ +# -*- coding: utf-8 -*- +"""The DeepSeek formatter module.""" +from typing import Any + +from pydantic import Field + +from ._formatter_base import FormatterBase +from .._logging import logger +from ..message import ( + Msg, + TextBlock, + DataBlock, + ThinkingBlock, + HintBlock, + ToolCallBlock, + ToolResultBlock, +) + + +class DeepSeekChatFormatter(FormatterBase): + """The DeepSeek formatter class for chatbot scenario, where only a user + and an agent are involved. We use the `role` field to identify different + entities in the conversation. + """ + + input_types: list[str] = Field( + default_factory=lambda: ["text/plain"], + description=( + 'The supported input types. Defaults to ``["text/plain"]`` ' + "(DeepSeek does not support multimodal input)." + ), + ) + + async def format( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Format message objects into DeepSeek API format. + + Args: + msgs (`list[Msg]`): + The list of message objects to format. + + Returns: + `list[dict[str, Any]]`: + The formatted messages as a list of dictionaries. + """ + self.assert_list_of_msgs(msgs) + + messages: list[dict] = [] + for msg in msgs: + content_blocks: list = [] + reasoning_content_blocks: list = [] + tool_calls = [] + + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + content_blocks.append({"type": "text", "text": block.text}) + + elif isinstance(block, ThinkingBlock): + reasoning_content_blocks.append(block.thinking) + + elif isinstance(block, HintBlock): + if ( + content_blocks + or tool_calls + or reasoning_content_blocks + ): + content_text = "\n".join( + b.get("text", "") for b in content_blocks + ) + msg_flush_hint: dict[str, Any] = { + "role": msg.role, + "content": content_text + or (None if tool_calls else ""), + } + if msg.role == "assistant": + msg_flush_hint["reasoning_content"] = ( + "\n".join(reasoning_content_blocks) + if reasoning_content_blocks + else "" + ) + if tool_calls: + msg_flush_hint["tool_calls"] = tool_calls + messages.append(msg_flush_hint) + content_blocks = [] + reasoning_content_blocks = [] + tool_calls = [] + + if isinstance(block.hint, str): + messages.append( + {"role": "user", "content": block.hint}, + ) + else: + hint_text_parts: list[str] = [] + for sub in block.hint: + if isinstance(sub, TextBlock): + hint_text_parts.append(sub.text) + elif isinstance(sub, DataBlock): + hint_text_parts.append( + f"[{sub.source.media_type} attached, " + "not supported by this provider]", + ) + if hint_text_parts: + messages.append( + { + "role": "user", + "content": "\n".join(hint_text_parts), + }, + ) + + elif isinstance(block, ToolCallBlock): + tool_calls.append( + { + "id": block.id, + "type": "function", + "function": { + "name": block.name, + "arguments": block.input, + }, + }, + ) + + elif isinstance(block, ToolResultBlock): + if ( + content_blocks + or tool_calls + or reasoning_content_blocks + ): + content_text = "\n".join( + b.get("text", "") for b in content_blocks + ) + msg_flush: dict[str, Any] = { + "role": msg.role, + "content": content_text + or (None if tool_calls else ""), + } + if msg.role == "assistant": + msg_flush["reasoning_content"] = ( + "\n".join(reasoning_content_blocks) + if reasoning_content_blocks + else "" + ) + if tool_calls: + msg_flush["tool_calls"] = tool_calls + messages.append(msg_flush) + content_blocks = [] + reasoning_content_blocks = [] + tool_calls = [] + + textual_output, _ = self.convert_tool_result_to_string( + block.output, + ) + messages.append( + { + "role": "tool", + "tool_call_id": block.id, + "content": textual_output, + "name": block.name, + }, + ) + + else: + logger.warning( + "Unsupported block type %s in the message, skipped.", + type(block), + ) + + content_msg = "\n".join(b.get("text", "") for b in content_blocks) + + msg_deepseek: dict[str, Any] = { + "role": msg.role, + "content": content_msg or (None if tool_calls else ""), + } + + # DeepSeek requires `reasoning_content` to be present on ALL + # assistant messages in multi-turn conversations that use thinking + # mode. When switching from a non-thinking model, historical + # messages will have no ThinkingBlock, so we always include the + # field (as None) for assistant messages to keep the context valid. + if msg.role == "assistant": + msg_deepseek["reasoning_content"] = ( + "\n".join(reasoning_content_blocks) + if reasoning_content_blocks + else "" + ) + + if tool_calls: + msg_deepseek["tool_calls"] = tool_calls + + if ( + msg_deepseek["content"] + or msg_deepseek.get("tool_calls") + or reasoning_content_blocks + ): + messages.append(msg_deepseek) + + return messages + + +class DeepSeekMultiAgentFormatter(FormatterBase): + """ + DeepSeek formatter for multi-agent conversations, where more than + a user and an agent are involved. + """ + + conversation_history_prompt: str = Field( + default=( + "# Conversation History\n" + "The content between tags contains " + "your conversation history\n" + ), + description="The prompt to use for the conversation history section.", + ) + + input_types: list[str] = Field( + default_factory=lambda: ["text/plain"], + description=( + 'The supported input types. Defaults to ``["text/plain"]`` ' + "(DeepSeek does not support multimodal input)." + ), + ) + + async def format(self, msgs: list[Msg]) -> list[dict[str, Any]]: + """Format input messages into the structure required by the DeepSeek + API for multi-agent conversations.""" + self.assert_list_of_msgs(msgs) + + formatted_msgs = [] + start_index = 0 + if len(msgs) > 0 and msgs[0].role == "system": + formatted_msgs.append( + await self._format_system_message(msgs[0]), + ) + start_index = 1 + + is_first_agent_message = True + async for typ, group in self._group_messages(msgs[start_index:]): + match typ: + case "tool_sequence": + formatted_msgs.extend( + await self._format_tool_sequence(group), + ) + case "agent_message": + formatted_msgs.extend( + await self._format_agent_message( + group, + is_first_agent_message, + ), + ) + is_first_agent_message = False + + return formatted_msgs + + async def _format_tool_sequence( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Given a sequence of tool call/result messages, format them into + the required format for the DeepSeek API.""" + return await DeepSeekChatFormatter( + input_types=self.input_types, + ).format(msgs) + + async def _format_agent_message( + self, + msgs: list[Msg], + is_first: bool = True, + ) -> list[dict[str, Any]]: + """Given a sequence of messages without tool calls/results, format + them into the required format for the DeepSeek API.""" + + if is_first: + conversation_history_prompt = self.conversation_history_prompt + else: + conversation_history_prompt = "" + + formatted_msgs: list[dict] = [] + accumulated_text = [] + + for msg in msgs: + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + accumulated_text.append(f"{msg.name}: {block.text}") + + conversation_blocks_text = "" + if accumulated_text: + conversation_blocks_text = ( + conversation_history_prompt + + "\n" + + "\n".join(accumulated_text) + + "\n" + ) + + if conversation_blocks_text: + formatted_msgs.append( + { + "role": "user", + "content": conversation_blocks_text, + }, + ) + + return formatted_msgs + + @staticmethod + async def _format_system_message( + msg: Msg, + ) -> dict[str, Any]: + """Format system message for DeepSeek API.""" + return { + "role": "system", + "content": msg.get_text_content(), + } diff --git a/src/agentscope/formatter/_formatter_base.py b/src/agentscope/formatter/_formatter_base.py new file mode 100644 index 0000000000000000000000000000000000000000..12edcd3e0253eaf62843ce4622d304a751b63f4e --- /dev/null +++ b/src/agentscope/formatter/_formatter_base.py @@ -0,0 +1,217 @@ +# -*- coding: utf-8 -*- +"""The formatter module.""" +import base64 +import mimetypes +import tempfile +from abc import abstractmethod +from fnmatch import fnmatch +from typing import Any, List, AsyncGenerator + +import shortuuid +from pydantic import BaseModel, Field + +from ..message import ( + Msg, + DataBlock, + TextBlock, + URLSource, + Base64Source, +) + + +class FormatterBase(BaseModel): + """The base class for formatters.""" + + input_types: list[str] = Field( + default_factory=lambda: ["text/plain"], + description=( + "The supported input types, aligned with the model card's " + "``input_types`` field. Entries other than ``text/plain`` and " + "``application/x-thinking`` are treated as media-type patterns " + "(glob-style, e.g. ``image/*``, ``audio/mp3``) that control which " + "``DataBlock``\\s are forwarded to the API." + ), + ) + """The supported input types for this formatter, aligned with the model + card's ``input_types`` field.""" + + @property + def supported_input_media_types(self) -> list[str]: + """Derive the accepted media-type patterns from :attr:`input_types` by + excluding ``text/plain`` and ``application/x-thinking``.""" + return [ + t + for t in self.input_types + if t not in ("text/plain", "application/x-thinking") + ] + + @abstractmethod + async def format(self, *args: Any, **kwargs: Any) -> list[dict[str, Any]]: + """Format the Msg objects to a list of dictionaries that satisfy the + API requirements.""" + + @staticmethod + def assert_list_of_msgs(msgs: list[Msg]) -> None: + """Assert that the input is a list of Msg objects. + + Args: + msgs (`list[Msg]`): + A list of Msg objects to be validated. + """ + if not isinstance(msgs, list): + raise TypeError("Input must be a list of Msg objects.") + + for msg in msgs: + if not isinstance(msg, Msg): + raise TypeError( + f"Expected Msg object, got {type(msg)} instead.", + ) + + def convert_tool_result_to_string( + self, + output: str | List[TextBlock | DataBlock], + ) -> tuple[str, list[TextBlock | DataBlock]]: + """Turn the tool result list into a textual output to be compatible + with the LLM API that doesn't support multimodal data in the tool + result. + + For URL-based images, the URL is included in the list. For + base64-encoded images, the local file path where the image is saved + is included in the returned list. + + Args: + output (`str | List[TextBlock | DataBlock]`): + The output of the tool response, including text and multimodal + data like images and audio. + + Returns: + `tuple[str, list[TextBlock | DataBlock]]`: + A tuple containing the textual representation of the tool + result and a list of blocks to be promoted as a user message. + """ + + if isinstance(output, str): + return output, [] + + textual_output = [] + multimodal_data: list = [] + + for block in output: + if isinstance(block, TextBlock): + textual_output.append(block.text) + + elif isinstance(block, DataBlock): + main_type = block.source.media_type.split("/")[0] + + if any( + fnmatch(block.source.media_type, _) + for _ in self.supported_input_media_types + ): + # If supported, promote the block + + # Create an identifier for such multimodal data for + # accurate reference (in terms of order, position, etc.) + identifier = shortuuid.uuid() + + textual_output.append( + f"A(n) {main_type} file is returned " + f"and will be presented to you with the identifier " + f"[{identifier}].", + ) + multimodal_data.extend( + [ + TextBlock( + text=f"- {identifier} ({main_type} file): ", + ), + block, + ], + ) + + # For unsupported media types, if it's a URL, include it in + # the textual output; if it's base64 data, save it locally + # and include the file path in the textual output. + # Note if you don't want to save the local file, you should + # transform the base64 data in the tool execution hook + # rather than changing the formatter. + elif isinstance(block.source, URLSource): + textual_output.append( + f"A(n) {main_type} file is returned " + f"and can be accessed at the URL: {block.source.url}." + f"", + ) + + elif isinstance(block.source, Base64Source): + # Have to save the base64 data locally + extension = mimetypes.guess_extension( + block.source.media_type, + ) + with tempfile.NamedTemporaryFile( + suffix=extension, + delete=False, + ) as temp_file: + decoded_data = base64.b64decode(block.source.data) + temp_file.write(decoded_data) + textual_output.append( + f"A(n) {main_type} file is " + f"returned and saved locally at: {temp_file.name}." + f"", + ) + + # Add system reminder tags if there is multimodal data to be promoted + if multimodal_data: + multimodal_data = [ + TextBlock( + text="The multimodal data and their " + "identifiers are listed as follows:", + ), + *multimodal_data, + TextBlock( + text="", + ), + ] + + return "\n".join(textual_output), multimodal_data + + @staticmethod + async def _group_messages(msgs: list[Msg]) -> AsyncGenerator: + """Group messages into tool sequences and agent messages. + + Args: + msgs (`list[Msg]`): + A list of Msg objects to be grouped. + """ + group_type = None + group = [] + for msg in msgs: + if group_type is None: + if msg.get_content_blocks( + "tool_call", + ) or msg.get_content_blocks("tool_result"): + group_type = "tool_sequence" + else: + group_type = "agent_message" + group.append(msg) + continue + + if group_type == "tool_sequence": + if msg.has_content_blocks( + "tool_call", + ) or msg.has_content_blocks("tool_result"): + group.append(msg) + else: + yield group_type, group + group = [msg] + group_type = "agent_message" + + elif group_type == "agent_message": + if msg.has_content_blocks( + "tool_call", + ) or msg.has_content_blocks("tool_result"): + yield group_type, group + group = [msg] + group_type = "tool_sequence" + else: + group.append(msg) + + if group_type: + yield group_type, group diff --git a/src/agentscope/formatter/_gemini_formatter.py b/src/agentscope/formatter/_gemini_formatter.py new file mode 100644 index 0000000000000000000000000000000000000000..e8358fab5f7ca26c0b0d56317f545f079748e638 --- /dev/null +++ b/src/agentscope/formatter/_gemini_formatter.py @@ -0,0 +1,432 @@ +# -*- coding: utf-8 -*- +"""Google Gemini API formatter in agentscope.""" +import base64 +import fnmatch +import json +from abc import ABC +from typing import Any + +import requests +from pydantic import Field + +from ._formatter_base import FormatterBase +from .._logging import logger +from ..message import ( + Msg, + TextBlock, + ThinkingBlock, + HintBlock, + DataBlock, + ToolCallBlock, + ToolResultBlock, + URLSource, + Base64Source, +) + + +class _GeminiFormatterBase(FormatterBase, ABC): + """Base class for Gemini formatters, providing shared data block + formatting logic.""" + + def _format_gemini_data_block( + self, + block: DataBlock, + ) -> dict[str, Any] | None: + """Format a DataBlock into Gemini API format. + + Args: + block (`DataBlock`): + The data block to format. + + Returns: + `dict[str, Any] | None`: + The formatted data block in Gemini ``inline_data`` format, + or None if the media type is not supported. + """ + source = block.source + media_type = source.media_type + + # Check if media type is supported + if not any( + fnmatch.fnmatch(media_type, pattern) + for pattern in self.supported_input_media_types + ): + logger.warning( + "Media type %s is not supported, skipped.", + media_type, + ) + return None + + return self._format_media_source(source) + + @staticmethod + def _format_media_source( + source: URLSource | Base64Source, + ) -> dict[str, Any]: + """Format a media source into Gemini API ``inline_data`` format. + + Args: + source (`URLSource | Base64Source`): + The media source to format. + + Returns: + `dict[str, Any]`: + The formatted media source. + """ + if isinstance(source, Base64Source): + return { + "inline_data": { + "data": source.data, + "mime_type": source.media_type, + }, + } + elif isinstance(source, URLSource): + url = str(source.url) + if url.startswith("file://"): + # Local file - read and convert to base64 + file_path = url.removeprefix("file://") + with open(file_path, "rb") as f: + data = base64.b64encode(f.read()).decode("utf-8") + return { + "inline_data": { + "data": data, + "mime_type": source.media_type, + }, + } + else: + # Remote URL - download and convert to base64 + response = requests.get(url, timeout=30) + response.raise_for_status() + data = base64.b64encode(response.content).decode("utf-8") + return { + "inline_data": { + "data": data, + "mime_type": source.media_type, + }, + } + else: + raise ValueError(f"Unsupported source type: {type(source)}") + + +class GeminiChatFormatter(_GeminiFormatterBase): + """The Gemini formatter class for chatbot scenario, where only a user + and an agent are involved. We use the `role` field to identify different + entities in the conversation. + """ + + input_types: list[str] = Field( + default_factory=lambda: [ + "text/plain", + "image/*", + "audio/*", + "video/*", + ], + description=( + "The supported input types. " + 'Defaults to ``["text/plain", "image/*", "audio/*", "video/*"]``.' + ), + ) + + async def format( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Format message objects into Gemini API required format. + + Args: + msgs (`list[Msg]`): + The list of message objects to format. + + Returns: + `list[dict[str, Any]]`: + The formatted messages as a list of dictionaries. + """ + self.assert_list_of_msgs(msgs) + + messages: list[dict] = [] + i = 0 + while i < len(msgs): + msg = msgs[i] + parts: list = [] + + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + parts.append({"text": block.text}) + + elif isinstance(block, ThinkingBlock): + # Gemini API requires `thought: true` to mark a part as a + # thinking/reasoning block so the model can distinguish it + # from normal text and maintain reasoning continuity. + parts.append({"thought": True, "text": block.thinking}) + + elif isinstance(block, HintBlock): + if parts: + role = "model" if msg.role == "assistant" else "user" + messages.append({"role": role, "parts": parts}) + parts = [] + + if isinstance(block.hint, str): + messages.append( + { + "role": "user", + "parts": [{"text": block.hint}], + }, + ) + else: + hint_parts: list[dict] = [] + for sub in block.hint: + if isinstance(sub, TextBlock): + hint_parts.append({"text": sub.text}) + elif isinstance(sub, DataBlock): + formatted_sub = self._format_gemini_data_block( + sub, + ) + if formatted_sub: + hint_parts.append(formatted_sub) + if hint_parts: + messages.append( + {"role": "user", "parts": hint_parts}, + ) + + elif isinstance(block, DataBlock): + formatted = self._format_gemini_data_block(block) + if formatted: + parts.append(formatted) + + elif isinstance(block, ToolCallBlock): + parts.append( + { + "function_call": { + "id": block.id, + "name": block.name, + "args": json.loads(block.input or "{}"), + }, + }, + ) + + elif isinstance(block, ToolResultBlock): + if parts: + role = "model" if msg.role == "assistant" else "user" + messages.append({"role": role, "parts": parts}) + parts = [] + + ( + textual_output, + multimodal_data, + ) = self.convert_tool_result_to_string(block.output) + + messages.append( + { + "role": "user", + "parts": [ + { + "function_response": { + "id": block.id, + "name": block.name, + "response": { + "output": textual_output, + }, + }, + }, + ], + }, + ) + + if multimodal_data: + promo_parts = [] + for item in multimodal_data: + if isinstance(item, TextBlock): + promo_parts.append({"text": item.text}) + elif isinstance(item, DataBlock): + fmt_item = self._format_gemini_data_block( + item, + ) + if fmt_item is not None: + promo_parts.append(fmt_item) + if promo_parts: + messages.append( + {"role": "user", "parts": promo_parts}, + ) + + else: + logger.warning( + "Unsupported block type %s in the message, skipped.", + type(block), + ) + + # Gemini uses "model" instead of "assistant" + role = "model" if msg.role == "assistant" else "user" + + if parts: + messages.append( + { + "role": role, + "parts": parts, + }, + ) + + i += 1 + + return messages + + +class GeminiMultiAgentFormatter(_GeminiFormatterBase): + """The multi-agent formatter for Google Gemini API, where more than a + user and an agent are involved. + + .. note:: This formatter will combine previous messages (except tool + calls/results) into a history section in the first system message with + the conversation history prompt. + + .. note:: For tool calls/results, they will be presented as separate + messages as required by the Gemini API. Therefore, the tool calls/ + results messages are expected to be placed at the end of the input + messages. + + .. tip:: Telling the assistant's name in the system prompt is very + important in multi-agent conversations. So that LLM can know who it + is playing as. + + """ + + conversation_history_prompt: str = Field( + default=( + "# Conversation History\n" + "The content between tags contains " + "your conversation history\n" + ), + description="The prompt to use for the conversation history section.", + ) + + input_types: list[str] = Field( + default_factory=lambda: [ + "text/plain", + "image/*", + "audio/*", + "video/*", + ], + description=( + "The supported input types. " + 'Defaults to ``["text/plain", "image/*", "audio/*", "video/*"]``.' + ), + ) + + async def format(self, msgs: list[Msg]) -> list[dict[str, Any]]: + """Format input messages into the structure required by the Gemini + API for multi-agent conversations.""" + self.assert_list_of_msgs(msgs) + + formatted_msgs = [] + start_index = 0 + if len(msgs) > 0 and msgs[0].role == "system": + formatted_msgs.append( + await self._format_system_message(msgs[0]), + ) + start_index = 1 + + is_first_agent_message = True + async for typ, group in self._group_messages(msgs[start_index:]): + match typ: + case "tool_sequence": + formatted_msgs.extend( + await self._format_tool_sequence(group), + ) + case "agent_message": + formatted_msgs.extend( + await self._format_agent_message( + group, + is_first_agent_message, + ), + ) + is_first_agent_message = False + + return formatted_msgs + + async def _format_tool_sequence( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Given a sequence of tool call/result messages, format them into + the required format for the Gemini API.""" + return await GeminiChatFormatter( + input_types=self.input_types, + ).format(msgs) + + async def _format_agent_message( + self, + msgs: list[Msg], + is_first: bool = True, + ) -> list[dict[str, Any]]: + """Given a sequence of messages without tool calls/results, format + them into the required format for the Gemini API.""" + + if is_first: + conversation_history_prompt = self.conversation_history_prompt + else: + conversation_history_prompt = "" + + formatted_msgs: list[dict] = [] + conversation_parts: list[dict] = [] + accumulated_text: list[str] = [] + + for msg in msgs: + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + accumulated_text.append(f"{msg.name}: {block.text}") + + elif isinstance(block, DataBlock): + # Flush accumulated text first + if accumulated_text: + conversation_parts.append( + {"text": "\n".join(accumulated_text)}, + ) + accumulated_text = [] + + formatted = self._format_gemini_data_block(block) + if formatted: + conversation_parts.append(formatted) + + if accumulated_text: + conversation_parts.append( + {"text": "\n".join(accumulated_text)}, + ) + + # Add prompt and tags around conversation history + if conversation_parts: + if conversation_parts[0].get("text"): + conversation_parts[0]["text"] = ( + conversation_history_prompt + + "\n" + + conversation_parts[0]["text"] + ) + else: + conversation_parts.insert( + 0, + {"text": conversation_history_prompt + "\n"}, + ) + + if conversation_parts[-1].get("text"): + conversation_parts[-1]["text"] += "\n" + else: + conversation_parts.append({"text": ""}) + + formatted_msgs.append( + { + "role": "user", + "parts": conversation_parts, + }, + ) + + return formatted_msgs + + @staticmethod + async def _format_system_message(msg: Msg) -> dict[str, Any]: + """Format system message for the Gemini API.""" + return { + "role": "user", + "parts": [ + { + "text": msg.get_text_content(), + }, + ], + } diff --git a/src/agentscope/formatter/_moonshot_formatter.py b/src/agentscope/formatter/_moonshot_formatter.py new file mode 100644 index 0000000000000000000000000000000000000000..36d38e0d3562687375e8a97c18ce2463400b0e11 --- /dev/null +++ b/src/agentscope/formatter/_moonshot_formatter.py @@ -0,0 +1,417 @@ +# -*- coding: utf-8 -*- +"""The Moonshot AI formatter for agentscope.""" +import base64 +from typing import Any + +import requests +from pydantic import Field + +from ._openai_formatter import _OpenAIFormatterBase +from .._logging import logger +from ..message import ( + Msg, + URLSource, + Base64Source, + TextBlock, + DataBlock, + ThinkingBlock, + HintBlock, + ToolCallBlock, + ToolResultBlock, +) + + +def _moonshot_format_image_source( + source: URLSource | Base64Source, +) -> dict[str, Any]: + """Convert an image source to Moonshot ``image_url`` format. + + Moonshot's vision API only accepts base64 data URIs or file IDs — raw + remote URLs are rejected. This helper downloads remote ``http(s)://`` + URLs and converts them to base64 data URIs, while ``file://`` URLs and + ``Base64Source`` go through the same conversion as the OpenAI base. + """ + if isinstance(source, Base64Source): + url = f"data:{source.media_type};base64,{source.data}" + + elif isinstance(source, URLSource): + url_str = str(source.url) + if url_str.startswith("file://"): + local_path = url_str.removeprefix("file://") + with open(local_path, "rb") as f: + encoded = base64.b64encode(f.read()).decode("utf-8") + url = f"data:{source.media_type};base64,{encoded}" + else: + response = requests.get(url_str, timeout=30) + response.raise_for_status() + encoded = base64.b64encode(response.content).decode("utf-8") + url = f"data:{source.media_type};base64,{encoded}" + + else: + raise ValueError(f"Unsupported image source type: {type(source)}") + + return { + "type": "image_url", + "image_url": {"url": url}, + } + + +class MoonshotChatFormatter(_OpenAIFormatterBase): + """The Moonshot AI formatter for chatbot scenario. + + Moonshot's API is OpenAI-compatible, but thinking models (``kimi-k2.6``, + ``kimi-k2-thinking``) return a ``reasoning_content`` field alongside + ``content`` in assistant messages. This formatter preserves that field + when re-sending assistant messages back to the API so that the + *Preserved Thinking* feature works correctly in multi-turn conversations. + """ + + input_types: list[str] = Field( + default_factory=lambda: ["text/plain", "image/*", "audio/*"], + description=( + "The supported input types. " + 'Defaults to ``["text/plain", "image/*", "audio/*"]``.' + ), + ) + + def _format_image_source( + self, + source: URLSource | Base64Source, + ) -> dict[str, Any]: + return _moonshot_format_image_source(source) + + # pylint: disable=too-many-branches + async def format( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Format messages into the Moonshot / OpenAI-compatible API format. + + Behaves identically to :class:`OpenAIChatFormatter` except that + :class:`ThinkingBlock` content is placed into the ``reasoning_content`` + field of the assistant message dict (required for Preserved Thinking). + + Args: + msgs (`list[Msg]`): + The list of message objects to format. + + Returns: + `list[dict[str, Any]]`: + The formatted messages as a list of dictionaries. + """ + self.assert_list_of_msgs(msgs) + + messages: list[dict] = [] + i = 0 + while i < len(msgs): + msg = msgs[i] + content_blocks: list[dict] = [] + reasoning_parts: list[str] = [] + tool_calls: list[dict] = [] + + for block in msg.get_content_blocks(): + if isinstance(block, ThinkingBlock): + # Preserve reasoning_content for multi-turn + # Preserved Thinking (kimi-k2.6 / kimi-k2-thinking) + reasoning_parts.append(block.thinking) + + elif isinstance(block, TextBlock): + content_blocks.append({"type": "text", "text": block.text}) + + elif isinstance(block, DataBlock): + formatted = self._format_openai_data_block(block) + if formatted is not None: + content_blocks.append(formatted) + + elif isinstance(block, HintBlock): + if content_blocks or tool_calls or reasoning_parts: + msg_moonshot = { + "role": msg.role, + "name": msg.name, + "content": content_blocks or None, + } + if msg.role == "assistant": + msg_moonshot["reasoning_content"] = ( + "\n".join(reasoning_parts) + if reasoning_parts + else "" + ) + if tool_calls: + msg_moonshot["tool_calls"] = tool_calls + messages.append(msg_moonshot) + content_blocks = [] + reasoning_parts = [] + tool_calls = [] + + if isinstance(block.hint, str): + messages.append( + { + "role": "user", + "content": [ + {"type": "text", "text": block.hint}, + ], + }, + ) + else: + hint_parts: list[dict] = [] + for sub in block.hint: + if isinstance(sub, TextBlock): + hint_parts.append( + {"type": "text", "text": sub.text}, + ) + elif isinstance(sub, DataBlock): + formatted_sub = self._format_openai_data_block( + sub, + ) + if formatted_sub is not None: + hint_parts.append(formatted_sub) + if hint_parts: + messages.append( + {"role": "user", "content": hint_parts}, + ) + + elif isinstance(block, ToolCallBlock): + tool_calls.append( + { + "id": block.id, + "type": "function", + "function": { + "name": block.name, + "arguments": block.input, + }, + }, + ) + + elif isinstance(block, ToolResultBlock): + if content_blocks or tool_calls or reasoning_parts: + msg_flush = { + "role": msg.role, + "name": msg.name, + "content": content_blocks or None, + } + if msg.role == "assistant": + msg_flush["reasoning_content"] = ( + "\n".join(reasoning_parts) + if reasoning_parts + else "" + ) + if tool_calls: + msg_flush["tool_calls"] = tool_calls + messages.append(msg_flush) + content_blocks = [] + reasoning_parts = [] + tool_calls = [] + + ( + textual_output, + multimodal_data, + ) = self.convert_tool_result_to_string(block.output) + + messages.append( + { + "role": "tool", + "tool_call_id": block.id, + "content": textual_output, + "name": block.name, + }, + ) + + if multimodal_data: + promo_content = [] + for item in multimodal_data: + if isinstance(item, TextBlock): + promo_content.append( + {"type": "text", "text": item.text}, + ) + elif isinstance(item, DataBlock): + fmt_item = self._format_openai_data_block( + item, + ) + if fmt_item is not None: + promo_content.append(fmt_item) + if promo_content: + messages.append( + { + "role": "user", + "name": "system-reminder", + "content": promo_content, + }, + ) + + else: + logger.warning( + "Unsupported block type %s in the message, skipped.", + type(block), + ) + + msg_moonshot = { + "role": msg.role, + "name": msg.name, + "content": content_blocks or None, + } + + # Moonshot's Preserved Thinking requires `reasoning_content` on ALL + # assistant messages in multi-turn conversations (None when no + # thinking took place), so that the model can continue its chain + # of thought correctly. + if msg.role == "assistant": + msg_moonshot["reasoning_content"] = ( + "\n".join(reasoning_parts) if reasoning_parts else "" + ) + + if tool_calls: + msg_moonshot["tool_calls"] = tool_calls + + if ( + msg_moonshot["content"] + or msg_moonshot.get("tool_calls") + or reasoning_parts + ): + messages.append(msg_moonshot) + + i += 1 + + return messages + + +class MoonshotMultiAgentFormatter(_OpenAIFormatterBase): + """Formatter for the Moonshot AI API in multi-agent conversations. + + Moonshot's API is OpenAI-compatible, so the multi-agent history collapsing + strategy is the same as :class:`OpenAIMultiAgentFormatter`. Tool + sequences are delegated to :class:`MoonshotChatFormatter` so that + ``reasoning_content`` is preserved correctly for multi-turn + *Preserved Thinking* conversations. + + .. note:: Telling the assistant's name in the system prompt is important + in multi-agent conversations so that the model knows which role it + is playing. + """ + + conversation_history_prompt: str = Field( + default=( + "# Conversation History\n" + "The content between tags contains " + "your conversation history\n" + ), + description="The prompt to use for the conversation history section.", + ) + + input_types: list[str] = Field( + default_factory=lambda: ["text/plain", "image/*", "audio/*"], + description=( + "The supported input types. " + 'Defaults to ``["text/plain", "image/*", "audio/*"]``.' + ), + ) + + def _format_image_source( + self, + source: URLSource | Base64Source, + ) -> dict[str, Any]: + return _moonshot_format_image_source(source) + + async def format(self, msgs: list[Msg]) -> list[dict[str, Any]]: + """Format input messages into the Moonshot AI API format for + multi-agent conversations. + + Non-tool messages from all agents are collapsed into a single user + message with ```` tags. Tool call / result + sequences are delegated to :class:`MoonshotChatFormatter`. + + Args: + msgs (`list[Msg]`): + The list of message objects to format. + + Returns: + `list[dict[str, Any]]`: + The formatted messages as a list of dictionaries. + """ + self.assert_list_of_msgs(msgs) + + formatted_msgs: list[dict] = [] + start_index = 0 + if msgs and msgs[0].role == "system": + formatted_msgs.append( + await self._format_system_message(msgs[0]), + ) + start_index = 1 + + is_first_agent_message = True + async for typ, group in self._group_messages(msgs[start_index:]): + if typ == "tool_sequence": + formatted_msgs.extend( + await MoonshotChatFormatter( + input_types=self.input_types, + ).format(group), + ) + elif typ == "agent_message": + formatted_msgs.extend( + await self._format_agent_message( + group, + is_first_agent_message, + ), + ) + is_first_agent_message = False + + return formatted_msgs + + async def _format_tool_sequence( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Format a sequence of tool-related messages using + MoonshotChatFormatter.""" + return await MoonshotChatFormatter( + input_types=self.input_types, + ).format(msgs) + + async def _format_agent_message( + self, + msgs: list[Msg], + is_first: bool = True, + ) -> list[dict[str, Any]]: + """Collapse agent messages into a ```` user message.""" + if is_first: + conversation_history_prompt = self.conversation_history_prompt + else: + conversation_history_prompt = "" + + accumulated_text: list[str] = [] + media_blocks: list[dict] = [] + + for msg in msgs: + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + accumulated_text.append(f"{msg.name}: {block.text}") + elif isinstance(block, DataBlock): + formatted = self._format_openai_data_block(block) + if formatted is not None: + media_blocks.append(formatted) + + if not accumulated_text and not media_blocks: + return [] + + history_text = "\n".join(accumulated_text) + if history_text: + history_text = ( + conversation_history_prompt + + "\n" + + history_text + + "\n" + ) + + content_list: list[dict[str, Any]] = [] + if history_text: + content_list.append({"type": "text", "text": history_text}) + content_list.extend(media_blocks) + + return [{"role": "user", "content": content_list}] + + @staticmethod + async def _format_system_message(msg: Msg) -> dict[str, Any]: + """Format a system message for the Moonshot AI API.""" + return { + "role": "system", + "content": msg.get_text_content(), + } diff --git a/src/agentscope/formatter/_ollama_formatter.py b/src/agentscope/formatter/_ollama_formatter.py new file mode 100644 index 0000000000000000000000000000000000000000..d70b708eab0f71a96e27aaf82e7f1b9fa6f95886 --- /dev/null +++ b/src/agentscope/formatter/_ollama_formatter.py @@ -0,0 +1,442 @@ +# -*- coding: utf-8 -*- +"""The Ollama formatter module.""" +import base64 +import fnmatch +import json +from abc import ABC +from typing import Any + +import requests +from pydantic import Field + +from ._formatter_base import FormatterBase +from .._logging import logger +from ..message import ( + Msg, + TextBlock, + HintBlock, + DataBlock, + ToolCallBlock, + ToolResultBlock, + ThinkingBlock, + URLSource, + Base64Source, +) + + +class _OllamaFormatterBase(FormatterBase, ABC): + """Base class for Ollama formatters, providing shared data block + formatting logic.""" + + def _format_ollama_data_block( + self, + block: DataBlock, + ) -> str | None: + """Format a DataBlock into Ollama API format (base64 string). + + Args: + block (`DataBlock`): + The data block to format. + + Returns: + `str | None`: + Base64 encoded data as a string, or None if the media type + is not supported. + """ + source = block.source + media_type = source.media_type + + # Check if media type is supported + if not any( + fnmatch.fnmatch(media_type, pattern) + for pattern in self.supported_input_media_types + ): + logger.warning( + "Media type %s is not supported, skipped.", + media_type, + ) + return None + + # Ollama only supports images + if not media_type.startswith("image/"): + logger.warning( + "Ollama only supports image data, got %s, skipped.", + media_type, + ) + return None + + return self._format_image_source(source) + + @staticmethod + def _format_image_source(source: URLSource | Base64Source) -> str: + """Format an image source into Ollama API format (base64 string). + + Args: + source (`URLSource | Base64Source`): + The image source to format. + + Returns: + `str`: + Base64 encoded image data. + """ + if isinstance(source, Base64Source): + return source.data + elif isinstance(source, URLSource): + url = str(source.url) + if url.startswith("file://"): + # Local file - read and convert to base64 + file_path = url.removeprefix("file://") + with open(file_path, "rb") as f: + data = base64.b64encode(f.read()).decode("utf-8") + return data + else: + # Remote URL - download and convert to base64 + response = requests.get(url, timeout=30) + response.raise_for_status() + data = base64.b64encode(response.content).decode("utf-8") + return data + else: + raise ValueError(f"Unsupported source type: {type(source)}") + + +class OllamaChatFormatter(_OllamaFormatterBase): + """The Ollama formatter class for chatbot scenario, where only a user + and an agent are involved. We use the `role` field to identify different + participants in the conversation. + """ + + input_types: list[str] = Field( + default_factory=lambda: ["text/plain", "image/*"], + description=( + "The supported input types. " + 'Defaults to ``["text/plain", "image/*"]``.' + ), + ) + + async def format( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Format message objects into Ollama API format. + + Args: + msgs (`list[Msg]`): + The list of message objects to format. + + Returns: + `list[dict[str, Any]]`: + The formatted messages as a list of dictionaries. + """ + self.assert_list_of_msgs(msgs) + + messages: list[dict] = [] + for msg in msgs: + content_parts = [] + images = [] + + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + content_parts.append(block.text) + + elif isinstance(block, HintBlock): + if content_parts or images: + msg_flush = { + "role": msg.role, + "content": "\n".join(content_parts), + } + if images: + msg_flush["images"] = images + messages.append(msg_flush) + content_parts = [] + images = [] + + if isinstance(block.hint, str): + messages.append( + {"role": "user", "content": block.hint}, + ) + else: + hint_text_parts: list[str] = [] + hint_images: list[str] = [] + for sub in block.hint: + if isinstance(sub, TextBlock): + hint_text_parts.append(sub.text) + elif isinstance(sub, DataBlock): + formatted_sub = self._format_ollama_data_block( + sub, + ) + if formatted_sub: + hint_images.append(formatted_sub) + if hint_text_parts or hint_images: + hint_msg: dict[str, Any] = { + "role": "user", + "content": "\n".join(hint_text_parts), + } + if hint_images: + hint_msg["images"] = hint_images + messages.append(hint_msg) + + elif isinstance(block, DataBlock): + formatted_image = self._format_ollama_data_block(block) + if formatted_image: + images.append(formatted_image) + + elif isinstance(block, ThinkingBlock): + # Ollama does not use reasoning content in the context + # — skip thinking blocks silently. + pass + + elif isinstance(block, ToolCallBlock): + messages.append( + { + "role": msg.role, + "content": "\n".join(content_parts) + if content_parts + else "", + "tool_calls": [ + { + "function": { + "name": block.name, + # Ollama SDK expects a dict, not a + # JSON string. + "arguments": json.loads( + block.input or "{}", + ), + }, + }, + ], + }, + ) + content_parts = [] + images = [] + + elif isinstance(block, ToolResultBlock): + if content_parts or images: + msg_flush = { + "role": msg.role, + "content": "\n".join(content_parts), + } + if images: + msg_flush["images"] = images + messages.append(msg_flush) + content_parts = [] + images = [] + + ( + textual_output, + multimodal_data, + ) = self.convert_tool_result_to_string(block.output) + + # Ollama expects tool results as a separate "tool" role + # message, regardless of the containing Msg's role. + messages.append( + { + "role": "tool", + "content": textual_output, + }, + ) + + # If there's multimodal data, append an extra user message. + if multimodal_data: + user_images = [] + user_content_parts = [] + for data_block in multimodal_data: + if isinstance(data_block, DataBlock): + formatted_image = ( + self._format_ollama_data_block( + data_block, + ) + ) + if formatted_image: + user_images.append(formatted_image) + elif isinstance(data_block, TextBlock): + user_content_parts.append(data_block.text) + + user_msg = { + "role": "user", + "content": "\n".join(user_content_parts) + if user_content_parts + else textual_output, + } + if user_images: + user_msg["images"] = user_images + messages.append(user_msg) + + else: + logger.warning( + "Unsupported block type %s in the message, skipped.", + type(block), + ) + + # Add the message if there's content or images + if content_parts or images: + msg_ollama: dict[str, Any] = { + "role": msg.role, + "content": "\n".join(content_parts) + if content_parts + else "", + } + if images: + msg_ollama["images"] = images + messages.append(msg_ollama) + + return messages + + +class OllamaMultiAgentFormatter(_OllamaFormatterBase): + """ + Ollama formatter for multi-agent conversations, where more than + a user and an agent are involved. + """ + + conversation_history_prompt: str = Field( + default=( + "# Conversation History\n" + "The content between tags contains " + "your conversation history\n" + ), + description="The prompt to use for the conversation history section.", + ) + + input_types: list[str] = Field( + default_factory=lambda: ["text/plain", "image/*"], + description=( + "The supported input types. " + 'Defaults to ``["text/plain", "image/*"]``.' + ), + ) + + async def format(self, msgs: list[Msg]) -> list[dict[str, Any]]: + """Format input messages into the structure required by the Ollama + API for multi-agent conversations.""" + self.assert_list_of_msgs(msgs) + + formatted_msgs = [] + start_index = 0 + if len(msgs) > 0 and msgs[0].role == "system": + formatted_msgs.append( + await self._format_system_message(msgs[0]), + ) + start_index = 1 + + is_first_agent_message = True + async for typ, group in self._group_messages(msgs[start_index:]): + match typ: + case "tool_sequence": + formatted_msgs.extend( + await self._format_tool_sequence(group), + ) + case "agent_message": + formatted_msgs.extend( + await self._format_agent_message( + group, + is_first_agent_message, + ), + ) + is_first_agent_message = False + + return formatted_msgs + + async def _format_tool_sequence( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Format a sequence of tool-related messages.""" + return await OllamaChatFormatter( + input_types=self.input_types, + ).format(msgs) + + async def _format_agent_message( + self, + msgs: list[Msg], + is_first: bool, + ) -> list[dict[str, Any]]: + """Format agent messages into conversation history format.""" + conversation_blocks: list[dict] = [] + accumulated_text: list[str] = [] + images: list[str] = [] + + for msg in msgs: + msg_text_parts = [] + if msg.name: + msg_text_parts.append(f"{msg.name}:") + + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + msg_text_parts.append(block.text) + elif isinstance(block, DataBlock): + formatted_image = self._format_ollama_data_block(block) + if formatted_image: + images.append(formatted_image) + elif isinstance(block, (HintBlock, ThinkingBlock)): + pass # Ollama does not use hint/thinking blocks + else: + logger.warning( + "Unsupported block type %s in agent message, skipped.", + type(block), + ) + + if msg_text_parts: + accumulated_text.append("\n".join(msg_text_parts)) + + if accumulated_text: + conversation_blocks.append( + {"text": "\n".join(accumulated_text)}, + ) + + if conversation_blocks and is_first: + if conversation_blocks[0].get("text"): + conversation_blocks[0]["text"] = ( + self.conversation_history_prompt + + "\n" + + conversation_blocks[0]["text"] + ) + + else: + conversation_blocks.insert( + 0, + { + "text": self.conversation_history_prompt + + "\n", + }, + ) + + if conversation_blocks[-1].get("text"): + conversation_blocks[-1]["text"] += "\n" + + else: + conversation_blocks.append({"text": ""}) + + conversation_blocks_text = "\n".join( + conversation_block.get("text", "") + for conversation_block in conversation_blocks + ) + + user_message: dict[str, Any] = { + "role": "user", + "content": conversation_blocks_text, + } + if images: + user_message["images"] = images + + formatted_msgs = [] + if conversation_blocks: + formatted_msgs.append(user_message) + + return formatted_msgs + + @staticmethod + async def _format_system_message(msg: Msg) -> dict[str, Any]: + """Format a system message.""" + text_parts = [] + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + text_parts.append(block.text) + else: + logger.warning( + "Unsupported block type %s in system message, skipped.", + type(block), + ) + return { + "role": "system", + "content": "\n".join(text_parts), + } diff --git a/src/agentscope/formatter/_openai_formatter.py b/src/agentscope/formatter/_openai_formatter.py new file mode 100644 index 0000000000000000000000000000000000000000..3ef4cb9459a0f9f03ead8cbf67453b3d3617edad --- /dev/null +++ b/src/agentscope/formatter/_openai_formatter.py @@ -0,0 +1,513 @@ +# -*- coding: utf-8 -*- +"""The OpenAI formatter for agentscope.""" +import base64 +from abc import ABC +from fnmatch import fnmatch +from typing import Any +from urllib.parse import urlparse + +import requests +from pydantic import Field + +from ._formatter_base import FormatterBase +from .._logging import logger +from ..message import ( + Msg, + URLSource, + TextBlock, + DataBlock, + Base64Source, + ToolCallBlock, + ToolResultBlock, + HintBlock, + ThinkingBlock, +) + + +class _OpenAIFormatterBase(FormatterBase, ABC): + """Base class for OpenAI formatters, providing shared data block + formatting logic.""" + + def _format_openai_data_block( + self, + block: DataBlock, + ) -> dict[str, Any] | None: + """Format a DataBlock into the required format for OpenAI API. + + For image blocks, URLs are returned as-is (or converted to base64 for + local ``file://`` paths). For audio blocks, data is always converted + to base64 as required by the OpenAI input_audio format. + + Args: + block (`DataBlock`): + The DataBlock to format. + + Returns: + `dict[str, Any] | None`: + A dictionary in OpenAI API format, or ``None`` if the block + should be skipped. + """ + if not any( + fnmatch(block.source.media_type, pattern) + for pattern in self.supported_input_media_types + ): + logger.warning( + "Unsupported media type %s for OpenAI API. " + "Supported types: %s. This block will be skipped.", + block.source.media_type, + ", ".join(self.supported_input_media_types), + ) + return None + + main_type = block.source.media_type.split("/")[0] + + if main_type == "image": + return self._format_image_source(block.source) + + if main_type == "audio": + return self._format_audio_source(block.source) + + logger.warning( + "Unsupported main media type %s for OpenAI API. " + "This block will be skipped.", + main_type, + ) + return None + + def _format_image_source( + self, + source: URLSource | Base64Source, + ) -> dict[str, Any]: + """Convert an image source to OpenAI image_url format. + + Local ``file://`` URLs are read from disk and converted to base64 + data URIs. Remote URLs are passed through unchanged. Subclasses may + override this to apply provider-specific handling (e.g. forcing + remote URLs to be downloaded and base64-encoded for APIs that don't + accept raw HTTPS URLs). + + Args: + source (`URLSource | Base64Source`): + The image source to convert. + + Returns: + `dict[str, Any]`: + A dictionary with ``"type": "image_url"`` in OpenAI format. + """ + if isinstance(source, Base64Source): + url = f"data:{source.media_type};base64,{source.data}" + + elif isinstance(source, URLSource): + url_str = str(source.url) + if url_str.startswith("file://"): + # Local file — read and encode as base64 data URI + local_path = url_str.removeprefix("file://") + with open(local_path, "rb") as f: + encoded = base64.b64encode(f.read()).decode("utf-8") + url = f"data:{source.media_type};base64,{encoded}" + else: + # Remote URL — pass through as-is + url = url_str + + else: + raise ValueError(f"Unsupported image source type: {type(source)}") + + return { + "type": "image_url", + "image_url": {"url": url}, + } + + @staticmethod + def _format_audio_source( + source: URLSource | Base64Source, + ) -> dict[str, Any]: + """Convert an audio source to OpenAI input_audio format. + + Local ``file://`` URLs are read from disk. Remote URLs are downloaded. + Only ``wav`` and ``mp3`` formats are supported by the OpenAI API. + + Args: + source (`URLSource | Base64Source`): + The audio source to convert. + + Returns: + `dict[str, Any]`: + A dictionary with ``"type": "input_audio"`` in OpenAI format. + """ + if isinstance(source, Base64Source): + media_type = source.media_type + if media_type not in ["audio/wav", "audio/mp3"]: + raise TypeError( + f"Unsupported audio media type: {media_type}, " + "only audio/wav and audio/mp3 are supported.", + ) + return { + "type": "input_audio", + "input_audio": { + "data": source.data, + "format": media_type.split("/")[-1], + }, + } + + if isinstance(source, URLSource): + url_str = str(source.url) + if url_str.startswith("file://"): + # Local file + local_path = url_str.removeprefix("file://") + extension = local_path.rsplit(".", 1)[-1].lower() + if extension not in ["wav", "mp3"]: + raise TypeError( + f"Unsupported audio file extension: {extension}, " + "wav and mp3 are supported.", + ) + with open(local_path, "rb") as f: + data = base64.b64encode(f.read()).decode("utf-8") + else: + # Remote URL — download and encode + parsed = urlparse(url_str) + extension = parsed.path.rsplit(".", 1)[-1].lower() + if extension not in ["wav", "mp3"]: + raise TypeError( + f"Unsupported audio file extension: {extension}, " + "wav and mp3 are supported.", + ) + response = requests.get(url_str, timeout=30) + response.raise_for_status() + data = base64.b64encode(response.content).decode("utf-8") + + return { + "type": "input_audio", + "input_audio": { + "data": data, + "format": extension, + }, + } + + raise TypeError(f"Unsupported audio source type: {type(source)}.") + + +class OpenAIChatFormatter(_OpenAIFormatterBase): + """The OpenAI formatter class for chatbot scenario, where only a user + and an agent are involved. We use the `name` field in OpenAI API to + identify different entities in the conversation. + """ + + input_types: list[str] = Field( + default_factory=lambda: ["text/plain", "image/*", "audio/*"], + description=( + "The supported input types. " + 'Defaults to ``["text/plain", "image/*", "audio/*"]``.' + ), + ) + + async def format( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Format message objects into OpenAI API required format. + + Args: + msgs (`list[Msg]`): + The list of Msg objects to format. + + Returns: + `list[dict[str, Any]]`: + A list of dictionaries, where each dictionary has "name", + "role", and "content" keys. + """ + self.assert_list_of_msgs(msgs) + + messages: list[dict] = [] + i = 0 + while i < len(msgs): + msg = msgs[i] + content_blocks = [] + tool_calls = [] + + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + content_blocks.append({"type": "text", "text": block.text}) + + elif isinstance(block, DataBlock): + formatted = self._format_openai_data_block( + block, + ) + if formatted is not None: + content_blocks.append(formatted) + + elif isinstance(block, HintBlock): + if content_blocks or tool_calls: + msg_openai = { + "role": msg.role, + "name": msg.name, + "content": content_blocks or None, + } + if tool_calls: + msg_openai["tool_calls"] = tool_calls + messages.append(msg_openai) + content_blocks = [] + tool_calls = [] + + if isinstance(block.hint, str): + messages.append( + { + "role": "user", + "content": [ + {"type": "text", "text": block.hint}, + ], + }, + ) + else: + hint_parts: list[dict] = [] + for sub in block.hint: + if isinstance(sub, TextBlock): + hint_parts.append( + {"type": "text", "text": sub.text}, + ) + elif isinstance(sub, DataBlock): + formatted_sub = self._format_openai_data_block( + sub, + ) + if formatted_sub is not None: + hint_parts.append(formatted_sub) + if hint_parts: + messages.append( + {"role": "user", "content": hint_parts}, + ) + + elif isinstance(block, ToolCallBlock): + tool_calls.append( + { + "id": block.id, + "type": "function", + "function": { + "name": block.name, + "arguments": block.input, + }, + }, + ) + + elif isinstance(block, ToolResultBlock): + if content_blocks or tool_calls: + msg_openai_flush = { + "role": msg.role, + "name": msg.name, + "content": content_blocks or None, + } + if tool_calls: + msg_openai_flush["tool_calls"] = tool_calls + messages.append(msg_openai_flush) + content_blocks = [] + tool_calls = [] + + ( + textual_output, + multimodal_data, + ) = self.convert_tool_result_to_string(block.output) + + messages.append( + { + "role": "tool", + "tool_call_id": block.id, + "content": textual_output, + "name": block.name, + }, + ) + + if multimodal_data: + promo_content = [] + for item in multimodal_data: + if isinstance(item, TextBlock): + promo_content.append( + {"type": "text", "text": item.text}, + ) + elif isinstance(item, DataBlock): + fmt_item = self._format_openai_data_block( + item, + ) + if fmt_item is not None: + promo_content.append(fmt_item) + if promo_content: + messages.append( + { + "role": "user", + "name": "system-reminder", + "content": promo_content, + }, + ) + + elif isinstance(block, ThinkingBlock): + # OpenAI API does not accept reasoning/thinking content + # in conversation history — skip thinking blocks silently. + pass + + else: + logger.warning( + "Unsupported block type %s in the message, skipped.", + type(block), + ) + + msg_openai = { + "role": msg.role, + "name": msg.name, + "content": content_blocks or None, + } + + if tool_calls: + msg_openai["tool_calls"] = tool_calls + + # When both content and tool_calls are None, skipped + if msg_openai["content"] or msg_openai.get("tool_calls"): + messages.append(msg_openai) + + # Move to next message + i += 1 + + return messages + + +class OpenAIMultiAgentFormatter(_OpenAIFormatterBase): + """ + OpenAI formatter for multi-agent conversations, where more than + a user and an agent are involved. + + .. tip:: This formatter is compatible with OpenAI API and + OpenAI-compatible services like vLLM, Azure OpenAI, and others. + """ + + conversation_history_prompt: str = Field( + default=( + "# Conversation History\n" + "The content between tags contains " + "your conversation history\n" + ), + description="The prompt to use for the conversation history section.", + ) + + input_types: list[str] = Field( + default_factory=lambda: ["text/plain", "image/*", "audio/*"], + description=( + "The supported input types. " + 'Defaults to ``["text/plain", "image/*", "audio/*"]``.' + ), + ) + + async def format(self, msgs: list[Msg]) -> list[dict[str, Any]]: + """Format input messages into the structure required by the OpenAI API + for multi-agent conversations.""" + self.assert_list_of_msgs(msgs) + + formatted_msgs = [] + start_index = 0 + if len(msgs) > 0 and msgs[0].role == "system": + formatted_msgs.append( + await self._format_system_message(msgs[0]), + ) + start_index = 1 + + is_first_agent_message = True + async for typ, group in self._group_messages(msgs[start_index:]): + match typ: + case "tool_sequence": + formatted_msgs.extend( + await self._format_tool_sequence(group), + ) + case "agent_message": + formatted_msgs.extend( + await self._format_agent_message( + group, + is_first_agent_message, + ), + ) + is_first_agent_message = False + + return formatted_msgs + + async def _format_tool_sequence( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Given a sequence of tool call/result messages, format them into + the required format for the OpenAI API.""" + return await OpenAIChatFormatter( + input_types=self.input_types, + ).format(msgs) + + async def _format_agent_message( + self, + msgs: list[Msg], + is_first: bool = True, + ) -> list[dict[str, Any]]: + """Given a sequence of messages without tool calls/results, format + them into the required format for the OpenAI API.""" + + if is_first: + conversation_history_prompt = self.conversation_history_prompt + else: + conversation_history_prompt = "" + + formatted_msgs: list[dict] = [] + conversation_blocks: list = [] + accumulated_text = [] + media_blocks: list[dict] = [] + + for msg in msgs: + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + accumulated_text.append(f"{msg.name}: {block.text}") + + elif isinstance(block, DataBlock): + formatted = self._format_openai_data_block( + block, + ) + if formatted is not None: + media_blocks.append(formatted) + + if accumulated_text: + conversation_blocks.append( + {"text": "\n".join(accumulated_text)}, + ) + + if conversation_blocks: + if conversation_blocks[0].get("text"): + conversation_blocks[0]["text"] = ( + conversation_history_prompt + + "\n" + + conversation_blocks[0]["text"] + ) + else: + conversation_blocks.insert( + 0, + {"text": conversation_history_prompt + "\n"}, + ) + + if conversation_blocks[-1].get("text"): + conversation_blocks[-1]["text"] += "\n" + else: + conversation_blocks.append({"text": ""}) + + conversation_blocks_text = "\n".join( + b.get("text", "") for b in conversation_blocks + ) + + content_list: list[dict[str, Any]] = [] + if conversation_blocks_text: + content_list.append( + {"type": "text", "text": conversation_blocks_text}, + ) + content_list.extend(media_blocks) + + if content_list: + formatted_msgs.append({"role": "user", "content": content_list}) + + return formatted_msgs + + @staticmethod + async def _format_system_message( + msg: Msg, + ) -> dict[str, Any]: + """Format system message for OpenAI API.""" + return { + "role": "system", + "content": msg.get_text_content(), + } diff --git a/src/agentscope/formatter/_openai_response_formatter.py b/src/agentscope/formatter/_openai_response_formatter.py new file mode 100644 index 0000000000000000000000000000000000000000..502836d66ebb3c0536baa8574799bdcecaa13cbc --- /dev/null +++ b/src/agentscope/formatter/_openai_response_formatter.py @@ -0,0 +1,486 @@ +# -*- coding: utf-8 -*- +"""Formatters for the OpenAI Responses API.""" +from abc import ABC +from typing import Any + +from pydantic import Field + +from ._openai_formatter import _OpenAIFormatterBase +from .._logging import logger +from ..message import ( + Msg, + TextBlock, + DataBlock, + ToolCallBlock, + ToolResultBlock, + HintBlock, + ThinkingBlock, +) + + +class _OpenAIResponseFormatterBase(_OpenAIFormatterBase, ABC): + """Base class for OpenAI Responses API formatters. + + Provides the shared ``_format_response_data_block`` helper used by both + :class:`OpenAIResponseFormatter` (chat) and + :class:`OpenAIResponseMultiAgentFormatter` (multi-agent). + """ + + input_types: list[str] = Field( + default_factory=lambda: ["text/plain", "image/*"], + description=( + "The supported input types. " + 'Defaults to ``["text/plain", "image/*"]``. ' + "Audio is not supported by the Responses API." + ), + ) + + def _format_response_data_block( + self, + block: DataBlock, + ) -> dict[str, Any] | None: + """Format a DataBlock into the Response API format. + + The Responses API uses different content types from the Chat + Completions API: + + * ``image_url`` → ``input_image`` + * ``input_audio`` → skipped (the Responses API does not support + audio input yet; use Chat Completions API instead). See + https://developers.openai.com/api/docs/guides/audio + + Args: + block (`DataBlock`): + The DataBlock to format. + + Returns: + `dict[str, Any] | None`: + A dictionary in the Responses API format, or ``None`` when the + block type is unsupported. + """ + # Intercept audio blocks before the generic formatter rejects them + # with a less helpful "Unsupported media type" warning. The Responses + # API does not support audio input yet; use Chat Completions API with + # an audio-capable model instead. + # https://developers.openai.com/api/docs/guides/audio + media_type = getattr(block.source, "media_type", "") or "" + if media_type.split("/", 1)[0] == "audio": + logger.warning( + "Audio input is not supported by the OpenAI Responses API. " + "Use OpenAIChatModel with an audio-capable model instead. " + "This audio block will be skipped.", + ) + return None + + base_result = self._format_openai_data_block(block) + if base_result is None: + return None + + if base_result.get("type") == "image_url": + return { + "type": "input_image", + "image_url": base_result["image_url"]["url"], + } + + return base_result + + +class OpenAIResponseFormatter(_OpenAIResponseFormatterBase): + """Formatter for the OpenAI Responses API in chat (single-agent) mode. + + Produces input items compatible with ``client.responses.create( + input=...)``. + Compared with the Chat Completions format, the key differences are: + + * Text content blocks use ``input_text`` instead of ``text``. + * Image content blocks use ``input_image`` instead of ``image_url``. + * Assistant tool-call messages become top-level ``function_call`` items. + * Tool result messages become ``function_call_output`` items. + * Reasoning items (``ThinkingBlock`` with ``reasoning_item_id``) are + echoed back verbatim as required by reasoning models (e.g. ``o1``). + """ + + # pylint: disable=too-many-branches + async def format( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Format message objects into OpenAI Response API input items. + + Args: + msgs (`list[Msg]`): + The list of Msg objects to format. + + Returns: + `list[dict[str, Any]]`: + A list of input items for ``client.responses.create``. + """ + self.assert_list_of_msgs(msgs) + + items: list[dict] = [] + i = 0 + while i < len(msgs): + msg = msgs[i] + content_parts: list[dict] = [] + function_calls: list[dict] = [] + + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + content_parts.append( + {"type": "input_text", "text": block.text}, + ) + + elif isinstance(block, DataBlock): + formatted = self._format_response_data_block(block) + if formatted is not None: + content_parts.append(formatted) + + elif isinstance(block, HintBlock): + if function_calls: + if content_parts: + items.append( + { + "role": msg.role, + "content": content_parts, + }, + ) + content_parts = [] + items.extend(function_calls) + function_calls = [] + elif content_parts: + items.append( + { + "role": msg.role, + "content": content_parts, + }, + ) + content_parts = [] + + if isinstance(block.hint, str): + items.append( + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": block.hint, + }, + ], + }, + ) + else: + hint_parts: list[dict] = [] + for sub in block.hint: + if isinstance(sub, TextBlock): + hint_parts.append( + { + "type": "input_text", + "text": sub.text, + }, + ) + elif isinstance(sub, DataBlock): + formatted_sub = ( + self._format_response_data_block( + sub, + ) + ) + if formatted_sub is not None: + hint_parts.append(formatted_sub) + if hint_parts: + items.append( + {"role": "user", "content": hint_parts}, + ) + + elif isinstance(block, ThinkingBlock): + # When reasoning_item_id is present the block originated + # from a Responses API "reasoning" output item. The API + # requires that such items are echoed back verbatim in + # multi-turn history (especially when they precede a + # function_call). Without the ID we skip silently. + reasoning_item_id = getattr( + block, + "reasoning_item_id", + None, + ) + if reasoning_item_id: + if content_parts: + items.append( + { + "role": msg.role, + "content": content_parts, + }, + ) + content_parts = [] + # summary may be empty when the model did not produce + # reasoning summary text (e.g. o4-mini with streaming) + summary = ( + [{"type": "summary_text", "text": block.thinking}] + if block.thinking + else [] + ) + items.append( + { + "type": "reasoning", + "id": reasoning_item_id, + "summary": summary, + "content": [], + }, + ) + + elif isinstance(block, ToolCallBlock): + # The Responses API distinguishes two identifiers on a + # function_call item: + # id → fc_xxx: the item identifier used when + # echoing the item in multi-turn history + # call_id → call_xxx: the identifier that must be + # echoed in the matching function_call_output + # For other APIs (Chat Completions, DashScope …) only one + # ID exists; call_id extra field is None and we fall back + # to id for both fields. + function_calls.append( + { + "type": "function_call", + "id": block.id, + "call_id": getattr(block, "call_id", None) + or block.id, + "name": block.name, + "arguments": block.input, + }, + ) + + elif isinstance(block, ToolResultBlock): + if function_calls: + if content_parts: + items.append( + { + "role": msg.role, + "content": content_parts, + }, + ) + content_parts = [] + items.extend(function_calls) + function_calls = [] + elif content_parts: + items.append( + { + "role": msg.role, + "content": content_parts, + }, + ) + content_parts = [] + + ( + textual_output, + multimodal_data, + ) = self.convert_tool_result_to_string(block.output) + + items.append( + { + "type": "function_call_output", + "call_id": block.id, + "output": textual_output, + }, + ) + + if multimodal_data: + promo_content = [] + for item in multimodal_data: + if isinstance(item, TextBlock): + promo_content.append( + { + "type": "input_text", + "text": item.text, + }, + ) + elif isinstance(item, DataBlock): + fmt_item = self._format_response_data_block( + item, + ) + if fmt_item is not None: + promo_content.append(fmt_item) + if promo_content: + items.append( + { + "role": "user", + "content": promo_content, + }, + ) + + else: + logger.warning( + "Unsupported block type %s in the message, skipped.", + type(block), + ) + + if function_calls: + if content_parts: + items.append( + { + "role": msg.role, + "content": content_parts, + }, + ) + items.extend(function_calls) + elif content_parts: + items.append( + { + "role": msg.role, + "content": content_parts, + }, + ) + + i += 1 + + return items + + +class OpenAIResponseMultiAgentFormatter(_OpenAIResponseFormatterBase): + """Formatter for the OpenAI Responses API in multi-agent mode. + + Handles conversations where more than a user and a single agent are + involved. Tool call/result sequences are formatted with the Responses API + ``function_call`` / ``function_call_output`` items (delegated to + :class:`OpenAIResponseFormatter`). Agent conversation history messages + are wrapped inside ```` tags and presented as a single + ``user`` input item. + """ + + conversation_history_prompt: str = Field( + default=( + "# Conversation History\n" + "The content between tags contains " + "your conversation history\n" + ), + description="The prompt to use for the conversation history section.", + ) + + async def format( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Format messages for multi-agent Responses API conversations. + + Args: + msgs (`list[Msg]`): + The list of Msg objects to format. + + Returns: + `list[dict[str, Any]]`: + A list of input items for ``client.responses.create``. + """ + self.assert_list_of_msgs(msgs) + + formatted_msgs: list[dict] = [] + start_index = 0 + if len(msgs) > 0 and msgs[0].role == "system": + formatted_msgs.append( + await self._format_system_message(msgs[0]), + ) + start_index = 1 + + is_first_agent_message = True + async for typ, group in self._group_messages(msgs[start_index:]): + if typ == "tool_sequence": + formatted_msgs.extend( + await self._format_tool_sequence(group), + ) + elif typ == "agent_message": + formatted_msgs.extend( + await self._format_agent_message( + group, + is_first_agent_message, + ), + ) + is_first_agent_message = False + + return formatted_msgs + + async def _format_tool_sequence( + self, + msgs: list[Msg], + ) -> list[dict[str, Any]]: + """Format a sequence of tool call/result messages using the Responses + API format. + + Args: + msgs (`list[Msg]`): + The tool call/result messages to format. + + Returns: + `list[dict[str, Any]]`: + A list of Responses API input items. + """ + return await OpenAIResponseFormatter( + input_types=self.input_types, + ).format(msgs) + + async def _format_agent_message( + self, + msgs: list[Msg], + is_first: bool = True, + ) -> list[dict[str, Any]]: + """Format a sequence of agent messages as a history-wrapped user item. + + Args: + msgs (`list[Msg]`): + The agent messages to format. + is_first (`bool`, defaults to ``True``): + Whether this is the first agent message group, which triggers + the conversation history prompt. + + Returns: + `list[dict[str, Any]]`: + A list containing at most one user input item. + """ + conversation_history_prompt = ( + self.conversation_history_prompt if is_first else "" + ) + + accumulated_text: list[str] = [] + media_blocks: list[dict] = [] + + for msg in msgs: + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + accumulated_text.append(f"{msg.name}: {block.text}") + elif isinstance(block, DataBlock): + formatted = self._format_response_data_block(block) + if formatted is not None: + media_blocks.append(formatted) + + if not accumulated_text and not media_blocks: + return [] + + history_text = "\n".join(accumulated_text) + wrapped = ( + conversation_history_prompt + + "\n" + + history_text + + "\n" + ) + + content_list: list[dict[str, Any]] = [ + {"type": "input_text", "text": wrapped}, + ] + content_list.extend(media_blocks) + + return [{"role": "user", "content": content_list}] + + @staticmethod + async def _format_system_message( + msg: Msg, + ) -> dict[str, Any]: + """Format a system message for the Responses API. + + Args: + msg (`Msg`): + The system message to format. + + Returns: + `dict[str, Any]`: + A dictionary with ``role`` and ``content`` keys. + """ + return { + "role": "system", + "content": msg.get_text_content(), + } diff --git a/src/agentscope/formatter/_xai_formatter.py b/src/agentscope/formatter/_xai_formatter.py new file mode 100644 index 0000000000000000000000000000000000000000..40a531175f9d767a00e867b8f41ecba7c535308b --- /dev/null +++ b/src/agentscope/formatter/_xai_formatter.py @@ -0,0 +1,485 @@ +# -*- coding: utf-8 -*- +"""The xAI formatter module. + +This formatter converts AgentScope ``Msg`` objects into the protobuf +``Message`` objects expected by the ``xai_sdk`` gRPC client. Unlike every +other formatter, the ``format()`` method returns a list of +``chat_pb2.Message`` proto objects rather than plain dicts, because the +``xai_sdk`` chat API accepts proto messages directly. +""" +import base64 +from typing import Any, List + +from pydantic import Field + +from ._formatter_base import FormatterBase +from .._logging import logger +from ..message import ( + Msg, + TextBlock, + ThinkingBlock, + ToolCallBlock, + ToolResultBlock, + DataBlock, + URLSource, + Base64Source, + HintBlock, +) + + +class XAIChatFormatter(FormatterBase): + """Formatter for the xAI chat model. + + Converts ``Msg`` objects into ``xai_sdk`` protobuf ``Message`` objects + that can be appended directly to a ``xai_sdk`` chat session. + + Unlike other formatters whose ``format()`` returns ``list[dict]``, this + formatter returns ``list[chat_pb2.Message]``. The type annotation is + intentionally widened to ``list[Any]`` to accommodate this difference. + """ + + input_types: list[str] = Field( + default=["text/plain", "image/jpeg", "image/png"], + description=( + "The supported input types. " + 'Defaults to ``["text/plain", "image/jpeg", "image/png"]``.' + ), + ) + + # pylint: disable=too-many-statements, too-many-branches + async def format( + self, + msgs: list[Msg], + **kwargs: Any, + ) -> List[Any]: + """Convert a list of ``Msg`` objects to ``xai_sdk`` proto messages. + + Args: + msgs (`list[Msg]`): + A list of ``Msg`` objects representing the conversation. + **kwargs (`Any`): + Unused; retained for interface compatibility. + + Returns: + `list[Any]`: + A list of ``chat_pb2.Message`` proto objects, ready to be + appended to a ``xai_sdk`` chat session via + ``chat.append()``. + """ + from xai_sdk.chat import ( + assistant, + image, + system, + tool_result, + user, + chat_pb2, + ) + + self.assert_list_of_msgs(msgs) + + xai_messages: List[Any] = [] + + for msg in msgs: + blocks = msg.get_content_blocks() + + text_blocks = [b for b in blocks if isinstance(b, TextBlock)] + + if msg.role == "system": + text = "\n".join(b.text for b in text_blocks) + xai_messages.append(system(text)) + + elif msg.role == "user": + content_args: list = [] + for block in blocks: + if isinstance(block, ThinkingBlock): + pass + elif isinstance(block, HintBlock): + if content_args: + xai_messages.append(user(*content_args)) + content_args = [] + if isinstance(block.hint, str): + xai_messages.append(user(block.hint)) + else: + hint_args = self._xai_user_args_from_blocks( + block.hint, + image, + ) + if hint_args: + xai_messages.append(user(*hint_args)) + elif isinstance(block, TextBlock): + content_args.append(block.text) + elif isinstance(block, DataBlock): + if block.source.media_type.startswith("image/"): + if isinstance(block.source, URLSource): + url_str = str(block.source.url) + if url_str.startswith("file://"): + # Local file — read and encode as data URI + local_path = url_str.removeprefix( + "file://", + ) + with open(local_path, "rb") as f: + encoded = base64.b64encode( + f.read(), + ).decode("utf-8") + content_args.append( + image( + f"data:{block.source.media_type};" + f"base64,{encoded}", + ), + ) + else: + content_args.append(image(url_str)) + elif isinstance(block.source, Base64Source): + content_args.append( + image( + f"data:{block.source.media_type};" + f"base64,{block.source.data}", + ), + ) + else: + logger.warning( + "Unsupported media type %s for xAI API. " + "Only image/jpeg and image/png are supported. " + "This block will be skipped.", + block.source.media_type, + ) + else: + logger.warning( + "Unsupported block type %s in user message, " + "skipped.", + type(block).__name__, + ) + + if content_args: + xai_messages.append(user(*content_args)) + + elif msg.role == "assistant": + pending_text: list[TextBlock] = [] + pending_tool_calls: list[ToolCallBlock] = [] + + for block in blocks: + if isinstance(block, ToolResultBlock): + # Convert each ToolResultBlock to a tool_result + # message. + if pending_tool_calls: + msg_proto = chat_pb2.Message() + msg_proto.role = chat_pb2.MessageRole.Value( + "ROLE_ASSISTANT", + ) + if pending_text: + c = msg_proto.content.add() + c.text = "\n".join( + b.text for b in pending_text + ) + pending_text = [] + for tc in pending_tool_calls: + proto_tc = msg_proto.tool_calls.add() + proto_tc.id = tc.id + proto_tc.type = chat_pb2.ToolCallType.Value( + "TOOL_CALL_TYPE_CLIENT_SIDE_TOOL", + ) + proto_tc.function.name = tc.name + proto_tc.function.arguments = tc.input + xai_messages.append(msg_proto) + pending_tool_calls = [] + elif pending_text: + text = "\n".join(b.text for b in pending_text) + if text: + xai_messages.append(assistant(text)) + pending_text = [] + + output_text = self._extract_result_text( + block.output, + ) + xai_messages.append( + tool_result( + output_text, + tool_call_id=block.id, + ), + ) + + elif isinstance(block, ToolCallBlock): + pending_tool_calls.append(block) + + elif isinstance(block, TextBlock): + pending_text.append(block) + + elif isinstance(block, ThinkingBlock): + pass + + elif isinstance(block, HintBlock): + if pending_tool_calls or pending_text: + if pending_tool_calls: + msg_proto = chat_pb2.Message() + msg_proto.role = chat_pb2.MessageRole.Value( + "ROLE_ASSISTANT", + ) + if pending_text: + c = msg_proto.content.add() + c.text = "\n".join( + b.text for b in pending_text + ) + pending_text = [] + for tc in pending_tool_calls: + proto_tc = msg_proto.tool_calls.add() + proto_tc.id = tc.id + proto_tc.type = ( + chat_pb2.ToolCallType.Value( + "TOOL_CALL_TYPE_CLIENT_SIDE_TOOL", + ) + ) + proto_tc.function.name = tc.name + proto_tc.function.arguments = tc.input + xai_messages.append(msg_proto) + pending_tool_calls = [] + elif pending_text: + text = "\n".join(b.text for b in pending_text) + if text: + xai_messages.append(assistant(text)) + pending_text = [] + if isinstance(block.hint, str): + xai_messages.append(user(block.hint)) + else: + hint_args = self._xai_user_args_from_blocks( + block.hint, + image, + ) + if hint_args: + xai_messages.append(user(*hint_args)) + + if pending_tool_calls: + # Assistant turn that triggered tool calls (history). + msg_proto = chat_pb2.Message() + msg_proto.role = chat_pb2.MessageRole.Value( + "ROLE_ASSISTANT", + ) + if pending_text: + c = msg_proto.content.add() + c.text = "\n".join(b.text for b in pending_text) + pending_text = [] + for tc in pending_tool_calls: + proto_tc = msg_proto.tool_calls.add() + proto_tc.id = tc.id + proto_tc.type = chat_pb2.ToolCallType.Value( + "TOOL_CALL_TYPE_CLIENT_SIDE_TOOL", + ) + proto_tc.function.name = tc.name + proto_tc.function.arguments = tc.input + xai_messages.append(msg_proto) + elif pending_text: + # Regular assistant text message. + text = "\n".join(b.text for b in pending_text) + if text: + xai_messages.append(assistant(text)) + + else: + logger.warning( + "Unsupported message role '%s', skipped.", + msg.role, + ) + + return xai_messages + + def _xai_user_args_from_blocks( + self, + blocks: list, + image: Any, + ) -> list: + """Convert a list of ``TextBlock | DataBlock`` into the positional + args expected by ``xai_sdk.chat.user(*args)``. + + DataBlocks that are not images (or use an unsupported media type) + are dropped with a warning, matching the existing user-role + DataBlock handling above. + + Args: + blocks (`list`): + The ``hint`` list from a :class:`HintBlock`. + image (`Any`): + The ``xai_sdk.chat.image`` constructor, passed in to keep + the local import contract identical to ``format``. + + Returns: + `list`: + Positional args for ``user(*args)``; empty if nothing + survived. + """ + args: list = [] + for sub in blocks: + if isinstance(sub, TextBlock): + args.append(sub.text) + elif isinstance(sub, DataBlock): + if not sub.source.media_type.startswith("image/"): + logger.warning( + "Unsupported media type %s for xAI API. " + "Only image/jpeg and image/png are supported. " + "This hint sub-block will be skipped.", + sub.source.media_type, + ) + continue + if isinstance(sub.source, URLSource): + url_str = str(sub.source.url) + if url_str.startswith("file://"): + local_path = url_str.removeprefix("file://") + with open(local_path, "rb") as f: + encoded = base64.b64encode(f.read()).decode( + "utf-8", + ) + args.append( + image( + f"data:{sub.source.media_type};" + f"base64,{encoded}", + ), + ) + else: + args.append(image(url_str)) + elif isinstance(sub.source, Base64Source): + args.append( + image( + f"data:{sub.source.media_type};" + f"base64,{sub.source.data}", + ), + ) + return args + + def _extract_result_text(self, output: Any) -> str: + """Extract a plain-text string from a ``ToolResultBlock`` output. + + Args: + output (`Any`): + The raw output of a ``ToolResultBlock``, which may be a + string, a list of blocks, or another type. + + Returns: + `str`: + A plain-text representation of the output. + """ + if output is None: + return "" + if isinstance(output, str): + return output + if isinstance(output, list): + parts = [] + for item in output: + if isinstance(item, TextBlock): + parts.append(item.text) + elif isinstance(item, str): + parts.append(item) + else: + parts.append(str(item)) + return "\n".join(parts) + return str(output) + + +class XAIMultiAgentFormatter(FormatterBase): + """Formatter for the xAI chat model in multi-agent conversations. + + Produces ``xai_sdk`` protobuf ``Message`` objects (same as + :class:`XAIChatFormatter`). Prior agent-to-agent messages are collapsed + into a single ``user`` message with ```` tags; tool + call / result sequences are delegated to :class:`XAIChatFormatter`. + + .. note:: ``format()`` returns ``list[Any]`` (protobuf messages), not + ``list[dict]``. + """ + + conversation_history_prompt: str = Field( + default=( + "# Conversation History\n" + "The content between tags contains " + "your conversation history\n" + ), + description="The prompt to use for the conversation history section.", + ) + + input_types: list[str] = Field( + default=["text/plain", "image/jpeg", "image/png"], + description=( + "The supported input types. " + 'Defaults to ``["text/plain", "image/jpeg", "image/png"]``.' + ), + ) + + async def format( + self, + msgs: list[Msg], + **kwargs: Any, + ) -> List[Any]: + """Convert a list of ``Msg`` objects to ``xai_sdk`` proto messages. + + Conversation history (non-tool messages) is collapsed into a single + ``user`` protobuf message containing ```` tags. + Tool sequences are formatted by :class:`XAIChatFormatter`. + + Args: + msgs (`list[Msg]`): + A list of ``Msg`` objects representing the conversation. + **kwargs (`Any`): + Unused; retained for interface compatibility. + + Returns: + `list[Any]`: + A list of ``chat_pb2.Message`` proto objects. + """ + from xai_sdk.chat import system, user + + self.assert_list_of_msgs(msgs) + + xai_messages: List[Any] = [] + start_index = 0 + + if msgs and msgs[0].role == "system": + text = msgs[0].get_text_content() + xai_messages.append(system(text)) + start_index = 1 + + is_first_agent_message = True + async for typ, group in self._group_messages(msgs[start_index:]): + if typ == "tool_sequence": + xai_messages.extend( + await XAIChatFormatter( + input_types=self.input_types, + ).format(group), + ) + elif typ == "agent_message": + history_text = self._build_history_text( + group, + is_first=is_first_agent_message, + ) + if history_text: + xai_messages.append(user(history_text)) + is_first_agent_message = False + + return xai_messages + + def _build_history_text(self, msgs: list[Msg], *, is_first: bool) -> str: + """Build a ```` text block from agent messages. + + Args: + msgs (`list[Msg]`): + Non-tool messages to collapse into history. + is_first (`bool`): + When ``True``, prepend + :attr:`conversation_history_prompt` before the tag. + + Returns: + `str`: + The formatted history string, or an empty string when there + is no text content. + """ + lines: list[str] = [] + for msg in msgs: + parts: list[str] = [] + if msg.name: + parts.append(f"{msg.name}:") + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + parts.append(block.text) + if parts: + lines.append(" ".join(parts)) + + if not lines: + return "" + + prefix = self.conversation_history_prompt if is_first else "" + return prefix + "\n" + "\n".join(lines) + "\n" diff --git a/src/agentscope/mcp/__init__.py b/src/agentscope/mcp/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..53032a6fc5d80e0dfced6b804c585fa2bb46b8e2 --- /dev/null +++ b/src/agentscope/mcp/__init__.py @@ -0,0 +1,13 @@ +# -*- coding: utf-8 -*- +"""The MCP module in AgentScope, that provides fine-grained control over +the MCP servers.""" + +from ._config import StdioMCPConfig, HttpMCPConfig +from ._mcp_client import MCPClient + + +__all__ = [ + "MCPClient", + "StdioMCPConfig", + "HttpMCPConfig", +] diff --git a/src/agentscope/mcp/_config.py b/src/agentscope/mcp/_config.py new file mode 100644 index 0000000000000000000000000000000000000000..3e7dc3ed272da8095cc4d16e3684684ce4bbb943 --- /dev/null +++ b/src/agentscope/mcp/_config.py @@ -0,0 +1,64 @@ +# -*- coding: utf-8 -*- +"""The MCP configurations.""" +from pathlib import Path +from typing import Literal + +from pydantic import BaseModel, Field + + +class StdioMCPConfig(BaseModel): + """The STDIO MCP server configuration.""" + + type: Literal["stdio_mcp"] = "stdio_mcp" + + command: str = Field( + title="Command", + description="The command to start the MCP server.", + ) + + args: list[str] | None = Field( + title="Args", + description="The command line arguments to pass to the MCP server.", + default=None, + ) + + env: dict[str, str] | None = Field( + title="Environment Variables", + default=None, + description="The environment variables to pass to the MCP server.", + ) + + cwd: str | Path | None = Field( + default=None, + title="CWD", + description="The working directory to use when spawning the process.", + ) + + encoding_error_handler: Literal["strict", "ignore", "replace"] = Field( + default="strict", + title="Encoding Error Handler", + description="The text encoding error handler.", + ) + + +class HttpMCPConfig(BaseModel): + """The HTTP MCP server configuration.""" + + type: Literal["http_mcp"] = "http_mcp" + + url: str = Field( + title="URL", + description="The URL of the MCP server.", + ) + + headers: dict[str, str] | None = Field( + title="Headers", + description="The additional headers to include in the HTTP request.", + default=None, + ) + + timeout: float | None = Field( + title="Timeout", + description="The HTTP request timeout in seconds.", + default=30.0, + ) diff --git a/src/agentscope/mcp/_mcp_client.py b/src/agentscope/mcp/_mcp_client.py new file mode 100644 index 0000000000000000000000000000000000000000..4bd05de54bb1b7adbbf6c47d943f0d1be486ecb5 --- /dev/null +++ b/src/agentscope/mcp/_mcp_client.py @@ -0,0 +1,438 @@ +# -*- coding: utf-8 -*- +"""Unified MCP client implementation for AgentScope.""" +import re +from contextlib import AsyncExitStack, _AsyncGeneratorContextManager +from typing import Any, TYPE_CHECKING + +import httpx +import mcp.types +from mcp import ClientSession, stdio_client, StdioServerParameters +from mcp.client.sse import sse_client +from mcp.client.streamable_http import streamable_http_client +from pydantic import Field, BaseModel, PrivateAttr + +from ._config import StdioMCPConfig, HttpMCPConfig +from .._logging import logger + +if TYPE_CHECKING: + from ..tool import MCPTool, ToolBase +else: + MCPTool = Any + ToolBase = Any + + +class MCPClient(BaseModel): + """The unified MCP client in AgentScope. + + This class provides a unified interface for MCP connections, handling both + stateful (persistent) and stateless (ephemeral) connections. + + - Stateful: Requires explicit connect() and close(), maintains session + - Stateless: No connect() needed, creates temporary session per call + + Private attributes: + - _client: The underlying MCP client context manager + - _session: The MCP ClientSession (for stateful connections only) + - _stack: AsyncExitStack for managing connection lifecycle + - _is_connected: Connection state flag + - _cached_tools: Cached list of tools + + Example: + + .. code-block:: python + + # Stateful connection (STDIO or HTTP) + client = MCPClient( + name="file_system", + is_stateful=True, + mcp_config=StdioMCPConfig( + command="mcp-server-filesystem" + ) + ) + await client.connect() + tools = await client.list_tools() + await client.close() + + # Stateless connection (HTTP only) + client = MCPClient( + name="weather_search", + is_stateful=False, + mcp_config=HttpMCPConfig( + url="https://api.weather.com/mcp" + ) + ) + # No connect() needed + tools = await client.list_tools() + + """ + + name: str = Field( + title="MCP Name", + description="The MCP name.", + ) + + is_stateful: bool = Field( + title="Stateful", + description=( + "Whether this is a stateful connection that requires explicit " + "connect() and close(). STDIO MCP must be stateful. HTTP MCP " + "can be either stateful or stateless." + ), + ) + + mcp_config: StdioMCPConfig | HttpMCPConfig = Field( + discriminator="type", + title="MCP Config", + description="The MCP server configuration.", + ) + + enable_tools: list[str] | None = None + """The tools enabled in this MCP, which will be returned in the + `list_tools` function. If `None`, all tools from the MCP server will be + returned.""" + + disable_tools: list[str] | None = None + """The tools disabled in this MCP, which will be filtered out in the + `list_tools` function.""" + + execution_timeout: float | None = None + """The execution timeout in seconds for calling the tools from this MCP.""" + + # Private attributes + _client: Any = PrivateAttr(default=None) + _session: ClientSession | None = PrivateAttr(default=None) + _stack: AsyncExitStack | None = PrivateAttr(default=None) + _is_connected: bool = PrivateAttr(default=False) + _cached_tools: list[mcp.types.Tool] | None = PrivateAttr(default=None) + + @property + def is_connected(self) -> bool: + """Whether the client is currently connected. + + Returns: + True if connected, False otherwise. + """ + return self._is_connected + + def model_post_init(self, __context: Any) -> None: + """Validate configuration and initialize client.""" + # MCP name is used to compose model-facing tool names + # (mcp__{name}__{tool}), which must match ^[a-zA-Z0-9_-]+$. + if not re.fullmatch(r"[a-zA-Z0-9_-]+", self.name): + raise ValueError( + f"MCPClient name '{self.name}' contains characters not " + f"allowed by LLM providers (only [a-zA-Z0-9_-] are " + f"permitted). Please rename it.", + ) + + # STDIO MCP must be stateful + if self.mcp_config.type == "stdio_mcp" and not self.is_stateful: + raise ValueError( + "STDIO MCP must be stateful (is_stateful=True).", + ) + + # Check arguments for self.enable_tools and disable_tools + if self.enable_tools is not None: + if not isinstance(self.enable_tools, list) or any( + not isinstance(_, str) for _ in self.enable_tools + ): + raise ValueError( + "Enable tools should be a list of strings, but got " + f"{self.enable_tools}.", + ) + + if self.disable_tools is not None: + if not isinstance(self.disable_tools, list) or any( + not isinstance(_, str) for _ in self.disable_tools + ): + raise ValueError( + "Disable tools should be a list of strings, but got " + f"{self.disable_tools}.", + ) + + if self.enable_tools is not None and self.disable_tools is not None: + intersection = set(self.enable_tools).intersection( + set(self.disable_tools), + ) + if len(intersection) != 0: + raise ValueError( + f"The tools in enable_tools and disable_tools " + f"should not overlap, but got {intersection}.", + ) + + # Initialize the underlying client + self._initialize_client() + + def _initialize_client(self) -> None: + """Pre-build the stdio client context manager. + + Only the stdio transport is materialised at construction time — + ``StdioServerParameters`` carry process-launch details that are + cheap to bind upfront. HTTP transports are built lazily inside + :meth:`connect` (or per-call inside :meth:`_get_client_gen` for + stateless mode), since each ``streamable_http_client`` / + ``sse_client`` is a one-shot context manager. + """ + if self.mcp_config.type == "stdio_mcp": + config = self.mcp_config + self._client = stdio_client( + StdioServerParameters( + command=config.command, + args=config.args or [], + env=config.env, + cwd=str(config.cwd) if config.cwd else None, + encoding="utf-8", + encoding_error_handler=config.encoding_error_handler, + ), + ) + + def _create_http_client( + self, + ) -> _AsyncGeneratorContextManager[Any]: + """Create an HTTP MCP client (SSE or streamable HTTP).""" + config = self.mcp_config + + # Determine transport from URL + if config.url.endswith("/sse") or config.url.endswith("/messages/"): + return sse_client( + url=config.url, + headers=config.headers, + timeout=config.timeout, + ) + + # StreamableHTTP transport + http_client = None + if config.headers or config.timeout: + http_client = httpx.AsyncClient( + headers=config.headers, + timeout=config.timeout, + ) + return streamable_http_client( + url=config.url, + http_client=http_client, + ) + + async def connect(self) -> None: + """Connect to the MCP server (for stateful connections only). + + For stateless connections, this method does nothing. + + Raises: + RuntimeError: If already connected. + """ + if not self.is_stateful: + logger.debug( + "Stateless MCP '%s' does not require explicit connect.", + self.name, + ) + return + + if self._is_connected: + raise RuntimeError( + f"MCP '{self.name}' is already connected. " + "Call close() before reconnecting.", + ) + + # Create HTTP client if needed + if self._client is None and self.mcp_config.type == "http_mcp": + self._client = self._create_http_client() + + self._stack = AsyncExitStack() + + try: + context = await self._stack.enter_async_context(self._client) + read_stream, write_stream = context[0], context[1] + self._session = ClientSession(read_stream, write_stream) + await self._stack.enter_async_context(self._session) + await self._session.initialize() + + self._is_connected = True + logger.info("MCP connected: %s", self.name) + except Exception: + await self._stack.aclose() + self._stack = None + raise + + async def close(self, ignore_errors: bool = True) -> None: + """Close the MCP connection (for stateful connections only). + + For stateless connections, this method does nothing. + + Args: + ignore_errors: Whether to ignore errors during cleanup. + + Raises: + RuntimeError: If not connected. + """ + if not self.is_stateful: + logger.debug( + "Stateless MCP '%s' does not require explicit close.", + self.name, + ) + return + + if not self._is_connected: + raise RuntimeError( + f"MCP '{self.name}' is not connected. " + "Call connect() first.", + ) + + try: + await self._stack.aclose() + except Exception as e: + if not ignore_errors: + raise e + logger.warning( + "Error closing MCP '%s': %s", + self.name, + str(e), + ) + finally: + self._stack = None + self._session = None + self._is_connected = False + logger.info("MCP closed: %s", self.name) + + def _get_client_gen(self) -> _AsyncGeneratorContextManager[Any]: + """Get client generator for stateless connections.""" + if self.mcp_config.type == "stdio_mcp": + return self._client + else: + return self._create_http_client() + + async def list_raw_tools(self) -> list[mcp.types.Tool]: + """List available tools from the MCP server in raw + :class:`mcp.types.Tool` form, applying ``enable_tools`` and + ``disable_tools`` filtering. + + The full (unfiltered) tool list is cached on ``_cached_tools`` so + :meth:`get_tool` can resolve names that were filtered out as well. + + Returns: + `list[mcp.types.Tool]`: + Raw MCP tool descriptors after filtering. + + Raises: + RuntimeError: If not connected (for stateful connections). + """ + if not self.is_stateful: + # Stateless: create temporary session + async with self._get_client_gen() as cli: + read_stream, write_stream = cli[0], cli[1] + async with ClientSession( + read_stream, + write_stream, + ) as session: + await session.initialize() + res = await session.list_tools() + self._cached_tools = res.tools + else: + # Stateful: use existing session + self._validate_connection() + res = await self._session.list_tools() + self._cached_tools = res.tools + + available_tools: list = self._cached_tools + if self.enable_tools is not None: + available_tools = [ + tool + for tool in available_tools + if tool.name in self.enable_tools + ] + if self.disable_tools is not None: + available_tools = [ + _ for _ in available_tools if _.name not in self.disable_tools + ] + return available_tools + + async def list_tools(self) -> list[ToolBase]: + """List available tools from the MCP server as wrapped + :class:`ToolBase` instances. If `enable_tools` and `disable_tools` + are not `None` in the constructor, the returned tools will be + filtered accordingly. + + Returns: + `list[ToolBase]`: + List of available MCP tools. + + Raises: + RuntimeError: If not connected (for stateful connections). + """ + raw_tools = await self.list_raw_tools() + return [await self.get_tool(_.name) for _ in raw_tools] + + async def get_tool( + self, + name: str, + ) -> MCPTool: + """Get a tool by name from the MCP server. + + The returned MCPTool object implements ToolProtocol and can be: + - Called directly: `await tool(arg1=val1)` + - Registered to toolkit: `toolkit.register_tool(tool)` + + Args: + name: The name of the tool function to get. + + Returns: + A tool object that implements ToolProtocol. + + Raises: + ValueError: If the tool is not found. + RuntimeError: If not connected (for stateful connections). + """ + # Avoid circular import by importing here + from ..tool import MCPTool + + # Fetch tools if not cached. Use list_raw_tools() to avoid the + # recursion list_tools() → get_tool() → list_tools(). + if self._cached_tools is None: + await self.list_raw_tools() + + # Find target tool + target_tool = None + for tool in self._cached_tools: + if tool.name == name: + target_tool = tool + break + + if target_tool is None: + raise ValueError( + f"Tool '{name}' not found in MCP server " f"'{self.name}'", + ) + + # Create MCPTool based on stateful/stateless + if not self.is_stateful: + # Stateless: pass client generator + return MCPTool( + mcp_name=self.name, + tool=target_tool, + client_gen=self._get_client_gen, + timeout=self.execution_timeout, + ) + else: + # Stateful: pass session + self._validate_connection() + return MCPTool( + mcp_name=self.name, + tool=target_tool, + session=self._session, + timeout=self.execution_timeout, + ) + + def _validate_connection(self) -> None: + """Validate connection state for stateful connections. + + Raises: + RuntimeError: If not connected or session not initialized. + """ + if not self._is_connected: + raise RuntimeError( + f"MCP '{self.name}' is not connected. " + "Call connect() first.", + ) + if not self._session: + raise RuntimeError( + f"MCP '{self.name}' session is not initialized. " + "Call connect() first.", + ) diff --git a/src/agentscope/message/__init__.py b/src/agentscope/message/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c03f43458b0329557580d49011e8b8a74147ee48 --- /dev/null +++ b/src/agentscope/message/__init__.py @@ -0,0 +1,39 @@ +# -*- coding: utf-8 -*- +"""The message module in agentscope.""" + +from ._block import ( + ContentBlock, + ContentBlockTypes, + TextBlock, + ThinkingBlock, + HintBlock, + ToolCallBlock, + ToolCallState, + ToolResultBlock, + ToolResultState, + DataBlock, + Base64Source, + URLSource, +) +from ._base import Msg, UserMsg, AssistantMsg, SystemMsg, Usage + + +__all__ = [ + "TextBlock", + "ThinkingBlock", + "HintBlock", + "ToolCallBlock", + "ToolCallState", + "ToolResultBlock", + "ToolResultState", + "DataBlock", + "Base64Source", + "URLSource", + "ContentBlock", + "ContentBlockTypes", + "Msg", + "UserMsg", + "AssistantMsg", + "SystemMsg", + "Usage", +] diff --git a/src/agentscope/message/_base.py b/src/agentscope/message/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..c3933df05277d87072f44df58f8703c72d3aa143 --- /dev/null +++ b/src/agentscope/message/_base.py @@ -0,0 +1,606 @@ +# -*- coding: utf-8 -*- +"""The message class in agentscope.""" +import base64 +from datetime import datetime +from typing import Literal, List, overload, Sequence, Self, TYPE_CHECKING, Any + +from pydantic import BaseModel, Field, model_validator + +from .._utils._common import _generate_id +from ._block import ( + TextBlock, + ThinkingBlock, + HintBlock, + DataBlock, + Base64Source, + URLSource, + ToolCallBlock, + ToolCallState, + ToolResultBlock, + ToolResultState, + ContentBlock, + ContentBlockTypes, +) +from .._logging import logger + +if TYPE_CHECKING: + from ..event import AgentEvent +else: + AgentEvent = Any + + +def _assert_user_content_blocks(content: Sequence[ContentBlock]) -> None: + """Assert that the content blocks in user message are valid.""" + for block in content: + if block.type not in ["text", "data"]: + raise ValueError( + "User message can only contain text blocks or data blocks.", + ) + + +def _assert_system_content_blocks( + content: Sequence[ContentBlock], +) -> None: + """Assert that the content blocks in system message are valid.""" + for block in content: + if block.type not in ["text"]: + raise ValueError("System message can only contain text blocks.") + + +def _to_blocks(content: str | list) -> list: + """Convert a plain string to a single-element TextBlock list.""" + if isinstance(content, str): + return [TextBlock(text=content)] + return content + + +class Usage(BaseModel): + """The token usage information of a message.""" + + input_tokens: int + """The number of input tokens.""" + output_tokens: int + """The number of output tokens.""" + + +class Msg(BaseModel): + """The message class in AgentScope, responsible for information storage + and transmission among different agents.""" + + name: str + """The name of the sender.""" + content: list[ContentBlock] + """The message content as a list of content blocks.""" + role: Literal["user", "assistant", "system"] + """The role of the sender.""" + id: str = Field(default_factory=_generate_id) + """The message identifier.""" + metadata: dict = Field(default_factory=dict) + """The metadata of the message""" + created_at: str = Field(default_factory=lambda: datetime.now().isoformat()) + """The creation time of the message""" + finished_at: str | None = Field(default=None) + """The finished time of the message""" + usage: Usage | None = Field(default=None) + """The token usage information of the message""" + + @model_validator(mode="after") + def validate_role_content(self) -> Self: + """Validate content blocks according to the role.""" + match self.role: + case "user": + _assert_user_content_blocks(self.content) + case "system": + _assert_system_content_blocks(self.content) + case "assistant": + pass + return self + + def has_content_blocks( + self, + block_type: ContentBlockTypes | list[ContentBlockTypes] | None = None, + ) -> bool: + """Check if the message has content blocks of the given type. + + Args: + block_type (`ContentBlockTypes | list[ContentBlockTypes] | None`, \ + optional): + The type of the block to be checked. If `None`, all blocks will + be checked. If a list is provided, it checks if there are + blocks of any types in the list. + + Returns: + `bool`: + `True` if there are content blocks of the given type, `False` + otherwise. + """ + if block_type is None: + return len(self.content) > 0 + + typs = [block_type] if isinstance(block_type, str) else block_type + return any(b.type in typs for b in self.content) + + def get_text_content(self, separator: str = "\n") -> str | None: + """Get the concatenated text from all TextBlocks.""" + gathered = [b.text for b in self.content if b.type == "text"] + return separator.join(gathered) if gathered else None + + @overload + def get_content_blocks( + self, + block_type: Literal["text"], + ) -> list[TextBlock]: + ... + + @overload + def get_content_blocks( + self, + block_type: Literal["thinking"], + ) -> list[ThinkingBlock]: + ... + + @overload + def get_content_blocks( + self, + block_type: Literal["tool_call"], + ) -> list[ToolCallBlock]: + ... + + @overload + def get_content_blocks( + self, + block_type: Literal["tool_result"], + ) -> list[ToolResultBlock]: + ... + + @overload + def get_content_blocks( + self, + block_type: Literal["data"], + ) -> list[DataBlock]: + ... + + @overload + def get_content_blocks( + self, + block_type: None = None, + ) -> list[ContentBlock]: + ... + + @overload + def get_content_blocks( + self, + block_type: Literal["hint"], + ) -> list[HintBlock]: + ... + + def get_content_blocks( + self, + block_type: ContentBlockTypes | List[ContentBlockTypes] | None = None, + ) -> Sequence[ContentBlock]: + """Get content blocks, optionally filtered by type. + + Args: + block_type (`ContentBlockTypes | List[ContentBlockTypes] | None`, \ + optional): + The type of the block to be extracted. If `None`, all blocks + will be returned. + + Returns: + `List[ContentBlock]`: + The content blocks. + """ + blocks: list[ContentBlock] = self.content or [] + if isinstance(block_type, str): + blocks = [b for b in blocks if b.type == block_type] + elif isinstance(block_type, list): + blocks = [b for b in blocks if b.type in block_type] + return blocks + + def _find_block( + self, + block_type: str, + block_id: str, + ) -> ContentBlock | None: + """Find a block in content by type and id.""" + for block in self.content: + if block.type == block_type and block.id == block_id: + return block + return None + + def append_event(self, event: AgentEvent) -> Self: + """Update the message by applying a streaming event. + + Mutates ``self.content``, ``self.finished_at``, and ``self.usage``: + content blocks are appended/updated by block-level events, + ``finished_at`` is stamped by ``REPLY_END``, and ``usage`` is + initialized then accumulated across each ``MODEL_CALL_END``. + Events whose ``reply_id`` does not match ``self.id`` are skipped with + a warning. Block-level delta/end events whose target block cannot be + found are also skipped with a warning. + + Args: + event (`AgentEvent`): + The event to apply. + """ + from ..event import EventType # local import to avoid circular dep + + if event.reply_id != self.id: + logger.warning( + "Event %s with reply_id %r does not match message id %r, " + "skipping.", + event.__class__.__name__, + event.reply_id, + self.id, + ) + return self + + match event.type: + case EventType.REPLY_END: + self.finished_at = event.created_at + + case EventType.MODEL_CALL_END: + if self.usage is None: + self.usage = Usage( + input_tokens=event.input_tokens, + output_tokens=event.output_tokens, + ) + else: + self.usage.input_tokens += event.input_tokens + self.usage.output_tokens += event.output_tokens + + case EventType.TEXT_BLOCK_START: + self.content.append(TextBlock(id=event.block_id, text="")) + + case EventType.TEXT_BLOCK_DELTA: + block = self._find_block("text", event.block_id) + if block is None: + logger.warning( + "TextBlock %r not found, skipping.", + event.block_id, + ) + else: + block.text += event.delta + + case EventType.TEXT_BLOCK_END: + pass + + case EventType.DATA_BLOCK_START: + self.content.append( + DataBlock( + id=event.block_id, + source=Base64Source( + data="", + media_type=event.media_type, + ), + ), + ) + + case EventType.DATA_BLOCK_DELTA: + block = self._find_block("data", event.block_id) + if block is None: + logger.warning( + "DataBlock %s not found, skipping.", + event.block_id, + ) + elif event.data: + # Each delta is an independently base64-encoded chunk + # (with its own padding); naive string concat would + # corrupt the byte stream. Decode, concat bytes, re-encode. + existing = ( + base64.b64decode(block.source.data) + if block.source.data + else b"" + ) + incoming = base64.b64decode(event.data) + block.source.data = base64.b64encode( + existing + incoming, + ).decode("ascii") + + case EventType.DATA_BLOCK_END: + pass + + case EventType.THINKING_BLOCK_START: + self.content.append( + ThinkingBlock(id=event.block_id, thinking=""), + ) + + case EventType.THINKING_BLOCK_DELTA: + block = self._find_block("thinking", event.block_id) + if block is None: + logger.warning( + "ThinkingBlock %r not found, skipping.", + event.block_id, + ) + else: + block.thinking += event.delta + + case EventType.THINKING_BLOCK_END: + pass + + case EventType.HINT_BLOCK: + # One-shot event — the full HintBlock content arrives in + # a single event, so just append it to ``content`` for + # persistence and replay. + self.content.append( + HintBlock( + id=event.block_id, + source=event.source, + hint=event.hint, + ), + ) + + case EventType.TOOL_CALL_START: + self.content.append( + ToolCallBlock( + id=event.tool_call_id, + name=event.tool_call_name, + input="", + ), + ) + + case EventType.TOOL_CALL_DELTA: + block = self._find_block("tool_call", event.tool_call_id) + if block is None: + logger.warning( + "ToolCallBlock %r not found, skipping.", + event.tool_call_id, + ) + else: + assert isinstance(block, ToolCallBlock) + block.input += event.delta + + case EventType.TOOL_CALL_END: + pass + + case EventType.TOOL_RESULT_START: + self.content.append( + ToolResultBlock( + id=event.tool_call_id, + name=event.tool_call_name, + output=[], + state=ToolResultState.RUNNING, + ), + ) + + case EventType.TOOL_RESULT_TEXT_DELTA: + block = self._find_block("tool_result", event.tool_call_id) + if block is None: + logger.warning( + "ToolResultBlock %r not found, skipping.", + event.tool_call_id, + ) + else: + assert isinstance(block, ToolResultBlock) + if isinstance(block.output, str): + block.output = [TextBlock(text=block.output)] + # Append the text + if not block.output or block.output[-1].type != "text": + block.output.append(TextBlock(text=event.delta)) + else: + block.output[-1].text += event.delta + + case EventType.TOOL_RESULT_DATA_DELTA: + block = self._find_block("tool_result", event.tool_call_id) + if block is None: + logger.warning( + "ToolResultBlock %r not found, skipping.", + event.tool_call_id, + ) + else: + assert isinstance(block, ToolResultBlock) + if isinstance(block.output, str): + block.output = [TextBlock(text=block.output)] + src = ( + Base64Source( + data=event.data, + media_type=event.media_type, + ) + if event.data is not None + else URLSource( + url=str(event.url), + media_type=event.media_type, + ) + ) + block.output.append( + DataBlock(id=event.block_id, source=src), + ) + + case EventType.TOOL_RESULT_END: + block = self._find_block("tool_result", event.tool_call_id) + if block is None: + logger.warning( + "ToolResultBlock %r not found, skipping.", + event.tool_call_id, + ) + else: + assert isinstance(block, ToolResultBlock) + block.state = event.state + block.metadata = event.metadata + # The paired ToolCallBlock's lifecycle ends with its + # result — flip it to FINISHED here so the SSE-rebuilt + # reply_msg matches ``agent.state.context``, which + # ``_update_tool_call_state`` mutates directly. + call_block = self._find_block("tool_call", event.tool_call_id) + if call_block is not None: + assert isinstance(call_block, ToolCallBlock) + call_block.state = ToolCallState.FINISHED + + case EventType.REQUIRE_USER_CONFIRM: + for tool_call in event.tool_calls: + b = self._find_block("tool_call", tool_call.id) + if b is not None: + assert isinstance(b, ToolCallBlock) + # Update the state + b.state = ToolCallState.ASKING + # Record the suggestions + b.suggested_rules = tool_call.suggested_rules + + case EventType.USER_CONFIRM_RESULT: + for result in event.confirm_results: + b = self._find_block("tool_call", result.tool_call.id) + if b is not None: + assert isinstance(b, ToolCallBlock) + b.state = ( + ToolCallState.ALLOWED + if result.confirmed + else ToolCallState.FINISHED + ) + + case EventType.REQUIRE_EXTERNAL_EXECUTION: + for tool_call in event.tool_calls: + b = self._find_block("tool_call", tool_call.id) + if b is not None: + assert isinstance(b, ToolCallBlock) + b.state = ToolCallState.SUBMITTED + + case EventType.EXTERNAL_EXECUTION_RESULT: + for result in event.execution_results: + self.content.append(result) + + return self + + +def UserMsg( + name: str, + content: str | list[TextBlock | DataBlock], + metadata: dict | None = None, + created_at: str | None = None, + finished_at: str | None = None, + id: str | None = None, # pylint: disable=redefined-builtin +) -> Msg: + """Create a user message with role ``"user"``. + + Args: + name (`str`): + The name of the sender. + content (`str | list[TextBlock | DataBlock]`): + The message content. A plain string will be automatically wrapped + in a :class:`TextBlock`. Only :class:`TextBlock` and + :class:`DataBlock` are allowed for user messages. + metadata (`dict | None`, optional): + Arbitrary key-value metadata attached to the message. Defaults to + an empty dict when not provided. + created_at (`str | None`, optional): + ISO-format timestamp for when the message was created. Defaults to + the current time when not provided. + finished_at (`str | None`, optional): + ISO-format timestamp for when the message was finished. Defaults to + the same value as ``created_at`` when not provided. + id (`str | None`, optional): + A unique identifier for the message. A random UUID hex string is + generated when not provided. + + Returns: + `Msg`: + A :class:`Msg` instance with ``role="user"``. + """ + created_at = created_at or datetime.now().isoformat() + if finished_at is None: + finished_at = created_at + return Msg( + name=name, + content=_to_blocks(content), + role="user", + metadata=metadata or {}, + created_at=created_at, + finished_at=finished_at, + id=id or _generate_id(), + ) + + +def AssistantMsg( + name: str, + content: str | list[ContentBlock], + metadata: dict | None = None, + created_at: str | None = None, + finished_at: str | None = None, + id: str | None = None, # pylint: disable=redefined-builtin + usage: Usage | None = None, +) -> Msg: + """Create an assistant message with role ``"assistant"``. + + Args: + name (`str`): + The name of the sender. + content (`str | list[ContentBlock]`): + The message content. A plain string will be automatically wrapped + in a :class:`TextBlock`. Any :class:`ContentBlock` subtype is + permitted for assistant messages. + metadata (`dict | None`, optional): + Arbitrary key-value metadata attached to the message. Defaults to + an empty dict when not provided. + created_at (`str | None`, optional): + ISO-format timestamp for when the message was created. Defaults to + the current time when not provided. + finished_at (`str | None`, optional): + ISO-format timestamp for when the message was finished. Not set by + default for assistant messages. + id (`str | None`, optional): + A unique identifier for the message. A random UUID hex string is + generated when not provided. + usage (`Usage | None`, optional): + The token usage information of the message. + + Returns: + `Msg`: + A :class:`Msg` instance with ``role="assistant"``. + """ + return Msg( + name=name, + content=_to_blocks(content), + role="assistant", + metadata=metadata or {}, + created_at=created_at or datetime.now().isoformat(), + finished_at=finished_at, + id=id or _generate_id(), + usage=usage, + ) + + +def SystemMsg( + name: str, + content: str | list[TextBlock], + metadata: dict | None = None, + created_at: str | None = None, + finished_at: str | None = None, + id: str | None = None, # pylint: disable=redefined-builtin +) -> Msg: + """Create a system message with role ``"system"``. + + Args: + name (`str`): + The name of the sender. + content (`str | list[TextBlock]`): + The message content. A plain string will be automatically wrapped + in a :class:`TextBlock`. Only :class:`TextBlock` is allowed for + system messages. + metadata (`dict | None`, optional): + Arbitrary key-value metadata attached to the message. Defaults to + an empty dict when not provided. + created_at (`str | None`, optional): + ISO-format timestamp for when the message was created. Defaults to + the current time when not provided. + finished_at (`str | None`, optional): + ISO-format timestamp for when the message was finished. Defaults to + the same value as ``created_at`` when not provided. + id (`str | None`, optional): + A unique identifier for the message. A random UUID hex string is + generated when not provided. + + Returns: + `Msg`: + A :class:`Msg` instance with ``role="system"``. + """ + created_at = created_at or datetime.now().isoformat() + if finished_at is None: + finished_at = created_at + return Msg( + name=name, + content=_to_blocks(content), + role="system", + metadata=metadata or {}, + created_at=created_at, + finished_at=finished_at, + id=id or _generate_id(), + ) diff --git a/src/agentscope/message/_block.py b/src/agentscope/message/_block.py new file mode 100644 index 0000000000000000000000000000000000000000..889246353b9f070fd46870be8da66dff27cb9877 --- /dev/null +++ b/src/agentscope/message/_block.py @@ -0,0 +1,207 @@ +# -*- coding: utf-8 -*- +"""The content blocks of messages.""" +from enum import StrEnum +from typing import Literal, List, TypeAlias, Any +from pydantic import BaseModel, Field, AnyUrl, field_serializer, ConfigDict + +from .._utils._common import _generate_id +from ..permission import PermissionRule + + +class TextBlock(BaseModel): + """The text block.""" + + type: Literal["text"] = "text" + """The type of the text block, which is always 'text'.""" + text: str + """The text content of the block.""" + id: str = Field(default_factory=_generate_id) + """The unique identifier of the block.""" + + +class ThinkingBlock(BaseModel): + """The thinking block. + + Allows extra provider-specific fields (e.g. Anthropic's ``signature``) + via ``extra="allow"`` so that model implementations can pass + arbitrary metadata without subclassing. + """ + + model_config = ConfigDict(extra="allow") + + type: Literal["thinking"] = "thinking" + """The type of the thinking block, which is always 'thinking'.""" + thinking: str + """The thinking content of the block.""" + id: str = Field(default_factory=_generate_id) + """The unique identifier of the block.""" + + +class Base64Source(BaseModel): + """The base64 source.""" + + type: Literal["base64"] = "base64" + """The type of the base64 source, which is always 'base64'.""" + data: str + """The base64-encoded data.""" + media_type: str + """The media type of the data, e.g., 'image/png', 'audio/mpeg', etc.""" + + +class URLSource(BaseModel): + """The URL source.""" + + type: Literal["url"] = "url" + """The type of the URL source, which is always 'url'.""" + url: AnyUrl + """A valid URI string conforming to RFC 3986.""" + media_type: str + """The media type of the data, e.g., 'image/png', 'audio/mpeg', etc.""" + + @field_serializer("url") + def serialize_url(self, url: AnyUrl) -> str: + """Serialize the URL to a string.""" + return str(url) + + +class DataBlock(BaseModel): + """The data block for binary content (images, audio, video, etc.).""" + + type: Literal["data"] = "data" + """The type of the data block, which is always 'data'.""" + id: str = Field(default_factory=_generate_id) + """The unique identifier of the block.""" + source: Base64Source | URLSource + """The source of the data, which can be either a base64-encoded string or + a URL.""" + name: str | None = None + """The name of the data block, which is optional.""" + + +class HintBlock(BaseModel): + """A block used to provide instructions or hints to the LLM during the + reasoning-acting loop. When passed to the LLM API, the hint block is + converted into a user message. + + The ``hint`` field can be a plain string (text-only) or a list of + :class:`TextBlock` / :class:`DataBlock` for multimodal content + (e.g. a background tool result containing both text and an image). + """ + + type: Literal["hint"] = "hint" + """The type of the hint block, which is always 'hint'.""" + hint: str | list[TextBlock | DataBlock] + """The hint content — plain text or a list of content blocks for + multimodal data.""" + id: str = Field(default_factory=_generate_id) + """The unique identifier of the block.""" + source: str | None = None + """The sender or origin of this hint. For team messages this is the + sender's display name (e.g. ``"alice"``); for system notifications + it may be ``"system"`` or ``None``.""" + + +class ToolCallState(StrEnum): + """The state of the tool call.""" + + PENDING = "pending" + ASKING = "asking" + ALLOWED = "allowed" + SUBMITTED = "submitted" + FINISHED = "finished" + + +class ToolCallBlock(BaseModel): + """The tool call block. + + Allows extra provider-specific fields (e.g. the OpenAI Responses API's + ``call_id``) via ``extra="allow"`` without requiring subclassing. + """ + + model_config = ConfigDict(use_enum_values=True, extra="allow") + + type: Literal["tool_call"] = "tool_call" + """The type of the tool call block, which is always 'tool_call'.""" + id: str + """The unique identifier of the tool call block.""" + name: str + """The name of the tool to be called.""" + input: str + """The raw JSON string input of the tool, accumulated during streaming.""" + state: ToolCallState = ToolCallState.PENDING + """The tool call state + - 'pending': the initial state when the tool call hasn't been processed + by the permission system + - 'asking': the tool call is asking and waiting for user confirmation + - 'allowed': allowed by the permission system/user and waits for execution + - 'submitted': the tool call has been submitted for external execution + and is waiting for results event + + Transitions + ----------- + pending + ├── permission DENY / input validation failed ──► finished + ├── permission ASK ──────────────────────────── ► asking + │ ├── user denied ───────────────────────► finished + │ └── user approved ─────────────────────► allowed + └── permission ALLOW ────────────────────────── ► allowed + + allowed + ├── local tool ── (execute) ─────────────────► finished + └── external tool ──────────────────────────── ► submitted + + submitted + └── ExternalExecutionResultEvent received ─────► finished + """ + suggested_rules: list[PermissionRule] = Field(default_factory=list) + """The suggestions for this tool call when asking user, used to maintain + the suggestions across requests.""" + + +class ToolResultState(StrEnum): + """The tool result state.""" + + SUCCESS = "success" + ERROR = "error" + INTERRUPTED = "interrupted" + DENIED = "denied" + RUNNING = "running" + + +class ToolResultBlock(BaseModel): + """The tool result block.""" + + model_config = ConfigDict(use_enum_values=True) + + type: Literal["tool_result"] = "tool_result" + """The type of the tool result block, which is always 'tool_result'.""" + id: str + """The unique identifier of the tool result block.""" + name: str + """The name of the tool.""" + output: str | List[TextBlock | DataBlock] + """The output of the tool, which can be a raw string of a list of + text and multimodal blocks.""" + state: ToolResultState = ToolResultState.RUNNING + """The execution state of the tool.""" + metadata: dict[str, Any] = Field(default_factory=dict) + """The metadata of the tool result block.""" + + +ContentBlock: TypeAlias = ( + TextBlock + | ThinkingBlock + | HintBlock + | ToolCallBlock + | ToolResultBlock + | DataBlock +) + +ContentBlockTypes: TypeAlias = Literal[ + "text", + "thinking", + "hint", + "tool_call", + "tool_result", + "data", +] diff --git a/src/agentscope/middleware/__init__.py b/src/agentscope/middleware/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..29bfb89e5200d8d55a8ee8aafa67327226095245 --- /dev/null +++ b/src/agentscope/middleware/__init__.py @@ -0,0 +1,19 @@ +# -*- coding: utf-8 -*- +"""Middleware system for AgentScope agents.""" + +from ._base import MiddlewareBase +from ._rag import RAGMiddleware +from ._budget import ReplyBudgetControlMiddleware +from ._longterm_memory import AgenticMemoryMiddleware, Mem0Middleware +from ._tracing import TracingMiddleware +from ._tts_middleware import TTSMiddleware + +__all__ = [ + "MiddlewareBase", + "AgenticMemoryMiddleware", + "Mem0Middleware", + "RAGMiddleware", + "TracingMiddleware", + "ReplyBudgetControlMiddleware", + "TTSMiddleware", +] diff --git a/src/agentscope/middleware/_base.py b/src/agentscope/middleware/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..6e0edbab95a465ce49dddb1cd1c7dfa2f01129a4 --- /dev/null +++ b/src/agentscope/middleware/_base.py @@ -0,0 +1,250 @@ +# -*- coding: utf-8 -*- +"""Base middleware class for AgentScope middleware system.""" +from typing import AsyncGenerator, Awaitable, Callable, TYPE_CHECKING + +from ..tool import ToolBase + +if TYPE_CHECKING: + from ..agent import Agent + from ..model import ChatResponse + + +class MiddlewareBase: # pylint: disable=unused-argument + """Base class for all middleware implementations. + + Middleware provides interception mechanisms at 5 key execution points + in the Agent lifecycle: + + **Onion Pattern Hooks** (with before/after logic): + - `on_reply`: Intercepts the entire reply process + - `on_reasoning`: Intercepts the reasoning/model call phase + - `on_acting`: Intercepts individual tool call execution + - `on_model_call`: Intercepts the raw model API call + + **Transformer Pattern Hook** (sequential pipeline): + - `on_system_prompt`: Transforms the system prompt string + + Each hook is optional - only implement the ones you need. The middleware + system will automatically detect which hooks are implemented at runtime. + + Example: + ```python + class LoggingMiddleware(MiddlewareBase): + async def on_reasoning( + self, + agent: Agent, + input_kwargs: dict, + next_handler: Callable[[], AsyncGenerator], + ) -> AsyncGenerator: + print(f"Before reasoning for agent {agent.name}") + async for event in next_handler(): + yield event + print(f"After reasoning for agent {agent.name}") + + agent = Agent( + ... + middlewares=[LoggingMiddleware()], + ... + ) + ``` + """ + + def is_implemented(self, hook_name: str) -> bool: + """Check if a hook method is implemented in the subclass. + + Args: + hook_name: Name of the hook method to check + + Returns: + True if the hook is implemented (overridden), False otherwise + """ + base_method = getattr(MiddlewareBase, hook_name, None) + sub_method = getattr(type(self), hook_name, None) + return base_method is not sub_method + + async def on_reply( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Hook for intercepting the reply process. + + Args: + agent: The Agent instance executing this middleware + input_kwargs: Dictionary containing: + - inputs: Msg | list[Msg] | UserConfirmResultEvent | + ExternalExecutionResultEvent | None — the unified inputs + that trigger this reply (new message(s), a resumption + event from a previous outside interaction, or None). + next_handler: Callable that executes the next middleware or + original method + + Yields: + AgentEvent | Msg: Events from the reply process + """ + raise RuntimeError( + f"{type(self).__name__} does not implement on_reply", + ) + yield # pylint: disable=unreachable + + async def on_reasoning( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Hook for intercepting the reasoning process. + + Args: + agent: The Agent instance executing this middleware + input_kwargs: Dictionary containing: + - tool_choice: ToolChoice (default None) + next_handler: Callable that executes the next middleware or + original method + + Yields: + Various events from the reasoning process + """ + raise RuntimeError( + f"{type(self).__name__} does not implement on_reasoning", + ) + yield # pylint: disable=unreachable + + async def on_acting( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Hook for intercepting the raw tool execution. + + This hook wraps **only** the ``toolkit.call_tool`` call — i.e. the + pure I/O execution layer. Permission checking, input validation, and + context writes are handled by the agent **outside** this hook and are + therefore not visible here. + + This separation makes it safe to offload the ``next_handler`` + coroutine to a background task: it will never mutate agent context + on its own. + + .. note:: + Tools with ``is_state_injected=True`` receive the live + ``agent.state`` object. Offloading such tools to a background + task may cause concurrent state mutations — guard against this + in your middleware implementation. + + Args: + agent (`Agent`): + The Agent instance executing this middleware. + input_kwargs (`dict`): + Dictionary containing: + + - ``tool_call`` (``ToolCallBlock``): the tool call to execute. + By the time this hook is invoked the tool call has already + been validated and permitted. + next_handler (`Callable[..., AsyncGenerator]`): + Callable that executes the next middleware or + ``_acting_impl``. + + Yields: + `ToolChunk | ToolResponse`: + Intermediate ``ToolChunk`` objects followed by a final + ``ToolResponse`` produced by the tool. + """ + raise RuntimeError( + f"{type(self).__name__} does not implement on_acting", + ) + yield # pylint: disable=unreachable + + async def on_model_call( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[ + ..., + Awaitable["ChatResponse" | AsyncGenerator["ChatResponse", None]], + ], + ) -> "ChatResponse" | AsyncGenerator["ChatResponse", None]: + """Hook for intercepting the model API call. + + Args: + agent: The Agent instance executing this middleware + input_kwargs: Dictionary containing: + - messages: list[Msg] + - tools: list[dict] + - tool_choice: ToolChoice + - current_model: The model instance used for this call + next_handler: Callable that executes the next middleware or + original method + + Returns: + ChatResponse or AsyncGenerator[ChatResponse, None] + """ + raise RuntimeError( + f"{type(self).__name__} does not implement on_model_call", + ) + + async def on_compress_context( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[..., Awaitable[None]], + ) -> None: + """Onion hook for `compress_context` function in `Agent` class + + Args: + agent (`Agent`): + The Agent instance executing this middleware + input_kwargs (`dict`): + Dictionary containing: + - context_config: ContextConfig | None + - instructions: HintBlock | None + next_handler (`Callable[..., Awaitable[None]]`): + Callable that executes the next middleware or + original method + """ + raise RuntimeError( + f"{type(self).__name__} does not implement on_compress_context", + ) + + async def on_system_prompt( + self, + agent: "Agent", + current_prompt: str, + ) -> str: + """Transform the system prompt string. + + This uses a transformer/pipeline pattern rather than onion pattern. + Multiple middlewares are applied sequentially, each receiving the + output of the previous one. + + Args: + agent: The Agent instance executing this middleware + current_prompt: The current system prompt string + + Returns: + str: The transformed system prompt + """ + raise RuntimeError( + f"{type(self).__name__} does not implement on_system_prompt", + ) + + async def list_tools(self) -> list[ToolBase]: + """List available tools provided by this middleware. Optional to + implement. + + Returns: + `list[ToolBase]`: + A list of tools provided by this middleware. + """ + return [] + + async def get_middleware_key(self) -> str: + """Get the unique key for this middleware, used to save middleware + states in `AgentState` instances. + + Optionally, middleware classes can override this method to + provide their own unique key. + """ + return self.__class__.__name__ diff --git a/src/agentscope/middleware/_budget.py b/src/agentscope/middleware/_budget.py new file mode 100644 index 0000000000000000000000000000000000000000..999e4533439f473587678f0564dbeb372f5c962b --- /dev/null +++ b/src/agentscope/middleware/_budget.py @@ -0,0 +1,207 @@ +# -*- coding: utf-8 -*- +"""Budget control middleware for AgentScope agents.""" +from typing import AsyncGenerator, Callable, TYPE_CHECKING + +from ..event import ModelCallEndEvent, ReplyStartEvent, ReplyEndEvent +from ..message import AssistantMsg, HintBlock +from ..tool import ToolChoice +from ._base import MiddlewareBase + +if TYPE_CHECKING: + from ..agent import Agent + +_DEFAULT_HINT_MESSAGE = ( + "You have reached the maximum token budget set by the " + "user. Now you MUST wrap up immediately and provide a final " + "concluding response without invoking any tools." + "" +) + + +class ReplyBudgetControlMiddleware(MiddlewareBase): + """Middleware that enforces a weighted token budget per reply. + + Tracks cumulative weighted token usage across all reasoning steps within + a single reply. The weighted cost is computed as:: + + cost = input_token_weight * input_tokens + output_token_weight * \ + output_tokens + + Once the accumulated cost reaches ``token_budget``, a hint message is + injected into the agent's context before the next reasoning step, and + ``tool_choice`` is forced to ``"none"`` so the agent wraps up without + invoking any further tools. + + Budget state is stored in + :attr:`~agentscope.agent.AgentState.middle_context` + keyed by the middleware key, so it persists across human-in-the-loop (HITL) + interruptions and resumptions. State is automatically cleaned up when the + reply ends via a :class:`~agentscope.event.ReplyEndEvent`. + + .. note:: + The middleware is stateless on the instance itself — all runtime state + lives in ``agent.state.middle_context``. This means the same middleware + instance can safely be shared across multiple agents. + + Example:: + + from agentscope.middleware import BudgetControlMiddleware + + agent = Agent( + ..., + middlewares=[ + BudgetControlMiddleware( + token_budget=10000, + input_token_weight=1.0, + output_token_weight=2.0, + ) + ], + ) + + """ + + def __init__( + self, + token_budget: float, + input_token_weight: float = 1, + output_token_weight: float = 1, + hint_message: str = _DEFAULT_HINT_MESSAGE, + ) -> None: + """Initialize the budget control middleware. + + Args: + token_budget (`float`): + Maximum weighted token cost allowed per reply. The cost for + each model call is computed as + ``input_token_weight * input_tokens + output_token_weight * + output_tokens``. Once the accumulated cost reaches this + threshold, the agent is instructed to wrap up without calling + any more tools. + input_token_weight (`float`, optional): + Multiplier applied to input tokens when computing the + weighted cost. Defaults to ``1``. + output_token_weight (`float`, optional): + Multiplier applied to output tokens when computing the + weighted cost. Defaults to ``1``. Set this higher than + ``input_token_weight`` to reflect that output tokens are + typically more expensive. + hint_message (`str`, optional): + The message injected into the agent's context when the budget + is exceeded. Defaults to a built-in wrap-up prompt. + """ + self.token_budget = token_budget + self.input_token_weight = input_token_weight + self.output_token_weight = output_token_weight + + self.hint_message = hint_message + + async def on_reply( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Manage per-reply budget state in ``agent.state.middle_context``. + + Initializes the weighted cost counter for the reply on + :class:`~agentscope.event.ReplyStartEvent`, accumulates cost on each + :class:`~agentscope.event.ModelCallEndEvent`, and removes the entry on + :class:`~agentscope.event.ReplyEndEvent`. + + Args: + agent (`Agent`): + The agent instance executing this middleware. + input_kwargs (`dict`): + Reply input kwargs (passed through unchanged). + next_handler (`Callable[..., AsyncGenerator]`): + Callable that executes the next middleware or ``_reply``. + + Yields: + Events from the reply process. + """ + + middleware_key = await self.get_middleware_key() + + async for event in next_handler(**input_kwargs): + if isinstance(event, ReplyStartEvent): + # Initialize the token counting number + if middleware_key not in agent.state.middle_context: + agent.state.middle_context[middleware_key] = {} + agent.state.middle_context[middleware_key][event.reply_id] = 0 + + elif isinstance(event, ReplyEndEvent): + # Clean up the token counting number + agent.state.middle_context[middleware_key].pop( + event.reply_id, + None, + ) + + elif isinstance(event, ModelCallEndEvent): + # Update the used tokens + if middleware_key not in agent.state.middle_context: + agent.state.middle_context[middleware_key] = {} + agent.state.middle_context[middleware_key][event.reply_id] += ( + self.input_token_weight * event.input_tokens + + self.output_token_weight * event.output_tokens + ) + + yield event + + async def on_reasoning( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Intercept each reasoning step to enforce the token budget. + + Before forwarding to the next handler, reads the accumulated weighted + cost for the current reply from ``agent.state.middle_context``. If the + budget is exhausted, appends a + :class:`~agentscope.message.HintBlock` to the last assistant message + in context (or creates a new + :class:`~agentscope.message.AssistantMsg`) and overrides + ``tool_choice`` to ``ToolChoice(mode="none")``. + + Args: + agent (`Agent`): + The agent instance executing this middleware. + input_kwargs (`dict`): + Dictionary containing ``tool_choice`` and other reasoning + kwargs forwarded to the next handler. + next_handler (`Callable[..., AsyncGenerator]`): + Callable that executes the next middleware or + ``_reasoning_impl``. + + Yields: + Events from the reasoning process. + """ + reply_id = agent.state.reply_id + middleware_key = await self.get_middleware_key() + used = agent.state.middle_context.get( + middleware_key, + {}, + ).get(reply_id, 0) + + # Insert hint block if exceeded budget + if used >= self.token_budget: + hint_block = HintBlock(hint=self.hint_message) + if ( + len(agent.state.context) > 0 + and agent.state.context[-1].role == "assistant" + and agent.state.context[-1].name == agent.name + ): + agent.state.context[-1].content.append(hint_block) + + else: + agent.state.context.append( + AssistantMsg( + id=agent.state.reply_id, + name=agent.name, + content=[hint_block], + ), + ) + input_kwargs["tool_choice"] = ToolChoice(mode="none") + + async for event in next_handler(**input_kwargs): + yield event diff --git a/src/agentscope/middleware/_longterm_memory/__init__.py b/src/agentscope/middleware/_longterm_memory/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f138dbb9f8ae7373296b929f6d861c3336fcb8ad --- /dev/null +++ b/src/agentscope/middleware/_longterm_memory/__init__.py @@ -0,0 +1,7 @@ +# -*- coding: utf-8 -*- +"""Long-term memory middlewares for AgentScope agents.""" + +from ._agentic_memory import AgenticMemoryMiddleware +from ._mem0 import Mem0Middleware + +__all__ = ["AgenticMemoryMiddleware", "Mem0Middleware"] diff --git a/src/agentscope/middleware/_longterm_memory/_agentic_memory/__init__.py b/src/agentscope/middleware/_longterm_memory/_agentic_memory/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fd55241815775ec82eaa6ff8a4a9e3e4da5ac455 --- /dev/null +++ b/src/agentscope/middleware/_longterm_memory/_agentic_memory/__init__.py @@ -0,0 +1,6 @@ +# -*- coding: utf-8 -*- +"""File-backed long-term memory middleware.""" + +from ._middleware import AgenticMemoryMiddleware + +__all__ = ["AgenticMemoryMiddleware"] diff --git a/src/agentscope/middleware/_longterm_memory/_agentic_memory/_middleware.py b/src/agentscope/middleware/_longterm_memory/_agentic_memory/_middleware.py new file mode 100644 index 0000000000000000000000000000000000000000..767949dc5ac621880469ec30352a9e39c58aa31a --- /dev/null +++ b/src/agentscope/middleware/_longterm_memory/_agentic_memory/_middleware.py @@ -0,0 +1,802 @@ +# -*- coding: utf-8 -*- +"""Filesystem-backed long-term memory middleware. + +The middleware keeps a workspace-local Markdown memory store, injects a +bounded ``MEMORY.md`` index into the system prompt, and can asynchronously +surface relevant topic files as hint blocks during the reasoning loop. +""" +from __future__ import annotations + +import asyncio +import re +from dataclasses import dataclass +from typing import TYPE_CHECKING, AsyncGenerator, Callable + +from pydantic import BaseModel, Field + +from ..._base import MiddlewareBase +from ...._logging import logger +from ...._utils._common import _estimate_tokens, _estimate_bytes +from ....message import Msg, SystemMsg, UserMsg, HintBlock +from ....model import ChatModelBase +from ....tool import BackendBase, LocalBackend + +if TYPE_CHECKING: + from ....agent import Agent + + +@dataclass(slots=True) +class _MemoryFileHeader: + """Lightweight header for one memory file, read from frontmatter only.""" + + filename: str + """Relative path under the memory directory (e.g. ``user_role.md``).""" + path: str + """Absolute path inside the backend.""" + description: str | None + """One-line description from frontmatter; ``None`` when absent.""" + type: str | None + """Memory type tag from frontmatter (user/feedback/project/reference).""" + mtime: float | None + """Modification time as a Unix timestamp; ``None`` when unavailable.""" + + +DEFAULT_MEMORY_INSTRUCTIONS = """# Auto Memory + +You have a persistent, file-based memory system at `{memory_dir}`. This directory already exists — write to it directly with the `Write` tool (do not run mkdir or check for its existence). + +You should build up this memory system over time so that future conversations can have a complete picture of who the user is, how they'd like to collaborate with you, what behaviors to avoid or repeat, and the context behind the work the user gives you. + +If the user explicitly asks you to remember something, save it immediately as whichever type fits best. If they ask you to forget something, find and remove the relevant entry. + +## Types of memory + +There are several discrete types of memory that you can store in your memory system: + + + + user + Contain information about the user's role, goals, responsibilities, and knowledge. Great user memories help you tailor your future behavior to the user's preferences and perspective. Your goal in reading and writing these memories is to build up an understanding of who the user is and how you can be most helpful to them specifically. For example, you should collaborate with a senior software engineer differently than a student who is coding for the very first time. Keep in mind, that the aim here is to be helpful to the user. Avoid writing memories about the user that could be viewed as a negative judgement or that are not relevant to the work you're trying to accomplish together. + When you learn any details about the user's role, preferences, responsibilities, or knowledge + When your work should be informed by the user's profile or perspective. For example, if the user is asking you to explain a part of the code, you should answer that question in a way that is tailored to the specific details that they will find most valuable or that helps them build their mental model in relation to domain knowledge they already have. + + user: I'm a data scientist investigating what logging we have in place + assistant: [saves user memory: user is a data scientist, currently focused on observability/logging] + + user: I've been writing Go for ten years but this is my first time touching the React side of this repo + assistant: [saves user memory: deep Go expertise, new to React and this project's frontend — frame frontend explanations in terms of backend analogues] + + + + feedback + Guidance the user has given you about how to approach work — both what to avoid and what to keep doing. These are a very important type of memory to read and write as they allow you to remain coherent and responsive to the way you should approach work in the project. Record from failure AND success: if you only save corrections, you will avoid past mistakes but drift away from approaches the user has already validated, and may grow overly cautious. + Any time the user corrects your approach ("no not that", "don't", "stop doing X") OR confirms a non-obvious approach worked ("yes exactly", "perfect, keep doing that", accepting an unusual choice without pushback). Corrections are easy to notice; confirmations are quieter — watch for them. In both cases, save what is applicable to future conversations, especially if surprising or not obvious from the code. Include *why* so you can judge edge cases later. + Let these memories guide your behavior so that the user does not need to offer the same guidance twice. + Lead with the rule itself, then a **Why:** line (the reason the user gave — often a past incident or strong preference) and a **How to apply:** line (when/where this guidance kicks in). Knowing *why* lets you judge edge cases instead of blindly following the rule. + + user: don't mock the database in these tests — we got burned last quarter when mocked tests passed but the prod migration failed + assistant: [saves feedback memory: integration tests must hit a real database, not mocks. Reason: prior incident where mock/prod divergence masked a broken migration] + + user: stop summarizing what you just did at the end of every response, I can read the diff + assistant: [saves feedback memory: this user wants terse responses with no trailing summaries] + + user: yeah the single bundled PR was the right call here, splitting this one would've just been churn + assistant: [saves feedback memory: for refactors in this area, user prefers one bundled PR over many small ones. Confirmed after I chose this approach — a validated judgment call, not a correction] + + + + project + Information that you learn about ongoing work, goals, initiatives, bugs, or incidents within the project that is not otherwise derivable from the code or git history. Project memories help you understand the broader context and motivation behind the work the user is doing within this working directory. + When you learn who is doing what, why, or by when. These states change relatively quickly so try to keep your understanding of this up to date. Always convert relative dates in user messages to absolute dates when saving (e.g., "Thursday" → "2026-03-05"), so the memory remains interpretable after time passes. + Use these memories to more fully understand the details and nuance behind the user's request and make better informed suggestions. + Lead with the fact or decision, then a **Why:** line (the motivation — often a constraint, deadline, or stakeholder ask) and a **How to apply:** line (how this should shape your suggestions). Project memories decay fast, so the why helps future-you judge whether the memory is still load-bearing. + + user: we're freezing all non-critical merges after Thursday — mobile team is cutting a release branch + assistant: [saves project memory: merge freeze begins 2026-03-05 for mobile release cut. Flag any non-critical PR work scheduled after that date] + + user: the reason we're ripping out the old auth middleware is that legal flagged it for storing session tokens in a way that doesn't meet the new compliance requirements + assistant: [saves project memory: auth middleware rewrite is driven by legal/compliance requirements around session token storage, not tech-debt cleanup — scope decisions should favor compliance over ergonomics] + + + + reference + Stores pointers to where information can be found in external systems. These memories allow you to remember where to look to find up-to-date information outside of the project directory. + When you learn about resources in external systems and their purpose. For example, that bugs are tracked in a specific project in Linear or that feedback can be found in a specific Slack channel. + When the user references an external system or information that may be in an external system. + + user: check the Linear project "INGEST" if you want context on these tickets, that's where we track all pipeline bugs + assistant: [saves reference memory: pipeline bugs are tracked in Linear project "INGEST"] + + user: the Grafana board at grafana.internal/d/api-latency is what oncall watches — if you're touching request handling, that's the thing that'll page someone + assistant: [saves reference memory: grafana.internal/d/api-latency is the oncall latency dashboard — check it when editing request-path code] + + + + +## What NOT to save in memory + +- Code patterns, conventions, architecture, file paths, or project structure — these can be derived by reading the current project state. +- Git history, recent changes, or who-changed-what — `git log` / `git blame` are authoritative. +- Debugging solutions or fix recipes — the fix is in the code; the commit message has the context. +- Anything already documented in CLAUDE.md files. +- Ephemeral task details: in-progress work, temporary state, current conversation context. + +These exclusions apply even when the user explicitly asks you to save. If they ask you to save a PR list or activity summary, ask what was *surprising* or *non-obvious* about it — that is the part worth keeping. + +## How to save memories + +Saving a memory is a two-step process: + +**Step 1** — write the memory to its own file (e.g., `user_role.md`, `feedback_testing.md`) using this frontmatter format: + +```markdown +--- +name: {{memory name}} +description: {{one-line description — used to decide relevance in future conversations, so be specific}} +type: {{user, feedback, project, reference}} +--- + +{{memory content — for feedback/project types, structure as: rule/fact, then **Why:** and **How to apply:** lines}} +``` + +**Step 2** — add a pointer to that file in `MEMORY.md`. `MEMORY.md` is an index, not a memory — each entry should be one line, under ~150 characters: - [Title](file.md) — one-line hook. It has no frontmatter. Never write memory content directly into MEMORY.md. + +- MEMORY.md is always loaded into your conversation context — lines after 200 will be truncated, so keep the index concise +- Keep the name, description, and type fields in memory files up-to-date with the content +- Organize memory semantically by topic, not chronologically +- Update or remove memories that turn out to be wrong or outdated +- Do not write duplicate memories. First check if there is an existing memory you can update before writing a new one. + +## When to access memories +- When memories seem relevant, or the user references prior-conversation work. +- You MUST access memory when the user explicitly asks you to check, recall, or remember. +- If the user says to ignore or not use memory: proceed as if MEMORY.md were empty. Do not apply remembered facts, cite, compare against, or mention memory content. +- Memory records can become stale over time. Use memory as context for what was true at a given point in time. Before answering the user or building assumptions based solely on information in memory records, verify that the memory is still correct and up-to-date by reading the current state of the files or resources. If a recalled memory conflicts with current information, trust what you observe now — and update or remove the stale memory rather than acting on it. + +## Before recommending from memory + +A memory that names a specific function, file, or flag is a claim that it existed *when the memory was written*. It may have been renamed, removed, or never merged. Before recommending it: +- If the memory names a file path: check the file exists. +- If the memory names a function or flag: grep for it. +- If the user is about to act on your recommendation (not just asking about history), verify first. + +"The memory says X exists" is not the same as "X exists now." + +A memory that summarizes repo state (activity logs, architecture snapshots) is frozen in time. If the user asks about *recent* or *current* state, prefer `git log` or reading the code over recalling the snapshot. + +## Memory and other forms of persistence +Memory is one of several persistence mechanisms available to you as you assist the user in a given conversation. The distinction is often that memory can be recalled in future conversations and should not be used for persisting information that is only useful within the scope of the current conversation. +- When to use or update a plan instead of memory: If you are about to start a non-trivial implementation task and would like to reach alignment with the user on your approach you should use a Plan rather than saving this information to memory. Similarly, if you already have a plan within the conversation and you have changed your approach persist that change by updating the plan rather than saving a memory. +- When to use or update tasks instead of memory: When you need to break your work in current conversation into discrete steps or keep track of your progress use tasks instead of saving to memory. Tasks are great for persisting information about the work that needs to be done in the current conversation, but memory should be reserved for information that will be useful in future conversations. + +## Searching past context + +When looking for past context: +1. Search topic files in your memory directory: +``` +Grep with pattern="" path="{memory_dir}" glob="*.md" +# or Bash command: +grep -rn "" {memory_dir} --include="*.md" +``` +Use narrow search terms (error messages, file paths, function names) rather than broad keywords. +""" # noqa: + +DEFAULT_RETRIEVAL_INSTRUCTIONS = ( + "You are selecting memory files that will be useful as context for " + "processing a user's query. You will be given the user's query and a " + "list of available memory files with their filenames and descriptions.\n\n" + "Return a list of filenames for the memories that will clearly be " + "useful (up to 5). Only include memories that you are certain will be " + "helpful based on their name and description.\n" + "- If you are unsure whether a memory will be useful, do not include " + "it. Be selective and discerning.\n" + "- If no memories would clearly be useful, return an empty list." +) + + +class _MemorySelection(BaseModel): + """Structured output schema for the memory relevance selector.""" + + selected_files: list[str] = Field( + description=( + "Filenames of the memory files to surface, relative to the " + "memory directory (e.g. 'user_role.md'). Up to 5 entries." + ), + ) + + +class AgenticMemoryMiddleware(MiddlewareBase): + """The agentic memory, where the LLM decides when and what to save, + together with an asyncio retrieval task in each reply. The memory is + stored and retrieval based on the Markdown files. + + The `AgenticMemoryMiddleware` supports different backends via the + `backend` argument in its constructor. + """ + + FILENAME_MEMORY_MD: str = "MEMORY.md" + + class Parameters(BaseModel): + """The user-tunable filesystem parameters.""" + + model_config = {"arbitrary_types_allowed": True} + + memory_max_tokens: int = Field( + default=4_000, + title="MEMORY.md Max Length", + description=( + "The maximum tokens of the MEMORY.md inserted into the system " + "prompt." + ), + ) + + memory_instructions: str = Field( + default=DEFAULT_MEMORY_INSTRUCTIONS, + title="Memory Instructions", + description=( + "The default instructions appended to the system prompt " + "before the MEMORY.md snapshots." + ), + ) + + retrieval_async: bool = Field( + default=True, + title="Async Retrieval", + description=( + "Whether to retrieve relevant memory files asynchronously " + "during the agent reply. If `True`, an async retrieval task " + "will be started." + ), + ) + + retrieval_model: ChatModelBase | None = Field( + default=None, + title="Retrieval Model", + description=( + "The LLM used to retrieve relevant memory files " + "asynchronously during the agent reply. If `None`, the " + "agent's model will be used." + ), + ) + + retrieval_max_tokens_per_md: int = Field( + default=2_000, + title="Retrieval Max Tokens Per File", + description=( + "Maximum tokens read from each memory file that is surfaced " + "by the relevance retrieval step. Keeps individual files from " + "flooding the context window." + ), + ) + + retrieval_max_files: int = Field( + default=200, + title="Retrieval Max Files", + description=( + "The maximum number of Markdown memory files to consider " + "during relevance selection." + ), + ) + + retrieval_max_tokens_per_frontmatter: int = Field( + default=256, + title="Retrieval Max Tokens per Frontmatter", + description=( + "The maximum number of tokens to read from the beginning of " + "each Markdown file when parsing frontmatter." + ), + ) + + retrieval_instructions: str = Field( + default=DEFAULT_RETRIEVAL_INSTRUCTIONS, + title="Retrieval Instructions", + description=( + "The instructions used to select relevant memory files for a " + "given user query in the asynchronous retrieval task." + ), + ) + + def __init__( + self, + *, + workdir: str, + memory_dir: str = "Memory", + parameters: Parameters | None = None, + backend: BackendBase | None = None, + ) -> None: + """Initialize filesystem-backed long-term memory behavior. + + Args: + workdir (`str`): + The working directory of this agent, used to store the agentic + searchable memory files. + memory_dir (`str`, defaults to ``"Memory"``): + The directory to store the long-term memory files, including + ``MEMORY.md``. + parameters (`AgenticMemoryMiddleware.Parameters | None`, \ + defaults to ``None``): + User-tunable parameters. When ``None``, defaults are used. + backend (`BackendBase | None`, optional): + The backend to switch between local and remote storage. + When ``None``, a local filesystem is used. + """ + self._workdir = workdir + self._memory_dir = memory_dir + self._parameters = parameters or self.Parameters() + self._backend = backend or LocalBackend() + + self._cached_input: str | None = None + # The in-flight asynchronous retrieval task started in ``on_reply`` + # and consumed in ``on_reasoning``. Kept on the instance so the + # reasoning hook can poll for completion across iterations. + self._retrieval_task: asyncio.Task | None = None + + @staticmethod + def _truncate_if_needed(content: str, max_length: int) -> str: + """Return ``content`` truncated to at most ``max_length`` tokens. + + Args: + content (`str`): + The content to truncate. + max_length (`int`): + The maximum estimated token count to keep. + + Returns: + `str`: + The original content when it already fits; otherwise a prefix + that fits within the requested token budget. + """ + if max_length <= 0: + return "" + + n_tokens = _estimate_tokens(content) + if n_tokens <= max_length: + return content + + index = int((max_length / n_tokens) * len(content)) + while index > 0 and _estimate_tokens(content[:index]) > max_length: + index = max(0, index - 10) + + return content[:index] + + async def on_system_prompt( + self, + agent: "Agent", + current_prompt: str, + ) -> str: + """Append memory instructions and a bounded ``MEMORY.md`` snapshot. + + Args: + agent (`Agent`): + The executing agent. Unused, but part of the middleware + contract. + current_prompt (`str`): + The system prompt produced by previous middleware. + + Returns: + `str`: + The prompt with filesystem memory instructions appended. + """ + await self._ensure_layout() + memory_md_content = await self._get_memory_md_content() or "" + + # Truncated by config + memory_md_truncated = self._truncate_if_needed( + memory_md_content, + self._parameters.memory_max_tokens, + ) + + if len(memory_md_truncated) != len(memory_md_content): + memory_md_path = self._get_memory_md_path() + remain_lines = len(memory_md_truncated.split("\n")) + omitted_lines = len(memory_md_content.split("\n")) - remain_lines + memory_md_truncated += ( + "\n<<>>\nThe remaining " + f"{omitted_lines} lines have been omitted due to context " + "length limits. Use the `Read` tool with offset " + f"`{remain_lines}` to access the rest of '{memory_md_path}'." + f"" + ) + + if not memory_md_truncated: + memory_md_truncated = ( + "Your MEMORY.md is currently empty. When you save new " + "memories, they will appear here." + ) + + memory_instructions = self._parameters.memory_instructions.replace( + "{memory_dir}", + self._get_memory_dir(), + ) + content = ( + f"{memory_instructions}\n" f"## MEMORY.md\n{memory_md_truncated}" + ) + + return f"{current_prompt}\n\n{content}" + + async def on_reply( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Cache the user input and kick off an asynchronous retrieval task + that runs concurrently with the agent reply. + + Args: + agent (`Agent`): + The executing agent whose model may be used for retrieval. + input_kwargs (`dict`): + Reply input kwargs forwarded unchanged. + next_handler (`Callable[..., AsyncGenerator]`): + The downstream middleware or core reply logic. + + Yields: + `Any`: + Items yielded by ``next_handler``. + """ + + if self._parameters.retrieval_async: + inputs = input_kwargs.get("inputs") + + msgs = None + if isinstance(inputs, list) and all( + isinstance(_, Msg) for _ in inputs + ): + msgs = inputs + elif isinstance(inputs, Msg): + msgs = [inputs] + + if msgs is not None: + self._cached_input = "\n".join( + [ + f"{_.name}: " + _.get_text_content() + for _ in msgs + if _.get_text_content() is not None + ], + ) + + # Start an asynchronous retrieval task that uses an LLM to decide + # which memory files are relevant to the current user input. The + # result is consumed by ``on_reasoning``. + if self._cached_input: + self._retrieval_task = asyncio.create_task( + self._retrieve_relevant_files(agent, self._cached_input), + ) + + try: + async for _ in next_handler(**input_kwargs): + yield _ + finally: + # Ensure the retrieval task does not outlive the reply. + if ( + self._retrieval_task is not None + and not self._retrieval_task.done() + ): + self._retrieval_task.cancel() + try: + await self._retrieval_task + except (asyncio.CancelledError, Exception): # noqa: BLE001 + pass + self._retrieval_task = None + self._cached_input = None + + async def on_reasoning( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Check if the retrieval finished and if yes, insert a hint block to + the content. + + Args: + agent (`Agent`): + The executing agent whose context may receive a hint block. + input_kwargs (`dict`): + Reasoning input kwargs forwarded unchanged. + next_handler (`Callable[..., AsyncGenerator]`): + The downstream middleware or core reasoning logic. + + Yields: + `Any`: + Items yielded by ``next_handler``. + """ + # Poll the in-flight retrieval task; if it has finished, consume its + # result and inject it into the agent context exactly once. + if self._retrieval_task is not None and self._retrieval_task.done(): + task = self._retrieval_task + self._retrieval_task = None + try: + retrieval_result = task.result() + except (asyncio.CancelledError, Exception): + retrieval_result = None + + if retrieval_result: + agent.state.append_context( + agent.name, + [ + HintBlock( + hint=retrieval_result, + ), + ], + ) + + async for event in next_handler(**input_kwargs): + yield event + + # ======================================================================== + # Helper functions + # ======================================================================== + + @staticmethod + def _format_manifest(headers: list[_MemoryFileHeader]) -> str: + """Format a list of memory file headers into a one-line-per-file + manifest string suitable for the selector prompt. + + Args: + headers (`list[_MemoryFileHeader]`): + The memory file headers to format. + + Returns: + `str`: + The formatted manifest. + """ + lines = [] + for h in headers: + tag = f"[{h.type}] " if h.type else "" + if h.mtime is not None: + from datetime import datetime + + ts = datetime.fromtimestamp(h.mtime).strftime("%Y-%m-%d") + else: + ts = "unknown" + desc = f": {h.description}" if h.description else "" + lines.append(f"- {tag}{h.filename} ({ts}){desc}") + return "\n".join(lines) + + async def _retrieve_relevant_files( + self, + agent: "Agent", + query: str, + ) -> str | None: + """Use an LLM to identify memory files relevant to ``query`` and + return their content as an injectable string. + + Args: + agent (`Agent`): + The agent whose model / memory store should be consulted. + query (`str`): + The cached user input used as the retrieval query. + + Returns: + `str | None`: + The formatted retrieval result ready to be injected into the + context, or ``None`` when nothing relevant was found. + """ + await self._ensure_layout() + + # 1. Scan available memory files (frontmatter only, cheap). + headers = await self._list_md_files() + if not headers: + return None + + valid_filenames = {h.filename for h in headers} + manifest = self._format_manifest(headers) + + # 2. Ask the model to select relevant files. + model = self._parameters.retrieval_model or agent.model + res = await model.generate_structured_output( + [ + SystemMsg( + name="system", + content=self._parameters.retrieval_instructions, + ), + UserMsg( + name="user", + content=( + f"Query: {query}\n\n" + f"Available memories:\n{manifest}" + ), + ), + ], + structured_model=_MemorySelection, + ) + + # 3. Validate: discard hallucinated filenames. + raw_selected: list[str] = res.content.get("selected_files", []) + selected = [f for f in raw_selected if f in valid_filenames][:5] + if not selected: + return None + + # 4. Read each selected file and format as an injectable string. + # Cap each file to avoid flooding the context window. + header_by_filename = {h.filename: h for h in headers} + parts: list[str] = [] + for filename in selected: + h = header_by_filename[filename] + try: + content = (await self._backend.read_file(h.path)).decode( + "utf-8", + errors="replace", + ) + except Exception: + continue + + # Truncate large files to avoid flooding the context window. + content = self._truncate_if_needed( + content, + self._parameters.retrieval_max_tokens_per_md, + ) + if h.mtime is not None: + import time + + days = max(0, int((time.time() - h.mtime) / 86_400)) + if days == 0: + age = "today" + elif days == 1: + age = "yesterday" + else: + age = f"{days} days ago" + header = f"Memory (saved {age}): {h.path}:" + else: + header = f"Memory: {h.path}:" + + parts.append(f"{header}\n\n{content}") + + if not parts: + return None + + return "\n\n---\n\n".join(parts) + + async def _ensure_layout(self) -> None: + """Create the memory directory and initial files idempotently. + + Existing human-edited documents are never replaced. The index file is + created only when absent so manual edits survive restarts. The memory + directory itself is materialized as a side effect of writing + ``MEMORY.md`` — :meth:`BackendBase.write_file` creates parent + directories — which avoids a platform-specific ``mkdir -p`` shell + invocation that is not portable on Windows. + """ + if not await self._backend.file_exists(self._get_memory_md_path()): + logger.info( + "Creating 'MEMORY.md' file in '%s'", + self._workdir, + ) + await self._backend.write_file( + self._get_memory_md_path(), + b"", + ) + + def _get_memory_dir(self) -> str: + """Get the memory directory. + + Returns: + `str`: + The backend path of the memory directory. + """ + return self._backend.join_path(self._workdir, self._memory_dir) + + def _get_memory_md_path(self) -> str: + """Get the ``MEMORY.md`` path. + + Returns: + `str`: + The backend path of the ``MEMORY.md`` index file. + """ + return self._backend.join_path( + self._get_memory_dir(), + self.FILENAME_MEMORY_MD, + ) + + async def _get_memory_md_content(self) -> str | None: + """Get the content of the ``MEMORY.md`` file. + + Returns: + `str | None`: + The decoded index file content, or ``None`` when the file does + not exist. + """ + if not await self._backend.file_exists(self._get_memory_md_path()): + return None + + return ( + await self._backend.read_file(self._get_memory_md_path()) + ).decode( + "utf-8", + errors="replace", + ) + + _FRONTMATTER_RE = re.compile( + r"^\s*---\s*\n(?P.*?)\n---\s*\n", + re.DOTALL, + ) + _FIELD_RE = re.compile(r"^(?P\w+)\s*:\s*(?P.+)$", re.MULTILINE) + + @classmethod + def _parse_frontmatter_fields(cls, content: str) -> dict[str, str]: + """Return a dict of YAML-like key/value pairs from the first + frontmatter block. Only scalar ``key: value`` lines are parsed; + nested structures are intentionally ignored. + + Args: + content (`str`): + The Markdown content prefix to inspect. + + Returns: + `dict[str, str]`: + Parsed frontmatter fields, or an empty dict when no leading + frontmatter block is present. + """ + m = cls._FRONTMATTER_RE.match(content) + if not m: + return {} + return { + fm.group("key"): fm.group("value").strip() + for fm in cls._FIELD_RE.finditer(m.group("body")) + } + + async def _list_md_files(self) -> list[_MemoryFileHeader]: + """Scan the memory directory for individual memory files. + + Returns: + `list[_MemoryFileHeader]`: + Memory file headers sorted newest-first and capped by + ``retrieval_max_files``. The system index file + (``MEMORY.md``) is excluded so only topic files are returned. + """ + memory_dir = self._get_memory_dir() + system_files = {self.FILENAME_MEMORY_MD} + + try: + all_entries = await self._backend.list_dir( + memory_dir, + recursive=True, + ) + except Exception: + return [] + + headers: list[_MemoryFileHeader] = [] + memory_dir_norm = self._backend.normpath(memory_dir) + memory_dir_prefix = self._backend.join_path(memory_dir_norm, "") + for entry in all_entries: + entry_path = self._backend.normpath(entry) + if self._backend.isabs(entry_path): + if not entry_path.startswith(memory_dir_prefix): + continue + filename = entry_path[len(memory_dir_prefix) :] + full_path = entry_path + else: + filename = entry_path + full_path = self._backend.join_path(memory_dir, filename) + + if not filename.endswith(".md") or filename in system_files: + continue + + try: + raw = await self._backend.read_file(full_path) + # Only parse the leading bytes to keep this cheap. + max_frontmatter_bytes = _estimate_bytes( + self._parameters.retrieval_max_tokens_per_frontmatter, + ) + snippet = raw[:max_frontmatter_bytes].decode( + "utf-8", + errors="replace", + ) + fields = self._parse_frontmatter_fields(snippet) + mtime = await self._backend.stat_mtime(full_path) + headers.append( + _MemoryFileHeader( + filename=filename, + path=full_path, + description=fields.get("description") or None, + type=fields.get("type") or None, + mtime=mtime, + ), + ) + except Exception: + continue + + headers.sort(key=lambda h: h.mtime or 0.0, reverse=True) + return headers[: self._parameters.retrieval_max_files] diff --git a/src/agentscope/middleware/_longterm_memory/_mem0/__init__.py b/src/agentscope/middleware/_longterm_memory/_mem0/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5daf2ed64d7e082bd3c4be27debd4b3d23601f5a --- /dev/null +++ b/src/agentscope/middleware/_longterm_memory/_mem0/__init__.py @@ -0,0 +1,6 @@ +# -*- coding: utf-8 -*- +"""mem0-backed long-term memory middleware.""" + +from ._middleware import Mem0Middleware + +__all__ = ["Mem0Middleware"] diff --git a/src/agentscope/middleware/_longterm_memory/_mem0/_agentscope_adapter.py b/src/agentscope/middleware/_longterm_memory/_mem0/_agentscope_adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..60d616483446264c40d1907cf460bfd7df0556ea --- /dev/null +++ b/src/agentscope/middleware/_longterm_memory/_mem0/_agentscope_adapter.py @@ -0,0 +1,439 @@ +# -*- coding: utf-8 -*- +"""Adapters that let mem0 drive its memory extraction with the user's +existing AgentScope chat / embedding model — instead of building yet +another OpenAI / Anthropic / Ollama client just for mem0. + +Two pieces: + +- :class:`AgentScopeLLM` — implements ``mem0.llms.base.LLMBase`` by + delegating to an ``agentscope.model.ChatModelBase`` instance. +- :class:`AgentScopeEmbedding` — implements + ``mem0.embeddings.base.EmbeddingBase`` by delegating to an + ``agentscope.embedding.EmbeddingModelBase`` instance. + +mem0 calls these synchronously; AgentScope models are async. We bridge +the two with ``asyncio.get_event_loop().run_until_complete()`` which +executes the coroutine on the caller's event loop. + +Use :func:`build_mem0_config` to produce a ``MemoryConfig`` wired to +AgentScope models: + + from mem0 import AsyncMemory + from agentscope.middleware._longterm_memory._mem0._agentscope_adapter \\ + import build_mem0_config + + mem0_client = AsyncMemory( + config=build_mem0_config( + chat_model=my_chat_model, + embedding_model=my_embedding_model, + ), + ) +""" +from __future__ import annotations + +import asyncio +import json +import threading +from collections.abc import AsyncGenerator +from typing import Any, TYPE_CHECKING + +from mem0.configs.embeddings.base import BaseEmbedderConfig +from mem0.configs.llms.base import BaseLlmConfig +from mem0.embeddings.base import EmbeddingBase +from mem0.llms.base import LLMBase + +from ....embedding import EmbeddingModelBase +from ....message import ( + AssistantMsg, + Msg, + SystemMsg, + UserMsg, +) +from ....model import ChatModelBase + +if TYPE_CHECKING: + from ....model import ChatResponse + + +# ---------------------------------------------------------------------- +# Sync → async bridge +# ---------------------------------------------------------------------- + + +class _AsyncBridge: + """Run async coroutines synchronously on a dedicated, long-lived + event loop. + + mem0's ``LLMBase.generate_response`` / ``EmbeddingBase.embed`` are + sync, but AgentScope models are async. The bridge owns one event + loop running on its own daemon thread and submits coroutines to it + via ``run_coroutine_threadsafe``, blocking for the result. The loop + stays alive for the bridge's lifetime, so async clients created + inside the model are reused and cleaned up against a live loop. + """ + + def __init__(self) -> None: + self._loop = asyncio.new_event_loop() + self._thread = threading.Thread( + target=self._loop.run_forever, + name="mem0-agentscope-bridge", + daemon=True, + ) + self._thread.start() + + def run(self, coro: Any) -> Any: + """Submit ``coro`` to the bridge loop and block for its result.""" + return asyncio.run_coroutine_threadsafe(coro, self._loop).result() + + +# ---------------------------------------------------------------------- +# LLM adapter +# ---------------------------------------------------------------------- + + +class AgentScopeLLM(LLMBase): + """mem0 ``LLMBase`` backed by an AgentScope ``ChatModelBase``. + + Pass your AgentScope model into ``config["model"]``; mem0's memory + extraction calls then route through it. Both streaming and + non-streaming AgentScope models are accepted (streaming responses are + drained and the final chunk is used). + """ + + def __init__( + self, + config: BaseLlmConfig | dict | None = None, + ) -> None: + """Initialize the AgentScope LLM for mem0.""" + super().__init__(config) + if self.config.model is None: + raise ValueError( + "AgentScopeLLM requires `model` in the config to be an " + "AgentScope ChatModelBase instance.", + ) + if not isinstance(self.config.model, ChatModelBase): + raise TypeError( + f"AgentScopeLLM `model` must be a ChatModelBase, got " + f"{type(self.config.model).__name__}.", + ) + self._agentscope_model: ChatModelBase = self.config.model + self._bridge = _AsyncBridge() + + # ----- LLMBase interface ----- + # pylint: disable=unused-argument + def generate_response( + self, + messages: list[dict[str, str]], + response_format: Any | None = None, # mem0 contract — unused + tools: list[dict] | None = None, + tool_choice: str = "auto", # mem0 contract — unused + ) -> str | dict: + """mem0 ``LLMBase`` entry — runs the AgentScope chat model + synchronously and returns str (or dict with tool_calls when + ``tools`` is given).""" + as_messages = _convert_messages_to_agentscope(messages) + if not as_messages: + raise ValueError( + "AgentScopeLLM received no usable messages " + "(empty list or all roles unrecognized).", + ) + + response = self._bridge.run( + _await_chat(self._agentscope_model, as_messages, tools), + ) + return _parse_chat_response(response, has_tool=bool(tools)) + + +async def _await_chat( + model: ChatModelBase, + messages: list[Msg], + tools: list[dict] | None, +) -> "ChatResponse": + """Call the AgentScope chat model, handling both streaming and + non-streaming returns.""" + result = await model(messages, tools=tools) + # Streaming model — drain the generator and keep the final chunk, + # which carries the complete content per AgentScope's streaming + # contract. ``isinstance`` (not ``hasattr``) — Pydantic BaseModel + # raises KeyError instead of AttributeError on missing dunder + # attrs, which ``hasattr`` does not catch. + if isinstance(result, AsyncGenerator): + last = None + async for chunk in result: + last = chunk + if last is None: + raise RuntimeError( + "AgentScope streaming model yielded no chunks.", + ) + return last + return result + + +def _convert_messages_to_agentscope( + messages: list[dict[str, str]], +) -> list[Msg]: + """mem0 hands us OpenAI-style ``[{"role", "content"}, ...]`` dicts; + AgentScope wants ``Msg`` objects.""" + out: list[Msg] = [] + for m in messages: + role = m.get("role") + content = m.get("content", "") + if role == "system": + out.append(SystemMsg(name="system", content=content)) + elif role == "user": + out.append(UserMsg(name="user", content=content)) + elif role == "assistant": + out.append(AssistantMsg(name="assistant", content=content)) + # unknown roles silently dropped — matches v1 behavior + return out + + +def _parse_chat_response( + response: "ChatResponse", + has_tool: bool, +) -> str | dict: + """Flatten an AgentScope ``ChatResponse`` into the str/dict shape + mem0 expects from ``LLMBase.generate_response``.""" + text_parts: list[str] = [] + thinking_parts: list[str] = [] + tool_parts: list[dict] = [] + + for block in response.content or []: + block_type = getattr(block, "type", None) + if block_type == "text": + text_parts.append(block.text or "") + elif block_type == "thinking": + thinking_parts.append(f"[Thinking: {block.thinking or ''}]") + elif block_type == "tool_call": + # AgentScope 2.0 stores tool args as a JSON string; mem0 + # expects an arguments' dict. + raw_input = block.input or "{}" + try: + arguments = json.loads(raw_input) + except json.JSONDecodeError: + arguments = raw_input + tool_parts.append( + {"name": block.name, "arguments": arguments}, + ) + # DataBlock and other types are not part of mem0's contract. + + text_out = "\n".join(thinking_parts + text_parts) + if has_tool: + return {"content": text_out, "tool_calls": tool_parts} + return text_out + + +# ---------------------------------------------------------------------- +# Embedding adapter +# ---------------------------------------------------------------------- + + +class AgentScopeEmbedding(EmbeddingBase): + """mem0 ``EmbeddingBase`` backed by an AgentScope + ``EmbeddingModelBase``.""" + + def __init__( + self, + config: BaseEmbedderConfig | dict | None = None, + ) -> None: + # mem0's EmbeddingBase (unlike LLMBase) does NOT auto-convert + # dict configs — it stores whatever is passed. Normalize here + # so callers can use the same dict-config style as the LLM. + if isinstance(config, dict): + config = BaseEmbedderConfig(**config) + super().__init__(config) + if self.config.model is None: + raise ValueError( + "AgentScopeEmbedding requires `model` in the config " + "to be an AgentScope EmbeddingModelBase instance.", + ) + if not isinstance(self.config.model, EmbeddingModelBase): + raise TypeError( + f"AgentScopeEmbedding `model` must be an " + f"EmbeddingModelBase, got " + f"{type(self.config.model).__name__}.", + ) + self._agentscope_model: EmbeddingModelBase = self.config.model + self._bridge = _AsyncBridge() + + # ----- EmbeddingBase interface ----- + # pylint: disable=unused-argument + def embed( + self, + text: str | list[str], + memory_action: str | None = None, # mem0 contract — unused + ) -> list[float]: + """mem0 ``EmbeddingBase`` entry — runs the AgentScope embedding + model synchronously and returns the first vector.""" + text_list = [text] if isinstance(text, str) else list(text) + response = self._bridge.run(self._agentscope_model(text_list)) + if not response.embeddings: + raise RuntimeError( + "AgentScope embedding model returned no embeddings.", + ) + # AgentScope EmbeddingResponse.embeddings is List[List[float]]; + # mem0 expects a single vector for a single-text call. + return response.embeddings[0] + + +# ---------------------------------------------------------------------- +# Build a mem0 MemoryConfig wired to AgentScope models +# ---------------------------------------------------------------------- + +# The provider name we register under in mem0's factory + config layer. +_AGENTSCOPE_PROVIDER = "agentscope" + + +def build_mem0_config( + *, + chat_model: ChatModelBase | None = None, + embedding_model: EmbeddingModelBase | None = None, + mem0_config: Any | None = None, +) -> Any: + """Build a ``mem0.configs.base.MemoryConfig`` wired to AgentScope + chat / embedding models. + + Three calling shapes: + + 1. ``build_mem0_config(chat_model=..., embedding_model=...)`` + — build a fresh config (mem0 defaults for vector_store / + history_db / reranker etc.) with both LLM and embedder routed + through AgentScope. + 2. ``build_mem0_config(mem0_config=cfg, chat_model=..., + embedding_model=...)`` — start from your customized + ``MemoryConfig`` and override only ``.llm`` / ``.embedder`` with + the AgentScope adapters. Use this when you want a non-default + vector store / history DB / reranker but still want AgentScope + to drive memory extraction. Either or both of ``chat_model`` + and ``embedding_model`` may be passed — fields you omit keep + whatever the input config had. + 3. ``build_mem0_config(mem0_config=cfg)`` — pass-through; the + AgentScope adapters are registered (cheap) but no fields are + overridden. + + Why a helper? Two private layers inside mem0 reject anything + outside its built-in provider list, so plugging in AgentScope + requires both: + + 1. Adding ``"agentscope"`` to ``LlmFactory.provider_to_class`` / + ``EmbedderFactory.provider_to_class`` so the factory can + construct our adapter classes. + 2. Bypassing the hardcoded provider whitelist in + ``LlmConfig.validate_config`` / ``EmbedderConfig.validate_config``. + Done by substituting subclasses whose validator only allows + ``"agentscope"``. Other provider names are rejected by these + subclasses, matching the fact that this helper is only for wiring + the AgentScope adapters. + + Args: + chat_model: + The AgentScope ``ChatModelBase`` mem0 should use for memory + extraction. Required when ``mem0_config`` is not given. + embedding_model: + The AgentScope ``EmbeddingModelBase`` mem0 should use to + embed memories. Its ``dimensions`` must match the vector + store's expected dim (mem0's default Qdrant expects 1536). + Required when ``mem0_config`` is not given. + mem0_config: + Optional pre-built ``MemoryConfig`` to use as the base. + When given, only the LLM / embedder slots are overridden + from ``chat_model`` / ``embedding_model`` — every other + field (``vector_store``, ``history_db_path``, ``reranker``, + ``custom_instructions``, ``version``) is preserved. + + Returns: + A ``MemoryConfig`` ready to pass to ``AsyncMemory(config=...)`` + or ``Memory(config=...)``. + """ + from mem0.configs.base import MemoryConfig + + _register_agentscope_provider() + llm_cfg_cls, emb_cfg_cls = _agentscope_config_classes() + + if mem0_config is None: + if chat_model is None or embedding_model is None: + raise ValueError( + "build_mem0_config requires `chat_model` and " + "`embedding_model` when `mem0_config` is not given.", + ) + return MemoryConfig( + llm=llm_cfg_cls( + provider=_AGENTSCOPE_PROVIDER, + config={"model": chat_model}, + ), + embedder=emb_cfg_cls( + provider=_AGENTSCOPE_PROVIDER, + config={"model": embedding_model}, + ), + ) + + # Use the user's config as base; partial-override .llm / .embedder + # only for fields they actually passed. Pydantic v2 doesn't + # re-validate on attribute assignment, so this sticks. + if chat_model is not None: + mem0_config.llm = llm_cfg_cls( + provider=_AGENTSCOPE_PROVIDER, + config={"model": chat_model}, + ) + if embedding_model is not None: + mem0_config.embedder = emb_cfg_cls( + provider=_AGENTSCOPE_PROVIDER, + config={"model": embedding_model}, + ) + return mem0_config + + +def _register_agentscope_provider() -> None: + """Plug the AgentScope adapter classes into mem0's factory dicts + under provider name ``"agentscope"``. Idempotent.""" + from mem0.utils.factory import EmbedderFactory, LlmFactory + + LlmFactory.provider_to_class[_AGENTSCOPE_PROVIDER] = ( + f"{__name__}.AgentScopeLLM", + BaseLlmConfig, + ) + EmbedderFactory.provider_to_class[ + _AGENTSCOPE_PROVIDER + ] = f"{__name__}.AgentScopeEmbedding" + + +def _agentscope_config_classes() -> tuple[type, type]: + """Return ``LlmConfig`` / ``EmbedderConfig`` subclasses whose + validator allows ``provider="agentscope"`` (and only that — every + other provider continues to be rejected with the same error mem0 + would have raised).""" + from pydantic import field_validator + + from mem0.embeddings.configs import EmbedderConfig + from mem0.llms.configs import LlmConfig + + class _AgentScopeLlmConfig(LlmConfig): + """``LlmConfig`` subclass that accepts the AgentScope provider.""" + + @field_validator("config") + @classmethod + def validate_config(cls, v: Any, values: Any) -> Any: + """Allow ``provider == "agentscope"``; reject everything + else with mem0's original error.""" + provider = values.data.get("provider") + if provider == _AGENTSCOPE_PROVIDER: + return v + raise ValueError(f"Unsupported LLM provider: {provider}") + + class _AgentScopeEmbedderConfig(EmbedderConfig): + """``EmbedderConfig`` subclass that accepts the AgentScope + provider.""" + + @field_validator("config") + @classmethod + def validate_config(cls, v: Any, values: Any) -> Any: + """Allow ``provider == "agentscope"``; reject everything + else with mem0's original error.""" + provider = values.data.get("provider") + if provider == _AGENTSCOPE_PROVIDER: + return v + raise ValueError( + f"Unsupported embedding provider: {provider}", + ) + + return _AgentScopeLlmConfig, _AgentScopeEmbedderConfig diff --git a/src/agentscope/middleware/_longterm_memory/_mem0/_middleware.py b/src/agentscope/middleware/_longterm_memory/_mem0/_middleware.py new file mode 100644 index 0000000000000000000000000000000000000000..fb8e53b31282277b79781e0476234d2a55753c35 --- /dev/null +++ b/src/agentscope/middleware/_longterm_memory/_mem0/_middleware.py @@ -0,0 +1,741 @@ +# -*- coding: utf-8 -*- +"""mem0-backed long-term memory middleware for AgentScope agents. + +Works with either ``mem0.AsyncMemory`` (open-source) or +``mem0.AsyncMemoryClient`` (hosted Platform). Both clients converge on +the same call shape: + +- ``search(query, filters={"user_id": ..., "agent_id": ...}, top_k=...)`` +- ``add(messages, user_id=..., agent_id=...)`` + +so one middleware class handles both. +""" +from __future__ import annotations + +import asyncio +import inspect +from typing import ( + Any, + AsyncGenerator, + Callable, + Literal, + TYPE_CHECKING, +) +from ..._base import MiddlewareBase +from ...._logging import logger +from ....event import ReplyStartEvent +from ....message import AssistantMsg, HintBlock, Msg +from ._tools import _build_memory_tools +from ._utils import ( + _extract_memory_texts, + _extract_query_text, + _mem0_extracted_anything, +) + +if TYPE_CHECKING: + from typing import TypeAlias + + from mem0 import AsyncMemory, AsyncMemoryClient + + from ....agent import Agent + from ....embedding import EmbeddingModelBase + from ....model import ChatModelBase + from ....tool import ToolBase + + # Explicit ``TypeAlias`` annotation tells mypy this is a type + # alias rather than a plain variable assignment — without it + # mypy 1.7+ refuses to use ``Mem0AsyncClient`` in annotation + # positions ("Variable is not valid as a type"). + Mem0AsyncClient: TypeAlias = AsyncMemory | AsyncMemoryClient + + +def _looks_async(method: Any) -> bool: + """Detect "async-callable" methods, including ones wrapped by a + sync ``@functools.wraps`` decorator that returns the underlying + coroutine. + + Plain ``inspect.iscoroutinefunction`` is too strict for mem0's + Platform client: ``AsyncMemoryClient.search`` is ``async def`` but + decorated by mem0's ``@api_error_handler``, which is a sync wrapper + that returns the coroutine produced by calling the underlying async + function. The wrapper itself isn't a coroutine function, so the + naive check rejects an otherwise-valid client. ``inspect.unwrap`` + walks ``__wrapped__`` (set by ``functools.wraps``) until it finds + the original async ``func``, which the strict check then accepts. + """ + if method is None: + return False + return inspect.iscoroutinefunction(inspect.unwrap(method)) + + +DEFAULT_MEMORY_SECTION_HEADER = "## Relevant memories from past conversations" +DEFAULT_MEMORY_SECTION_INTRO = ( + "The following memories about the user may be relevant. " + "Use them only if they are pertinent to the current request." +) + +DEFAULT_TOOL_INSTRUCTIONS = ( + "## Long-term memory\n\n" + "You have `search_memory` and `add_memory` tools available. Use " + "them whenever the conversation depends on (search) or contributes " + "(add) a durable fact about the user — see each tool's own " + "description for the exact input shape and usage guidance." +) + + +class Mem0Middleware(MiddlewareBase): + """AgentScope middleware that adds long-term memory backed by + `mem0 `_. + + Two construction paths: + + 1. **Models** — pass an AgentScope ``chat_model`` and + ``embedding_model`` (optionally with a custom ``mem0_config`` + as base for vector store / history DB / etc.). The middleware + builds an OSS ``AsyncMemory`` internally, wired so mem0's + memory extraction and embedding both go through your AgentScope + models. + 2. **Client** — pass a pre-built ``mem0.AsyncMemory`` / + ``mem0.AsyncMemoryClient`` when you want full control (e.g. + hosted Platform, sharing one mem0 across agents). When + ``client`` is given it takes precedence and the other backend + kwargs are ignored. + + Three control patterns are available via the ``mode`` parameter + (``"static_control"`` / ``"agent_control"`` / ``"both"``); see + the constructor's ``mode`` arg for what each does. + + Example (build OSS internally):: + + from agentscope.middleware import Mem0Middleware + from agentscope.tool import Toolkit + + mw = Mem0Middleware( + user_id="alice", + chat_model=my_chat_model, + embedding_model=my_embedding_model, + mode="both", + ) + agent = Agent( + ..., + toolkit=Toolkit(tools=await mw.list_tools()), + middlewares=[mw], + ) + + Example (hosted Platform with pre-built client):: + + from agentscope.middleware import Mem0Middleware + from agentscope.tool import Toolkit + from mem0 import AsyncMemoryClient + + mw = Mem0Middleware( + user_id="alice", + client=AsyncMemoryClient(api_key="m0-..."), + mode="both", + ) + agent = Agent( + ..., + toolkit=Toolkit(tools=await mw.list_tools()), + middlewares=[mw], + ) + """ + + def __init__( + self, + *, + user_id: str, + client: Mem0AsyncClient | None = None, + chat_model: "ChatModelBase | None" = None, + embedding_model: "EmbeddingModelBase | None" = None, + mem0_config: Any | None = None, + mode: Literal["static_control", "agent_control", "both"] = "both", + agent_id: str | None = None, + top_k: int = 5, + threshold: float | None = None, + scope_search_by_agent: bool = True, + await_write: bool = True, + memory_section_header: str = DEFAULT_MEMORY_SECTION_HEADER, + memory_section_intro: str = DEFAULT_MEMORY_SECTION_INTRO, + tool_instructions: str = DEFAULT_TOOL_INSTRUCTIONS, + ) -> None: + """Initialize the mem0 middleware. + + Three ways to wire up the mem0 backend: + + - **Models only** — pass ``chat_model`` + ``embedding_model`` + and the middleware builds a local OSS ``AsyncMemory`` wired + to them (mem0's default Qdrant for storage). + - **Models + ``mem0_config``** — same but starts from your + customized ``MemoryConfig``; only ``.llm`` / ``.embedder`` + slots are overridden with the AgentScope adapters, every + other field (vector store, history DB, reranker, ...) is + preserved. + - **Client** — pass a pre-built mem0 client (OSS / Platform / + custom). When ``client`` is given it takes absolute + precedence; ``chat_model`` / ``embedding_model`` / + ``mem0_config`` are ignored, with a warning listing any + ignored kwargs. + + Args: + user_id: + The mem0 ``user_id`` for memory namespacing. Required. + client: + A pre-built mem0 async client — + ``mem0.AsyncMemory`` (OSS) or + ``mem0.AsyncMemoryClient`` (Platform). Use this when + you want full control over the mem0 setup. + chat_model: + The AgentScope chat model mem0 should use for memory + extraction. Required if ``client`` is not given and + ``mem0_config`` does not already supply an LLM. + embedding_model: + The AgentScope embedding model mem0 should use to + embed memories. Required if ``client`` is not given + and ``mem0_config`` does not already supply an embedder. + Its ``dimensions`` must match mem0's vector store + (the default Qdrant expects 1536). + mem0_config: + Optional ``mem0.configs.base.MemoryConfig`` to use as + the base — lets you customize vector store / history + DB / reranker / etc. while still routing LLM and + embedding through AgentScope. Mutually exclusive with + ``client``. + mode: + How the agent interacts with mem0: + + - ``"static_control"``: middleware searches mem0 + before each reply, appends the retrieved memories + to ``agent.state.context`` as an + ``AssistantMsg(name="memory")``, and writes the new + exchange back after the reply. The agent never sees + mem0 as a tool. + - ``"agent_control"``: middleware exposes + ``search_memory`` / ``add_memory`` tools for the + agent to invoke on demand, plus a short nudge in + the system prompt. No automatic retrieval or + write-back. + - ``"both"``: both patterns at once — auto retrieval + AND on-demand tools. + + Defaults to ``"both"`` (matching AgentScope 1.x's + ``ReActAgent.long_term_memory_mode`` default). + agent_id: + Optional mem0 ``agent_id`` for finer-grained + namespacing. When ``None`` (default) mem0 receives no + ``agent_id`` filter and memories are scoped by + ``user_id`` only. + top_k: + Max number of memories retrieved per static-control + search. Also serves as the default ``top_k`` for the + ``search_memory`` tool (the agent can override). + threshold: + Minimum similarity score. ``None`` lets mem0 decide. + scope_search_by_agent: + When ``True`` (default) search filters include both + ``user_id`` and ``agent_id`` — memories are scoped to + the agent that created them. When ``False`` search uses + ``user_id`` only, so a user's memories are shared across + agents. + await_write: + When ``True`` (default) the post-turn ``add`` call is + awaited inline. When ``False`` it's fire-and-forget — + faster response but exceptions only surface in logs. + memory_section_header, memory_section_intro: + Strings used when injecting retrieved memories into + the model's messages list (``static_control`` / + ``both`` modes). + tool_instructions: + Markdown block appended to the agent's system prompt + in ``agent_control`` / ``both`` modes, advertising the + ``search_memory`` / ``add_memory`` tools to the LLM. + """ + is_empty_user_id = isinstance(user_id, str) and not user_id.strip() + if user_id is None or is_empty_user_id: + raise ValueError( + "Mem0Middleware requires a non-empty `user_id`.", + ) + if mode not in ("static_control", "agent_control", "both"): + raise ValueError( + f"Unknown mode {mode!r}; expected one of " + f"'static_control', 'agent_control', 'both'.", + ) + + client = self._resolve_client( + client=client, + chat_model=chat_model, + embedding_model=embedding_model, + mem0_config=mem0_config, + ) + self._client = client + + self._user_id = user_id + self._agent_id = agent_id + self._mode = mode + self._top_k = top_k + self._threshold = threshold + self._scope_search_by_agent = scope_search_by_agent + self._await_write = await_write + self._memory_section_header = memory_section_header + self._memory_section_intro = memory_section_intro + self._tool_instructions = tool_instructions + + # ------------------------------------------------------------------ + # Hook: on_reply + # ------------------------------------------------------------------ + async def on_reply( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + user_id = self._user_id + agent_id = self._agent_id + + # In pure agent_control mode the middleware is a no-op on the + # reply path — the agent decides when to invoke the memory + # tools — so just pass through. + if self._mode == "agent_control": + async for item in next_handler(**input_kwargs): + yield item + return + + # static_control / both, mirroring AgentScope 1.x's ReActAgent: + # 1. Pre-fetch memories from mem0 using the user's new query. + # 2. Once the agent has actually ingested the new user input + # into state.context (signaled by ReplyStartEvent — fires + # right after _handle_incoming_messages and before the + # reasoning loop), append the memory note. This places the + # note IMMEDIATELY AFTER the user message in context, same + # slot as v1's `_retrieve_from_long_term_memory` (which ran + # right after `self.memory.add(msg)`). + # 3. After the reply finishes, write the new exchange back. + # + # The memory note persists in state.context across turns. Long + # sessions will accumulate one per turn that retrieved + # anything; rely on ``compress_context`` or pop them yourself + # if that becomes a token concern. + # user_id / agent_id already resolved + cached above. + + inputs = input_kwargs.get("inputs") + query_text = _extract_query_text(inputs) + + memories: list[str] = [] + if query_text: + search_agent_id = agent_id if self._scope_search_by_agent else None + logger.info( + "mem0 search started: user_id=%s agent_id=%s chars=%d", + user_id, + search_agent_id, + len(query_text), + ) + try: + memories = await self._async_search( + query_text, + user_id=user_id, + agent_id=search_agent_id, + ) + logger.info( + "mem0 search finished: user_id=%s agent_id=%s memories=%d", + user_id, + search_agent_id, + len(memories), + ) + except Exception as e: # noqa: BLE001 + logger.warning( + "mem0 search failed for user_id=%s: %s", + user_id, + e, + ) + + final_msg: Msg | None = None + injected = False + try: + async for item in next_handler(**input_kwargs): + is_reply_start = isinstance(item, ReplyStartEvent) + if not injected and memories and is_reply_start: + agent.state.context.append( + self._build_memory_message(memories), + ) + injected = True + if isinstance(item, Msg) and item.role == "assistant": + final_msg = item + yield item + finally: + if query_text and final_msg is not None: + assistant_text = final_msg.get_text_content() + if assistant_text: + await self._dispatch_write( + [ + {"role": "user", "content": query_text}, + {"role": "assistant", "content": assistant_text}, + ], + user_id=user_id, + agent_id=agent_id, + ) + + # ------------------------------------------------------------------ + # Hook: on_system_prompt (advertise memory tools to the LLM) + # ------------------------------------------------------------------ + async def on_system_prompt( + self, + agent: "Agent", + current_prompt: str, + ) -> str: + """Append memory-tool instructions to the system prompt. + + Args: + agent (`Agent`): + The agent whose system prompt is being transformed. + current_prompt (`str`): + The system prompt produced by previous middleware. + + Returns: + `str`: + The unchanged prompt in static-control mode, otherwise the + prompt with memory-tool instructions appended. + """ + if self._mode == "static_control": + return current_prompt + return f"{current_prompt}\n\n{self._tool_instructions}" + + async def list_tools(self) -> list["ToolBase"]: + """List memory tools provided by this middleware. + + Returns: + `list[ToolBase]`: + The ``search_memory`` and ``add_memory`` tools in + agent-control modes, otherwise an empty list. + """ + if self._mode == "static_control": + return [] + return _build_memory_tools(self) + + # ================================================================== + # mem0 client construction + # ================================================================== + @staticmethod + def _resolve_client( + *, + client: "Mem0AsyncClient | None", + chat_model: "ChatModelBase | None", + embedding_model: "EmbeddingModelBase | None", + mem0_config: Any | None, + ) -> "Mem0AsyncClient": + """Resolve the constructor's mem0-backend kwargs into a single + async client. + + ``client=`` takes absolute precedence — if given, the other + three (``chat_model`` / ``embedding_model`` / ``mem0_config``) + are ignored (a warning is logged so the mismatch is not + invisible). Otherwise, the rest are combined by + :func:`build_mem0_config` into an ``AsyncMemory``. + """ + if client is not None: + ignored = [ + name + for name, value in ( + ("chat_model", chat_model), + ("embedding_model", embedding_model), + ("mem0_config", mem0_config), + ) + if value is not None + ] + if ignored: + logger.warning( + "Mem0Middleware: `client` was provided, so %s " + "%s ignored. Pass them via the mem0 client itself " + "(or omit `client` to let the middleware build one " + "for you).", + ", ".join(ignored), + "is" if len(ignored) == 1 else "are", + ) + if client is None: + no_models = chat_model is None and embedding_model is None + if mem0_config is None and no_models: + raise ValueError( + "Mem0Middleware needs one of: a pre-built `client`, " + "a `mem0_config`, or both `chat_model` and " + "`embedding_model`.", + ) + # When no mem0_config is given, models must come as a pair. + if mem0_config is None and ( + (chat_model is None) ^ (embedding_model is None) + ): + raise ValueError( + "Mem0Middleware: `chat_model` and " + "`embedding_model` must be passed together when " + "`mem0_config` is not given.", + ) + + from mem0 import AsyncMemory + + from ._agentscope_adapter import build_mem0_config + + client = AsyncMemory( + config=build_mem0_config( + chat_model=chat_model, + embedding_model=embedding_model, + mem0_config=mem0_config, + ), + ) + + if not _looks_async( + getattr(client, "search", None), + ) or not _looks_async(getattr(client, "add", None)): + raise TypeError( + "Mem0Middleware requires an async mem0 client " + "(`mem0.AsyncMemory` or `mem0.AsyncMemoryClient`). " + "The synchronous `Memory` / `MemoryClient` are not " + "supported.", + ) + return client + + # ================================================================== + # mem0 client adapters (OSS + Platform share this call shape) + # ================================================================== + async def _async_search( + self, + query: str, + *, + user_id: str, + agent_id: str | None, + top_k: int | None = None, + ) -> list[str]: + """Search mem0 and normalize the result into memory strings. + + Args: + query (`str`): + Search query text. + user_id (`str`): + mem0 ``user_id`` namespace. + agent_id (`str | None`): + Optional mem0 ``agent_id`` namespace. ``None`` searches by + user only. + top_k (`int | None`, optional): + Optional per-call result limit. Defaults to the middleware's + configured ``top_k``. + + Returns: + `list[str]`: + Retrieved memory texts extracted from mem0's response. + """ + filters: dict[str, Any] = {"user_id": user_id} + if agent_id: + filters["agent_id"] = agent_id + + kwargs: dict[str, Any] = { + "filters": filters, + "top_k": self._top_k if top_k is None else top_k, + } + if self._threshold is not None: + kwargs["threshold"] = self._threshold + + raw = await self._client.search(query, **kwargs) + return _extract_memory_texts(raw) + + async def _async_add( + self, + messages: list[dict[str, str]], + *, + user_id: str, + agent_id: str | None, + infer: bool = True, + ) -> dict | None: + """Add messages to mem0 using the shared async client call shape. + + Args: + messages (`list[dict[str, str]]`): + mem0-compatible message dictionaries. + user_id (`str`): + mem0 ``user_id`` namespace. + agent_id (`str | None`): + Optional mem0 ``agent_id`` namespace. + infer (`bool`, optional): + Whether mem0 should run memory extraction. ``False`` asks + mem0 to store the text directly. + + Returns: + `dict | None`: + The raw result returned by mem0. + """ + kwargs: dict[str, Any] = {"user_id": user_id} + if agent_id: + kwargs["agent_id"] = agent_id + if not infer: + # mem0 docstring: ``infer=False`` skips the LLM extraction + # step and stores the message text directly. + kwargs["infer"] = False + + return await self._client.add(messages, **kwargs) + + async def _async_add_with_fallback( + self, + text: str, + *, + user_id: str, + agent_id: str | None, + ) -> dict | None: + """Two-tier add strategy: try extraction first; if mem0's + extraction LLM returns no memories, fall back to ``infer=False`` + and save the raw text. Guarantees that ``add_memory`` always + persists *something* — matching AgentScope 1.x's + ``record_to_memory`` "always save" contract. + + Historical note — why this is 2 tiers, not 3 + --------------------------------------------- + AgentScope 1.x's ``record_to_memory`` had a 3-tier fallback: + (a) user role → (b) assistant role → (c) assistant + infer=False. + Tier (b) was meaningful against **old** mem0, which routed + user-role messages through ``USER_MEMORY_EXTRACTION_PROMPT`` + and assistant-role messages through + ``AGENT_MEMORY_EXTRACTION_PROMPT`` — two genuinely different + prompts, so switching role had a real chance of rescuing an + empty extraction. + + Current mem0 (v2.x) restructured this in + ``_add_to_vector_store``: + + parsed_messages = parse_messages(messages) + ... + is_agent_scoped = bool(filters.get("agent_id")) \\ + and not filters.get("user_id") + system_prompt = ADDITIVE_EXTRACTION_PROMPT + if is_agent_scoped: + system_prompt += AGENT_CONTEXT_SUFFIX + + Prompt selection now depends on the **filters dict**, not on + the message role. ``ADDITIVE_EXTRACTION_PROMPT`` itself says + explicitly "You extract from BOTH user and assistant messages" + (it just changes attribution framing). Since this middleware + always passes both ``user_id`` and ``agent_id`` (when scoping + is on), ``is_agent_scoped`` is always False and the same + ``ADDITIVE_EXTRACTION_PROMPT`` runs regardless of message role. + Retrying with role="assistant" is a no-op LLM call. + + So tier (b) is dropped. Tier (c) ``infer=False`` is still + valuable — it bypasses extraction entirely and saves raw + text, useful when mem0 decides nothing in the input is + memory-worthy but the caller wants the bytes persisted anyway. + """ + # 1. Normal path: let mem0's extraction LLM do its job. + result = await self._async_add( + [{"role": "user", "content": text, "name": "user"}], + user_id=user_id, + agent_id=agent_id, + ) + if _mem0_extracted_anything(result): + return result + + # 2. Raw save: extraction returned empty, persist the raw text + # so the caller's ``add_memory`` invocation isn't silently + # discarded. + return await self._async_add( + [{"role": "user", "content": text, "name": "user"}], + user_id=user_id, + agent_id=agent_id, + infer=False, + ) + + # ================================================================== + # Helpers + # ================================================================== + def _build_memory_message(self, memories: list[str]) -> Msg: + """Format retrieved ``memories`` as a synthetic hint message. + + The context entry uses an assistant-role ``Msg`` container because + user messages cannot carry ``HintBlock`` content. Formatters convert + the ``HintBlock`` itself into a user message before the model call. + + Args: + memories (`list[str]`): + Retrieved memory texts to expose to the model. + + Returns: + `Msg`: + An assistant-role message containing one ``HintBlock``. + """ + bullets = "\n".join(f"- {m}" for m in memories) + content = ( + f"{self._memory_section_header}\n" + f"{self._memory_section_intro}\n" + f"{bullets}" + ) + return AssistantMsg( + name="memory", + content=[HintBlock(hint=content)], + ) + + async def _dispatch_write( + self, + messages: list[dict[str, str]], + *, + user_id: str, + agent_id: str | None, + ) -> None: + """Persist a completed user/assistant exchange to mem0. + + When ``await_write`` is enabled, the add call is awaited inline so + errors can be logged before reply cleanup completes. Otherwise, to + write is scheduled as a background task and failures are logged there. + + Args: + messages (`list[dict[str, str]]`): + User/assistant message pair to persist. + user_id (`str`): + mem0 ``user_id`` namespace. + agent_id (`str | None`): + Optional mem0 ``agent_id`` namespace. + """ + if self._await_write: + try: + logger.info( + "mem0 write started: user_id=%s agent_id=%s messages=%d", + user_id, + agent_id, + len(messages), + ) + await self._async_add( + messages, + user_id=user_id, + agent_id=agent_id, + ) + logger.info( + "mem0 write finished for user_id=%s agent_id=%s", + user_id, + agent_id, + ) + except Exception as e: # noqa: BLE001 + logger.warning( + "mem0 add failed for user_id=%s: %s", + user_id, + e, + ) + else: + + async def _bg() -> None: + try: + logger.info( + "mem0 background write started for user_id=%s " + "agent_id=%s messages=%d", + user_id, + agent_id, + len(messages), + ) + await self._async_add( + messages, + user_id=user_id, + agent_id=agent_id, + ) + logger.info( + "mem0 bg write finished: user_id=%s agent_id=%s", + user_id, + agent_id, + ) + except Exception as e: # noqa: BLE001 + logger.warning( + "mem0 background add failed for user_id=%s: %s", + user_id, + e, + ) + + asyncio.create_task(_bg()) diff --git a/src/agentscope/middleware/_longterm_memory/_mem0/_tools.py b/src/agentscope/middleware/_longterm_memory/_mem0/_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..13057baabca53d8c529cf48ab0e259672bdd4c76 --- /dev/null +++ b/src/agentscope/middleware/_longterm_memory/_mem0/_tools.py @@ -0,0 +1,270 @@ +# -*- coding: utf-8 -*- +# pylint: disable=protected-access +"""Agent-control tools exposed by the mem0 middleware. + +These ``search_memory`` / ``add_memory`` tools are listed by +:class:`Mem0Middleware` when ``mode`` is ``"agent_control"`` or +``"both"``. Callers pass them into the agent's toolkit explicitly. +Each tool reads ``user_id`` / ``agent_id`` directly from the +middleware instance — both are plain strings set at construction +time, so no Agent instance is stored or referenced at call time. + +Shape and behavior mirror AgentScope 1.x's +``Mem0LongTermMemory.retrieve_from_memory`` / ``record_to_memory`` +(multi-keyword parallel search, fallback write, verbose result +text) — adapted to AgentScope 2.x's tool conventions +(custom ``ToolBase`` implementations; failures return a ``ToolChunk`` +with ``state=ERROR`` so the toolkit aggregates it properly). +""" +from __future__ import annotations + +import asyncio +from typing import Any, TYPE_CHECKING + +from ....message import TextBlock, ToolResultState +from ....permission import PermissionBehavior, PermissionDecision +from ....tool import ToolBase, ToolChunk + +if TYPE_CHECKING: + from ._middleware import Mem0Middleware + + +class _Mem0MemoryToolBase(ToolBase): + """Base class for mem0 tools that auto-allow themselves. + + Middleware-provided memory tools are part of the agent's standard + capabilities — prompting on every call would defeat the point. + """ + + is_external_tool: bool = False + is_state_injected: bool = False + is_mcp: bool = False + mcp_name: str | None = None + + def __init__( + self, + mw: "Mem0Middleware", + ) -> None: + self._mw = mw + + async def check_permissions( + self, + *_args: Any, + **_kwargs: Any, + ) -> PermissionDecision: + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="auto-allowed: mem0 long-term memory tool", + ) + + +class _SearchMemoryTool(_Mem0MemoryToolBase): + """Agent-callable mem0 search tool.""" + + name: str = "search_memory" + description: str = ( + "Retrieve memories based on short, targeted search keywords. " + "Each keyword is issued as an independent query; results are merged " + "and deduplicated." + ) + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "keywords": { + "type": "array", + "items": {"type": "string"}, + "description": ( + "Short, targeted search phrases such as a person's name, " + "a specific date, a location, or a phrase describing what " + "to retrieve from memory." + ), + }, + "limit": { + "type": "integer", + "description": ( + "Maximum number of memories to retrieve per keyword." + ), + "default": 5, + }, + }, + "required": ["keywords"], + } + is_concurrency_safe: bool = True + is_read_only: bool = True + + async def __call__( + self, + keywords: list[str], + limit: int = 5, + ) -> ToolChunk: + """Retrieve the memory based on the given keywords. + + Args: + keywords (list[str]): + Short, targeted search phrases (for example, a person's + name, a specific date, a location, or a phrase + describing something you want to retrieve from the + memory). Each keyword is issued as an independent query + against the memory store; results are merged and + deduplicated. + limit (int): + The maximum number of memories to retrieve per keyword. + Defaults to 5. + """ + if not keywords: + return _text_chunk("(no keywords supplied — nothing to search)") + + user_id = self._mw._user_id + agent_id = self._mw._agent_id + search_agent_id = agent_id if self._mw._scope_search_by_agent else None + + # Match v1: each keyword is an independent search, run them in + # parallel and merge. + try: + per_keyword = await asyncio.gather( + *[ + self._mw._async_search( + kw, + user_id=user_id, + agent_id=search_agent_id, + top_k=limit, + ) + for kw in keywords + ], + ) + except Exception as e: # noqa: BLE001 + return _error_chunk(f"Error retrieving memory: {e}") + + seen: set[str] = set() + merged: list[str] = [] + for results in per_keyword: + for r in results: + if r not in seen: + seen.add(r) + merged.append(r) + + if not merged: + return _text_chunk("(no relevant memories found)") + return _text_chunk("\n".join(f"- {m}" for m in merged)) + + +class _AddMemoryTool(_Mem0MemoryToolBase): + """Agent-callable mem0 write tool.""" + + name: str = "add_memory" + description: str = ( + "Record important, durable information that may be useful later. " + "Only the provided content is persisted; thinking is retained in the " + "tool result for auditability." + ) + input_schema: dict[str, Any] = { + "type": "object", + "properties": { + "thinking": { + "type": "string", + "description": ( + "Reasoning about why this information is worth " + "remembering. This is not persisted to mem0." + ), + }, + "content": { + "type": "array", + "items": {"type": "string"}, + "description": ( + "Specific facts to remember. Each item should be a " + "complete, standalone sentence." + ), + }, + }, + "required": ["thinking", "content"], + } + is_concurrency_safe: bool = False + is_read_only: bool = False + + async def __call__( + self, + thinking: str, + content: list[str], + ) -> ToolChunk: + """Use this function to record important information that you + may need later. The target content should be specific and + concise, e.g. who, when, where, do what, why, how, etc. + + Do NOT pass back content that appears earlier in the + conversation history as a previous ``search_memory`` tool + result — those facts are already in the store, re-adding them + wastes an extraction call. + + Args: + thinking (str): + Your reasoning about why this is worth remembering. + Stays in the agent transcript but is NOT persisted to + the memory store — only ``content`` is. Use it to + force yourself to think before writing. + content (list[str]): + The content to remember, as a list of strings (one + item per fact). Each item should be a complete, + standalone sentence — only this is sent to mem0 for + extraction. + """ + if not content: + return _error_chunk("`content` is empty — nothing to record.") + + user_id = self._mw._user_id + agent_id = self._mw._agent_id + + # Only the user-facing content goes into mem0. ``thinking`` is + # the agent's internal rationale — meta about the agent's + # decision, not a fact about the user — so feeding it to + # mem0's extraction LLM would muddy the stored memories with + # agent self-narration. We keep it in the tool response so the + # decision is auditable in the transcript. + text = "\n".join(content) + + try: + result = await self._mw._async_add_with_fallback( + text, + user_id=user_id, + agent_id=agent_id, + ) + except Exception as e: # noqa: BLE001 + return _error_chunk(f"Error recording memory: {e}") + + rationale = f" (rationale: {thinking})" if thinking else "" + return _text_chunk( + f"Successfully recorded to memory{rationale} → {result}", + ) + + +def _build_memory_tools( + mw: "Mem0Middleware", +) -> list[ToolBase]: + """Return the ``search_memory`` / ``add_memory`` tools bound to + ``mw``. + + The tool classes intentionally reach into ``mw``'s private state + (``_resolve_user_id`` / ``_async_search`` / ...) — + we live in the same package and ``mw`` is the natural place for + that state. + """ + return [ + _SearchMemoryTool(mw), + _AddMemoryTool(mw), + ] + + +def _text_chunk(message: str) -> ToolChunk: + """Wrap a text message as a normal tool chunk.""" + return ToolChunk( + content=[TextBlock(type="text", text=message)], + ) + + +def _error_chunk(message: str) -> ToolChunk: + """Wrap an error message as a ``ToolChunk(state=ERROR)`` so the + toolkit aggregates it as a failed tool call — the agent sees the + message and can decide whether to retry or move on.""" + return ToolChunk( + content=[TextBlock(type="text", text=message)], + state=ToolResultState.ERROR, + ) diff --git a/src/agentscope/middleware/_longterm_memory/_mem0/_utils.py b/src/agentscope/middleware/_longterm_memory/_mem0/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fb6df96431ae52e248bf636bcc50920f75adab36 --- /dev/null +++ b/src/agentscope/middleware/_longterm_memory/_mem0/_utils.py @@ -0,0 +1,76 @@ +# -*- coding: utf-8 -*- +"""Pure helper functions for the mem0 middleware. + +These are stateless adapters that translate between AgentScope and +mem0 data shapes. Keeping them out of the middleware class makes them +trivial to unit-test in isolation. +""" +from __future__ import annotations + +from typing import Any + +from ....event import ExternalExecutionResultEvent, UserConfirmResultEvent +from ....message import Msg + + +def _extract_query_text(inputs: Any) -> str | None: + """Pull a single text query out of the agent inputs. + + Returns ``None`` for resumption events or empty/non-user inputs, + in which case the middleware skips both retrieval and write-back. + """ + if inputs is None: + return None + if isinstance( + inputs, + (ExternalExecutionResultEvent, UserConfirmResultEvent), + ): + return None + + msgs = inputs if isinstance(inputs, list) else [inputs] + texts: list[str] = [] + for m in msgs: + if not isinstance(m, Msg) or m.role != "user": + continue + text = m.get_text_content() + if text: + texts.append(text) + return "\n".join(texts) if texts else None + + +def _mem0_extracted_anything(raw: Any) -> bool: + """Did mem0's add() actually extract any memories from the input? + + mem0 returns ``{"results": [...]}``; an empty list means its LLM + extractor decided nothing was worth storing. The + ``_async_add_with_fallback`` strategy uses this to decide whether + to try another role / disable inference. + """ + if not isinstance(raw, dict): + return False + results = raw.get("results") + return isinstance(results, list) and len(results) > 0 + + +def _extract_memory_texts(raw: Any) -> list[str]: + """Flatten a mem0 search response into a list of memory strings. + + Tolerates the common shapes: + - ``{"results": [{"memory": str, ...}, ...]}`` (current OSS / Platform) + - ``[{"memory": str, ...}, ...]`` (legacy / variations) + - ``{"results": [str, ...]}`` (defensive fallback) + """ + if raw is None: + return [] + results = raw.get("results", raw) if isinstance(raw, dict) else raw + if not isinstance(results, list): + return [] + out: list[str] = [] + for item in results: + if isinstance(item, str): + out.append(item) + elif isinstance(item, dict): + text = item.get("memory") or item.get("text") + if text: + out.append(str(text)) + return out diff --git a/src/agentscope/middleware/_rag.py b/src/agentscope/middleware/_rag.py new file mode 100644 index 0000000000000000000000000000000000000000..e959ce17f5455b876fa707f219b82f749dace087 --- /dev/null +++ b/src/agentscope/middleware/_rag.py @@ -0,0 +1,765 @@ +# -*- coding: utf-8 -*- +"""RAG middleware that brings knowledge-base search into the agent loop. + +The :class:`RAGMiddleware` wraps one or more +:class:`~agentscope.rag.KnowledgeBase` runtime handles — each carrying +its own embedding model, vector store, and (optional) metadata filter — +so a single agent can search across knowledge bases that were created +with *different* embedding models. + +Two modes are supported, selected via :class:`SearchConfig.mode`: + +- ``"agentic"`` — exposes a single ``search_knowledge`` tool via + :meth:`RAGMiddleware.list_tools`. The agent decides when (and which + knowledge bases) to query. Nothing is injected automatically. +- ``"static"`` — on the first reasoning step of each reply + (``agent.state.cur_iter == 0``) the middleware searches with the + fresh user turn as the query and injects the merged results into + ``agent.state.context`` as a :class:`~agentscope.message.HintBlock`. + Optionally surfaces a + :class:`~agentscope.event.HintBlockEvent` so the front-end can + display the matched snippets. + +User-tunable parameters are declared on :class:`SearchConfig`; the +JSON Schema served to the front-end is derived from it via +``model_json_schema()``. + +Document indexing (parsing, chunking, embedding, insertion) is *not* +this middleware's job — it belongs to the caller (or, in the hosted +service, to the knowledge-base manager) that constructs the +:class:`KnowledgeBase` instances passed in. +""" +import asyncio +import json +from copy import deepcopy +from typing import ( + TYPE_CHECKING, + Any, + AsyncGenerator, + Callable, + Literal, + Sequence, +) + +from pydantic import BaseModel, ConfigDict, Field, field_validator +from pydantic.json_schema import SkipJsonSchema + +from ._base import MiddlewareBase +from .._logging import logger +from ..event import HintBlockEvent +from ..message import ( + DataBlock, + HintBlock, + Msg, + TextBlock, + ToolResultState, +) +from ..permission import PermissionBehavior, PermissionDecision +from ..tool import ParamsBase, ToolBase, ToolChunk + +if TYPE_CHECKING: + from ..agent import Agent + from ..rag import KnowledgeBase, VectorSearchResult + + +_DEFAULT_HINT_TEMPLATE = ( + "The following content is retrieved from the " + "knowledge base(s) and may be helpful for the current " + "request:\n{context}" +) +# Wrapper around the formatted search results. Must contain a single +# ``{context}`` placeholder — :func:`_wrap_hint` splits on it. + +_HINT_SOURCE = json.dumps({"label": "KnowledgeBase", "sublabel": ""}) +# The ``source`` value stamped on injected hint blocks. Encoded as a +# JSON string so the front-end can parse a structured label out of it +# while the field stays a plain ``str`` everywhere else. + + +class _SearchParams(ParamsBase): + """The parameters accepted by ``_SearchKnowledgeTool``.""" + + query: str = Field( + description=( + "The query string to search the knowledge base(s) with. " + "Must be concise, explicit, and self-contained: AVOID " + "ambiguous references like `he`/`she`/`it`, " + "`today`/`yesterday`/`tomorrow`, `here`/`there`, etc. — " + "the search is purely semantic and has no conversational " + "context. Phrase the query as a complete statement of " + "what you want to find." + ), + ) + + knowledge_bases: list[str] | None = Field( + default=None, + description=( + "Optional subset of knowledge bases to query, by name. " + "When omitted (or `null`) every equipped knowledge base " + "is searched. Names must exactly match those listed in " + "the tool description." + ), + ) + + +class _SearchKnowledgeTool(ToolBase): + """The agentic-mode search tool exposed by :class:`RAGMiddleware`. + + A single tool fans out across every bound knowledge base; the + agent may also pass ``knowledge_bases`` to restrict the search to + a subset by name. + """ + + name: str = "search_knowledge" + """The tool name presented to the agent.""" + + is_read_only: bool = True + is_concurrency_safe: bool = True + is_external_tool: bool = False + is_state_injected: bool = False + is_mcp: bool = False + + def __init__( + self, + knowledge_bases: list["KnowledgeBase"], + top_k: int, + score_threshold: float | None, + ) -> None: + """Initialize the search tool. + + Args: + knowledge_bases (`list[KnowledgeBase]`): + The knowledge bases the agent may query. + top_k (`int`): + Maximum number of chunks returned per call, after + merging across knowledge bases. + score_threshold (`float | None`): + Minimum similarity score; forwarded unchanged to each + :meth:`KnowledgeBase.search` call. + """ + # ``ToolBase`` expects a list of *tool* middlewares; this tool + # has none of its own (the owning ``RAGMiddleware`` is an + # *agent* middleware, not a tool one). + super().__init__() + self._knowledge_bases = knowledge_bases + self._top_k = top_k + self._score_threshold = score_threshold + self.description = self._build_description() + self.input_schema = self._build_input_schema() + + def _build_description(self) -> str: + """Build the tool description from the currently equipped + knowledge bases so the agent has enough context to decide which + ones to query.""" + lines = [ + "Search the agent's equipped knowledge bases by semantic " + "similarity and return the most relevant chunks.", + "", + "## When to Use", + "- The user's question may be answered by content stored " + "in one of the listed knowledge bases (see *Equipped " + "Knowledge Bases* below).", + "- You need supporting facts, definitions, or documents " + "that are unlikely to be in your parametric knowledge.", + "", + "## Guidance", + "- Knowledge base names and descriptions are " + "user-supplied and may be terse, vague, or unrelated to " + "the actual contents. When in doubt, try the search — a " + "single call is cheap and an empty result is informative.", + "- Phrase `query` as a self-contained statement of what " + "you want to find. Avoid pronouns or relative time " + "references — the search is purely semantic and has no " + "conversational context.", + "- Set `knowledge_bases` only when the question clearly " + "matches one or two specific bases; otherwise leave it " + "unset to search them all.", + "", + "## Equipped Knowledge Bases", + ] + if self._knowledge_bases: + lines.append( + f"The agent is currently equipped with " + f"{len(self._knowledge_bases)} knowledge base(s):", + ) + lines.extend( + f"- **{kb.name}**: {kb.description}" + for kb in self._knowledge_bases + ) + else: + lines.append( + "No knowledge bases are currently equipped. **Do " + "not call this tool** — it will return nothing.", + ) + return "\n".join(lines) + + def _build_input_schema(self) -> dict: + """Build the JSON Schema, with ``knowledge_bases.items.enum`` + narrowed to the currently equipped knowledge-base names so + the LLM cannot invent unknown names.""" + schema: dict[str, Any] = _SearchParams.model_json_schema() + if self._knowledge_bases: + names = [kb.name for kb in self._knowledge_bases] + kb_schema = schema["properties"]["knowledge_bases"] + # Pydantic emits ``Optional[list[str]]`` as + # ``anyOf: [{type: array, items: ...}, {type: null}]``; we + # narrow the array branch's items. Fall back to a direct + # ``items`` field for the non-optional shape. + if "items" in kb_schema: + kb_schema["items"]["enum"] = names + else: + for variant in kb_schema.get("anyOf", []): + if variant.get("type") == "array" and "items" in variant: + variant["items"]["enum"] = names + break + return schema + + async def check_permissions( + self, + tool_input: dict[str, Any], + context: Any, + ) -> Any: + """Allow the engine to handle this read-only search. + + Args: + tool_input (`dict[str, Any]`): + The tool input data. + context (`PermissionContext`): + The permission context. + + Returns: + `PermissionDecision`: + An allow decision — knowledge-base search is read-only. + """ + del tool_input, context + return PermissionDecision( + behavior=PermissionBehavior.ALLOW, + message="Knowledge-base search is read-only.", + ) + + async def call( # type: ignore[override] + self, + query: str, + knowledge_bases: list[str] | None = None, + ) -> ToolChunk: + """Search the selected knowledge bases and return the results + as content blocks. + + Args: + query (`str`): + The natural-language search query. + knowledge_bases (`list[str] | None`, optional): + Optional subset of knowledge bases to query, by name. + ``None`` searches every equipped knowledge base. + + Returns: + `ToolChunk`: + The formatted search results, or a notice when + nothing relevant is found. + """ + if knowledge_bases is None: + targets = list(self._knowledge_bases) + else: + wanted = set(knowledge_bases) + targets = [kb for kb in self._knowledge_bases if kb.name in wanted] + + if not targets: + return ToolChunk( + content=[TextBlock(text="No relevant content found.")], + state=ToolResultState.SUCCESS, + is_last=True, + ) + + try: + results = await _search_across( + targets, + [query], + top_k=self._top_k, + score_threshold=self._score_threshold, + ) + except Exception as e: # pylint: disable=broad-except + logger.exception("search_knowledge failed.") + return ToolChunk( + content=[TextBlock(text=f"Search failed: {e}")], + state=ToolResultState.ERROR, + is_last=True, + ) + + blocks = _format_results(results) + if not blocks: + return ToolChunk( + content=[TextBlock(text="No relevant content found.")], + state=ToolResultState.SUCCESS, + is_last=True, + ) + return ToolChunk( + content=blocks, + state=ToolResultState.SUCCESS, + is_last=True, + ) + + +# --------------------------------------------------------------------- +# Shared helpers — used by both the tool (agentic mode) and the +# middleware's static-mode injection path. +# --------------------------------------------------------------------- + + +async def _search_across( + knowledge_bases: Sequence["KnowledgeBase"], + queries: Sequence[str | TextBlock | DataBlock], + top_k: int, + score_threshold: float | None, +) -> list["VectorSearchResult"]: + """Search every knowledge base concurrently and merge the results. + + Each knowledge base handles its own filtering — it silently drops + :class:`DataBlock` inputs when its embedding model is not + multimodal — so callers can pass the same query list to every + knowledge base without per-KB pre-filtering. Per-KB hits are + flattened, sorted by descending score, and truncated to ``top_k``. + + .. note:: + Scores from knowledge bases with different embedding models + are not strictly comparable; this merge sorts by raw score. + For mixed-embedding deployments where that matters, switch to + a rank-based fusion (e.g. RRF) — each per-KB + :meth:`KnowledgeBase.search` already returns ordered results. + + Args: + knowledge_bases (`list[KnowledgeBase]`): + The knowledge bases to query. + queries (`list[str | TextBlock | DataBlock]`): + The query inputs. + top_k (`int`): + Maximum number of results after merging. + score_threshold (`float | None`): + Forwarded to each :meth:`KnowledgeBase.search` call. + + Returns: + `list[VectorSearchResult]`: + At most ``top_k`` hits across all knowledge bases. + """ + if not queries or not knowledge_bases: + return [] + + queries_list = list(queries) + per_kb = await asyncio.gather( + *( + kb.search( + queries=queries_list, + top_k=top_k, + score_threshold=score_threshold, + ) + for kb in knowledge_bases + ), + ) + + merged = [r for sub in per_kb for r in sub] + merged.sort(key=lambda r: r.score, reverse=True) + return merged[:top_k] + + +def _format_results( + results: list["VectorSearchResult"], +) -> list[TextBlock | DataBlock]: + """Render search results as a numbered, cited list of blocks. + + Every result becomes ``[N] (source: ...)`` followed by its chunk + content (text inlined, multimodal blocks kept as standalone + :class:`DataBlock` entries). Adjacent text fragments are merged + into a single :class:`TextBlock` so downstream consumers see one + contiguous text block instead of many small ones. + + Args: + results (`list[VectorSearchResult]`): + The search results to format. + + Returns: + `list[TextBlock | DataBlock]`: + The formatted blocks; empty list when ``results`` is + empty. + """ + entries: list[TextBlock | DataBlock] = [] + last = len(results) + for index, result in enumerate(results, start=1): + prefix = f"[{index}] (source: {result.chunk.source})\n" + block = deepcopy(result.chunk.content) + if isinstance(block, TextBlock): + block.text = prefix + block.text + entries.append(block) + else: + entries.append(TextBlock(text=prefix)) + entries.append(block) + if index != last: + entries.append(TextBlock(text="\n\n")) + + # Coalesce adjacent TextBlocks into one to keep the consumer-facing + # block list tight. We build fresh TextBlocks rather than mutating + # the deep copied ones to avoid aliasing the originals. + merged: list[TextBlock | DataBlock] = [] + for entry in entries: + if ( + isinstance(entry, TextBlock) + and merged + and isinstance(merged[-1], TextBlock) + ): + merged[-1] = TextBlock(text=merged[-1].text + entry.text) + else: + merged.append(entry) + return merged + + +def _wrap_hint( + template: str, + blocks: list[TextBlock | DataBlock], +) -> str | list[TextBlock | DataBlock]: + """Substitute ``{context}`` in ``template`` with the rendered blocks. + + When every block is a :class:`TextBlock` the result is a single + ``str`` (cheap, no block overhead). When the blocks include + multimodal :class:`DataBlock` content the template is split on + ``{context}`` and the surrounding text is prepended / appended as + :class:`TextBlock` items, producing a ``list[TextBlock | + DataBlock]`` that preserves the binary payloads end-to-end. + + Args: + template (`str`): + Wrapper template with a single ``{context}`` placeholder. + blocks (`list[TextBlock | DataBlock]`): + The formatted search blocks to wrap. + + Returns: + `str | list[TextBlock | DataBlock]`: + A plain string for text-only payloads, or a list of blocks + when binary payloads are involved. + """ + if all(isinstance(b, TextBlock) for b in blocks): + joined = "\n".join(b.text for b in blocks) # type: ignore[union-attr] + return template.format(context=joined) + + prefix, _, end = template.partition("{context}") + wrapped: list[TextBlock | DataBlock] = list(blocks) + if prefix: + if isinstance(wrapped[0], TextBlock): + wrapped[0] = TextBlock(text=prefix + wrapped[0].text) + else: + wrapped.insert(0, TextBlock(text=prefix)) + if end: + if isinstance(wrapped[-1], TextBlock): + wrapped[-1] = TextBlock(text=wrapped[-1].text + end) + else: + wrapped.append(TextBlock(text=end)) + return wrapped + + +class RAGMiddleware(MiddlewareBase): + """Middleware that integrates knowledge-base search into the agent. + + Constructed from a list of :class:`~agentscope.rag.KnowledgeBase` + handles — each carrying its own embedding model, vector store, and + metadata filter. The middleware does not own these resources; it + only orchestrates search against them. + + .. code-block:: python + + # Automatic injection (static mode) + middleware = RAGMiddleware( + knowledge_bases=[kb1, kb2], + search_config=SearchConfig(mode="static"), + ) + + # Agent-driven search (agentic mode, the default) + middleware = RAGMiddleware(knowledge_bases=[kb1, kb2]) + + agent = Agent(..., middlewares=[middleware], ...) + """ + + class Parameters(BaseModel): + """User-tunable knowledge-base search parameters of + :class:`RAGMiddleware`. + + The fields here are exactly the keys the hosted service persists + into ``SessionKnowledgeConfig.parameters`` and the keys + :class:`RAGMiddleware` accepts as ``search_config``. Every field + is annotated with a ``title`` and ``description`` so the front-end + can render them as labels and tooltips via + ``model_json_schema()``. + """ + + model_config = ConfigDict(frozen=True) + + mode: Literal["static", "agentic"] = Field( + default="agentic", + title="Mode", + description=( + "Retrieval is either agentic, letting the Agent decide when " + "to retrieve, or static, triggering on every user input." + ), + ) + + top_k: int = Field( + default=5, + ge=1, + le=50, + title="Top K", + description=( + "Maximum number of chunks returned per search, across all " + "configured knowledge bases." + ), + ) + + score_threshold: float | None = Field( + default=None, + title="Score Threshold", + description=( + "Minimum similarity score for a hit to be kept. Leave " + "empty to disable filtering." + ), + ) + + emit_hint_event: bool = Field( + default=True, + title="Show matched chunks in chat", + description=( + "Emit a `HintBlockEvent` in static mode so the front-end " + "can display the matched snippets to the user." + ), + ) + + persist_hint: bool = Field( + default=False, + title="Persist Hint", + description=( + "In `static` mode, keep the injected hint block in the " + "agent context instead of removing it right after the " + "model call." + ), + ) + + hint_template: SkipJsonSchema[str] = Field( + default=_DEFAULT_HINT_TEMPLATE, + title="Hint template", + description=( + "Template wrapping the formatted search results in static " + "mode, with a `{context}` placeholder." + ), + ) + + # ``hint_template`` is intentionally hidden from the JSON Schema + # exposed to the dock UI: the wrapper text is part of the + # middleware's prompt contract and exposing it through the dock + # invites session-by-session prompt drift. It is still accepted + # for programmatic use. + + @field_validator("hint_template") + @classmethod + def _validate_hint_template(cls, value: str) -> str: + """Reject templates with anything other than exactly one + ``{context}`` placeholder — :func:`_wrap_hint` substitutes on + the first occurrence, so zero placeholders silently drop the + matched content and multiple placeholders duplicate it.""" + count = value.count("{context}") + if count != 1: + raise ValueError( + "hint_template must contain exactly one '{context}' " + f"placeholder; found {count}.", + ) + return value + + def __init__( + self, + knowledge_bases: list["KnowledgeBase"], + parameters: "RAGMiddleware.Parameters | None" = None, + ) -> None: + """Initialize the RAG middleware. + + Args: + knowledge_bases (`list[KnowledgeBase]`): + The knowledge bases this agent searches. + parameters (`RAGMiddleware.Parameters | None`, optional): + Search-time knobs (mode, top_k, score threshold, hint + behaviour). ``None`` uses the defaults of + :class:`SearchConfig`. + """ + self._knowledge_bases = knowledge_bases + self._parameters = parameters or RAGMiddleware.Parameters() + # Static-mode reply scratchpad: populated by ``on_reply`` and + # consumed by ``on_reasoning`` so the auto-search can use the + # original reply inputs (which are no longer in + # ``agent.state.context`` by the time ``on_reasoning`` runs in + # a tool-call loop). Cleared in ``on_reply``'s finally. + self._cached_inputs: list[TextBlock | DataBlock] | None = None + + # ------------------------------------------------------------------ + # Agentic mode — expose the search tool + # ------------------------------------------------------------------ + + async def list_tools(self) -> list[ToolBase]: + """Expose the search tool in ``"agentic"`` mode. + + Returns: + `list[ToolBase]`: + A single ``search_knowledge`` tool in ``"agentic"`` + mode; an empty list in ``"static"`` mode. + """ + if self._parameters.mode == "agentic": + return [ + _SearchKnowledgeTool( + self._knowledge_bases, + top_k=self._parameters.top_k, + score_threshold=self._parameters.score_threshold, + ), + ] + return [] + + # ------------------------------------------------------------------ + # Static mode — capture inputs in on_reply, search in on_reasoning + # ------------------------------------------------------------------ + + async def on_reply( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Cache reply inputs for the static-mode search. + + ``on_reasoning`` runs *after* the agent has potentially + consumed the inputs and started its own reply, so it cannot + recover them from ``agent.state.context`` reliably. This + method captures them on entry and clears the cache when the + reply finishes. + + Args: + agent (`Agent`): + The executing agent. Unused, but part of the + middleware contract. + input_kwargs (`dict`): + Reply input kwargs (``inputs``, plus any extras); + forwarded unchanged. + next_handler (`Callable[..., AsyncGenerator]`): + The downstream middleware or core reply logic. + + Yields: + `Any`: + Whatever ``next_handler`` yields. + """ + inputs = input_kwargs.get("inputs") + + msgs: list[Msg] | None = None + if isinstance(inputs, Msg): + msgs = [inputs] + elif isinstance(inputs, list) and all( + isinstance(m, Msg) for m in inputs + ): + msgs = inputs + + if msgs: + # Deepcopy because we are about to mutate the first text block of + # each message to prepend the speaker name — never touch the + # caller's message objects. + msgs = deepcopy(msgs) + blocks: list[TextBlock | DataBlock] = [] + for msg in msgs: + if not msg.content: + continue + speaker = f"{msg.name}: " + if isinstance(msg.content[0], TextBlock): + msg.content[0].text = speaker + msg.content[0].text + else: + blocks.append(TextBlock(text=speaker)) + blocks.extend(msg.content) + + self._cached_inputs = blocks + + try: + async for evt in next_handler(**input_kwargs): + yield evt + finally: + self._cached_inputs = None + + async def on_reasoning( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Inject a one-shot RAG hint on the first reasoning step. + + Only active in ``"static"`` mode and only when + ``agent.state.cur_iter == 0`` — i.e. the first reasoning cycle + of a reply. Subsequent reasoning iterations (tool-call + rounds) skip the search so the agent does not re-embed and + re-inject for every iteration. + + When ``persist_hint`` is ``False`` (the default) the injected + block is removed from the context right after the reasoning + step it participated in — keyed on the block's id so other + middlewares can append their own blocks to the same carrier + message without interfering. + + Args: + agent (`Agent`): + The executing agent whose ``state.context`` receives + the hint. + input_kwargs (`dict`): + Reasoning input kwargs; forwarded unchanged. + next_handler (`Callable[..., AsyncGenerator]`): + The downstream middleware or core reasoning logic. + + Yields: + `Any`: + An optional :class:`HintBlockEvent` followed by events + from downstream. + """ + hint: HintBlock | None = None + + if ( + self._parameters.mode == "static" + and agent.state.cur_iter == 0 + and self._cached_inputs + ): + try: + results = await _search_across( + self._knowledge_bases, + self._cached_inputs, + top_k=self._parameters.top_k, + score_threshold=self._parameters.score_threshold, + ) + except Exception: # pylint: disable=broad-except + logger.exception( + "Knowledge-base search failed; proceeding without " + "matched context.", + ) + results = [] + + blocks = _format_results(results) + if blocks: + hint = HintBlock( + hint=_wrap_hint(self._parameters.hint_template, blocks), + source=_HINT_SOURCE, + ) + agent.state.append_context(agent.name, [hint]) + if self._parameters.emit_hint_event: + yield HintBlockEvent( + reply_id=agent.state.reply_id, + block_id=hint.id, + source=hint.source, + hint=hint.hint, + ) + + try: + async for evt in next_handler(**input_kwargs): + yield evt + finally: + if hint is not None and not self._parameters.persist_hint: + # Remove the injected block from whichever message + # ``append_context`` placed it on. Reverse-scan + # because that carrier is always the latest message + # with our ``reply_id``. + for msg in reversed(agent.state.context): + if msg.id != agent.state.reply_id: + continue + msg.content = [b for b in msg.content if b.id != hint.id] + break diff --git a/src/agentscope/middleware/_tracing/__init__.py b/src/agentscope/middleware/_tracing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..15c753b228cafb1df6a1124bf2cee7766255c096 --- /dev/null +++ b/src/agentscope/middleware/_tracing/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- +"""The tracing interface class in agentscope.""" + +from ._trace import TracingMiddleware + +__all__ = [ + "TracingMiddleware", +] diff --git a/src/agentscope/middleware/_tracing/_attributes.py b/src/agentscope/middleware/_tracing/_attributes.py new file mode 100644 index 0000000000000000000000000000000000000000..cba959200d3b3006d2631bbffc02425e24217250 --- /dev/null +++ b/src/agentscope/middleware/_tracing/_attributes.py @@ -0,0 +1,201 @@ +# -*- coding: utf-8 -*- +"""The tracing types class in agentscope.""" +from opentelemetry.semconv._incubating.attributes import ( + gen_ai_attributes as GenAIAttributes, +) + + +class SpanAttributes: + """The span attributes.""" + + # GenAI Common Attributes + GEN_AI_CONVERSATION_ID = GenAIAttributes.GEN_AI_CONVERSATION_ID + """The gen ai conversation ID.""" + + GEN_AI_OPERATION_NAME = GenAIAttributes.GEN_AI_OPERATION_NAME + """The gen ai operation name.""" + + GEN_AI_PROVIDER_NAME = GenAIAttributes.GEN_AI_PROVIDER_NAME + """The gen ai provider name.""" + + # GenAI Request Attributes + GEN_AI_REQUEST_MODEL = GenAIAttributes.GEN_AI_REQUEST_MODEL + """The gen ai request model.""" + + GEN_AI_REQUEST_TEMPERATURE = GenAIAttributes.GEN_AI_REQUEST_TEMPERATURE + """The gen ai request temperature.""" + + GEN_AI_REQUEST_TOP_P = GenAIAttributes.GEN_AI_REQUEST_TOP_P + """The gen ai request top_p.""" + + GEN_AI_REQUEST_TOP_K = GenAIAttributes.GEN_AI_REQUEST_TOP_K + """The gen ai request top_k.""" + + GEN_AI_REQUEST_MAX_TOKENS = GenAIAttributes.GEN_AI_REQUEST_MAX_TOKENS + """The gen ai request max_tokens.""" + + GEN_AI_REQUEST_PRESENCE_PENALTY = ( + GenAIAttributes.GEN_AI_REQUEST_PRESENCE_PENALTY + ) + """The gen ai request presence_penalty.""" + + GEN_AI_REQUEST_FREQUENCY_PENALTY = ( + GenAIAttributes.GEN_AI_REQUEST_FREQUENCY_PENALTY + ) + """The gen ai request frequency_penalty.""" + + GEN_AI_REQUEST_STOP_SEQUENCES = ( + GenAIAttributes.GEN_AI_REQUEST_STOP_SEQUENCES + ) + """The gen ai request stop_sequences.""" + + GEN_AI_REQUEST_SEED = GenAIAttributes.GEN_AI_REQUEST_SEED + """The gen ai request seed.""" + + # GenAI Response Attributes + GEN_AI_RESPONSE_ID = GenAIAttributes.GEN_AI_RESPONSE_ID + """The gen ai response ID.""" + + GEN_AI_RESPONSE_FINISH_REASONS = ( + GenAIAttributes.GEN_AI_RESPONSE_FINISH_REASONS + ) + """The gen ai response finish reasons.""" + + # GenAI Usage Attributes + GEN_AI_USAGE_INPUT_TOKENS = GenAIAttributes.GEN_AI_USAGE_INPUT_TOKENS + """The gen ai usage input tokens.""" + + GEN_AI_USAGE_OUTPUT_TOKENS = GenAIAttributes.GEN_AI_USAGE_OUTPUT_TOKENS + """The gen ai usage output tokens.""" + + # GenAI Message Attributes + GEN_AI_INPUT_MESSAGES = GenAIAttributes.GEN_AI_INPUT_MESSAGES + """The gen ai input messages.""" + + GEN_AI_OUTPUT_MESSAGES = GenAIAttributes.GEN_AI_OUTPUT_MESSAGES + """The gen ai output messages.""" + + # GenAI Agent Attributes + GEN_AI_AGENT_ID = GenAIAttributes.GEN_AI_AGENT_ID + """The gen ai agent ID.""" + + GEN_AI_AGENT_NAME = GenAIAttributes.GEN_AI_AGENT_NAME + """The gen ai agent name.""" + + GEN_AI_AGENT_DESCRIPTION = GenAIAttributes.GEN_AI_AGENT_DESCRIPTION + """The gen ai agent description.""" + + # GenAI Tool Attributes + GEN_AI_TOOL_CALL_ID = GenAIAttributes.GEN_AI_TOOL_CALL_ID + """The gen ai tool call ID.""" + + GEN_AI_TOOL_NAME = GenAIAttributes.GEN_AI_TOOL_NAME + """The gen ai tool name.""" + + GEN_AI_TOOL_DESCRIPTION = GenAIAttributes.GEN_AI_TOOL_DESCRIPTION + """The gen ai tool description.""" + + GEN_AI_TOOL_CALL_ARGUMENTS = GenAIAttributes.GEN_AI_TOOL_CALL_ARGUMENTS + """The gen ai tool call arguments.""" + + GEN_AI_TOOL_CALL_RESULT = GenAIAttributes.GEN_AI_TOOL_CALL_RESULT + """The gen ai tool call result.""" + + GEN_AI_TOOL_DEFINITIONS = GenAIAttributes.GEN_AI_TOOL_DEFINITIONS + """The gen ai tool definitions.""" + + AGENTSCOPE_CACHE_INPUT_TOKENS = "agentscope.usage.cache_input_tokens" + """The number of input tokens read from prompt cache.""" + + AGENTSCOPE_CACHE_CREATION_INPUT_TOKENS = ( + "agentscope.usage.cache_creation_input_tokens" + ) + """The number of input tokens used to create prompt cache.""" + + AGENTSCOPE_REPLY_ID = "agentscope.agent.reply_id" + """The reply ID of the current agent reply. + + Shared by both calls in a HITL or external-execution chain, allowing + observers to group the two ``invoke_agent`` spans that belong to the + same logical reply. + """ + + AGENTSCOPE_HITL_PENDING_TOOLS = "agentscope.agent.hitl_pending_tools" + """JSON list of tool names that are waiting for human confirmation. + + Set on the first ``invoke_agent`` span when the agent pauses due to a + ``RequireUserConfirmEvent``. + """ + + AGENTSCOPE_EXTERNAL_EXECUTION_PENDING_TOOLS = ( + "agentscope.agent.external_execution_pending_tools" + ) + """JSON list of tool names submitted for external execution. + + Set on the first ``invoke_agent`` span when the agent pauses due to a + ``RequireExternalExecutionEvent``. + """ + + AGENTSCOPE_INCOMING_EVENT_TYPE = "agentscope.agent.incoming_event_type" + """Type of the continuation event passed to the second reply call. + + Possible values: ``"user_confirm_result"``, + ``"external_execution_result"``. + """ + + AGENTSCOPE_IS_EXTERNAL_EXECUTION = "agentscope.agent.is_external_execution" + """Marks a synthetic ``execute_tool`` span that represents a tool executed + externally (i.e. via ``ExternalExecutionResultEvent``). + + The timestamp reflects when the result was received, not when the tool + actually ran on the external system. + """ + + +class OperationNameValues: + """The operation name values.""" + + CHAT = GenAIAttributes.GenAiOperationNameValues.CHAT.value + """The chat operation name.""" + + INVOKE_AGENT = GenAIAttributes.GenAiOperationNameValues.INVOKE_AGENT.value + """The invoke agent operation name.""" + + EXECUTE_TOOL = GenAIAttributes.GenAiOperationNameValues.EXECUTE_TOOL.value + """The execute tool operation name.""" + + +class ProviderNameValues: + """The provider name values.""" + + DASHSCOPE = "dashscope" + """The dashscope provider name.""" + + OLLAMA = "ollama" + """The ollama provider name.""" + + DEEPSEEK = GenAIAttributes.GenAiProviderNameValues.DEEPSEEK.value + """The deepseek provider name.""" + + OPENAI = GenAIAttributes.GenAiProviderNameValues.OPENAI.value + """The openai provider name.""" + + ANTHROPIC = GenAIAttributes.GenAiProviderNameValues.ANTHROPIC.value + """The anthropic provider name.""" + + GCP_GEMINI = GenAIAttributes.GenAiProviderNameValues.GCP_GEMINI.value + """The gcp gemini provider name.""" + + MOONSHOT = "moonshot" + """The moonshot provider name.""" + + AZURE_AI_OPENAI = ( + GenAIAttributes.GenAiProviderNameValues.AZURE_AI_OPENAI.value + ) + """The azure openai provider name.""" + + AWS_BEDROCK = GenAIAttributes.GenAiProviderNameValues.AWS_BEDROCK.value + """The aws bedrock provider name.""" + + XAI = GenAIAttributes.GenAiProviderNameValues.X_AI.value + """The xAI (Grok) provider name.""" diff --git a/src/agentscope/middleware/_tracing/_converter.py b/src/agentscope/middleware/_tracing/_converter.py new file mode 100644 index 0000000000000000000000000000000000000000..9979c870ca377f33e12b66c0019e219527b429e1 --- /dev/null +++ b/src/agentscope/middleware/_tracing/_converter.py @@ -0,0 +1,121 @@ +# -*- coding: utf-8 -*- +"""Convert ContentBlock to OpenTelemetry GenAI part format.""" + +import json +from typing import Any, Dict + +from ...message import ( + ContentBlock, + TextBlock, + ThinkingBlock, + ToolCallBlock, + ToolResultBlock, + DataBlock, + Base64Source, + URLSource, +) + +from ._utils import _serialize_to_str + +# Valid OTel GenAI modality values for multimodal content parts. +_VALID_MODALITIES = frozenset({"image", "audio", "video"}) + + +def _get_modality(media_type: str) -> str: + """Derive OTel modality from a MIME type string (e.g. 'image/png').""" + prefix = media_type.split("/")[0] if media_type else "" + return prefix if prefix in _VALID_MODALITIES else "unknown" + + +def _convert_media_block( + source: Base64Source | URLSource, +) -> Dict[str, Any] | None: + """Convert a DataBlock source to OpenTelemetry GenAI part format. + + Args: + source (`Base64Source | URLSource`): + The data source of the DataBlock. + + Returns: + `Dict[str, Any] | None`: + Converted part Dictionary or None if the source type is invalid. + """ + modality = _get_modality(source.media_type) + + if isinstance(source, URLSource): + return { + "type": "uri", + "uri": str(source.url), + "modality": modality, + } + + if isinstance(source, Base64Source): + return { + "type": "blob", + "content": source.data, + "media_type": source.media_type, + "modality": modality, + } + + return None + + +def _convert_block_to_part(block: ContentBlock) -> Dict[str, Any] | None: + """Convert content block to OpenTelemetry GenAI part format. + + Converts text, thinking, tool_call, tool_result and data (media) blocks + to standardized parts. + + Args: + block (`ContentBlock`): + The content block object to convert. Supported block types: + - text: Text content block + - thinking: Reasoning/thinking content block + - tool_call: Tool call block with id, name, and input + - tool_result: Tool result block with id and output + - data: Binary data block (image, audio, video, etc.) + + Returns: + `Dict[str, Any] | None`: + Standardized part Dictionary in OpenTelemetry GenAI format, + or None if the block type is unsupported or cannot be converted. + """ + part: Dict[str, Any] | None = None + + if isinstance(block, TextBlock): + part = { + "type": "text", + "content": block.text, + } + elif isinstance(block, ThinkingBlock): + part = { + "type": "reasoning", + "content": block.thinking, + } + elif isinstance(block, ToolCallBlock): + try: + arguments = json.loads(block.input) + except (json.JSONDecodeError, TypeError): + arguments = block.input + part = { + "type": "tool_call", + "id": block.id, + "name": block.name, + "arguments": arguments, + } + elif isinstance(block, ToolResultBlock): + output = block.output + if isinstance(output, (list, dict)): + result = _serialize_to_str(output) + else: + result = str(output) + + part = { + "type": "tool_call_response", + "id": block.id, + "response": result, + } + elif isinstance(block, DataBlock): + part = _convert_media_block(block.source) + + return part diff --git a/src/agentscope/middleware/_tracing/_extractor.py b/src/agentscope/middleware/_tracing/_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..6f3bf7ad8142e96fd561a3a6a0a9cbed54bb734a --- /dev/null +++ b/src/agentscope/middleware/_tracing/_extractor.py @@ -0,0 +1,608 @@ +# -*- coding: utf-8 -*- +"""Extract attributes from AgentScope components for OpenTelemetry tracing.""" +import inspect +from typing import Any, Dict, TYPE_CHECKING + +from ...message import Msg, ToolCallBlock + +from ._attributes import ( + SpanAttributes, + OperationNameValues, + ProviderNameValues, +) +from ._converter import _convert_block_to_part +from ._utils import _serialize_to_str +from ...model import ChatResponse, ChatModelBase +from ...event import ( + ExternalExecutionResultEvent, + UserConfirmResultEvent, +) + +if TYPE_CHECKING: + from ...agent import Agent + from ...tool import Toolkit, ToolChoice + +_CLASS_NAME_MAP = { + "dashscope": ProviderNameValues.DASHSCOPE, + "openai": ProviderNameValues.OPENAI, + "anthropic": ProviderNameValues.ANTHROPIC, + "gemini": ProviderNameValues.GCP_GEMINI, + "ollama": ProviderNameValues.OLLAMA, + "deepseek": ProviderNameValues.DEEPSEEK, + "xai": ProviderNameValues.XAI, + "moonshot": ProviderNameValues.MOONSHOT, +} + +# Map base URL fragments to provider names for OpenAI-compatible APIs +_BASE_URL_PROVIDER_MAP = [ + ("api.openai.com", ProviderNameValues.OPENAI), + ("dashscope", ProviderNameValues.DASHSCOPE), + ("deepseek", ProviderNameValues.DEEPSEEK), + ("moonshot", ProviderNameValues.MOONSHOT), + ("generativelanguage.googleapis.com", ProviderNameValues.GCP_GEMINI), + ("openai.azure.com", ProviderNameValues.AZURE_AI_OPENAI), + ("amazonaws.com", ProviderNameValues.AWS_BEDROCK), + ("api.x.ai", ProviderNameValues.XAI), +] + + +def _get_common_attributes(session_id: str = "") -> Dict[str, str]: + """Get common attributes for all spans. + + Args: + session_id (`str`): + The session ID to set as conversation ID. + + Returns: + `Dict[str, str]`: + Common span attributes including conversation ID. + """ + return { + SpanAttributes.GEN_AI_CONVERSATION_ID: ( + session_id if session_id else "[no_session_id]" + ), + } + + +def _get_provider_name(instance: "ChatModelBase") -> str: + """Get provider name from ChatModelBase instance. + + Maps ChatModelBase class names to provider names, with special handling + for OpenAI-compatible APIs that may use different base URLs. + This follows the implementation pattern from agentscope-java PR #73. + + Args: + instance (`ChatModelBase`): + The chat model instance to get the provider name for. + + Returns: + `str`: + Provider name (e.g., "openai", "dashscope", "anthropic") + """ + classname = instance.__class__.__name__ + + # For other model types, use direct mapping + prefix_key = ( + classname.removesuffix("ChatModel") + .removesuffix("MultiAgentModel") + .removesuffix("ResponseModel") + .lower() + ) + + # Special handling for OpenAI-compatible models — inspect base_url + # from credential to distinguish the actual provider. + if prefix_key == "openai": + base_url = getattr(instance.credential, "base_url", None) + if base_url: + base_url = str(base_url) + for url_fragment, provider_name in _BASE_URL_PROVIDER_MAP: + if url_fragment in base_url: + return provider_name + return ProviderNameValues.OPENAI + + return _CLASS_NAME_MAP.get(prefix_key, "unknown") + + +def _get_tool_definitions( + tools: list[dict[str, Any]] | None, + tool_choice: "ToolChoice | None", +) -> str | None: + """Extract and serialize tool definitions for tracing. + + Converts AgentScope/OpenAI nested tool format to OpenTelemetry GenAI + flat format for tracing. + + Args: + tools (`list[dict[str, Any]] | None`, optional): + List of tool definitions in OpenAI format with nested + structure: ``[{"type": "function", "function": {...}}]`` + tool_choice (`ToolChoice | None`, optional): + Tool choice configuration with ``mode`` and optional ``tools`` + fields. If mode is ``"none"``, returns None to indicate tools + should not be traced. + + Returns: + `str | None`: + Serialized tool definitions in flat format: + ``[{"type": "function", "name": ..., "parameters": ...}]`` + or None if tools should not be traced (e.g., tools is None/empty + or tool_choice is "none"). + """ + # No tools provided + if tools is None or not isinstance(tools, list) or len(tools) == 0: + return None + + # Tool choice is explicitly "none" (model should not use tools) + if tool_choice is not None and tool_choice.mode == "none": + return None + + try: + # Convert nested format to flat format for OpenTelemetry GenAI + # TODO: Currently only supports "function" type tools. If other tool + # types are added in the future (e.g., "retrieval", "code_interpreter", + # "browser"), this conversion logic needs to be updated to handle them. + flat_tools = [] + for tool in tools: + if not isinstance(tool, dict) or "function" not in tool: + continue + + func_def = tool["function"] + flat_tool = { + "type": tool.get("type", "function"), + "name": func_def.get("name"), + "description": func_def.get("description"), + "parameters": func_def.get("parameters"), + } + # Remove None values + flat_tool = {k: v for k, v in flat_tool.items() if v is not None} + flat_tools.append(flat_tool) + + if flat_tools: + return _serialize_to_str(flat_tools) + return None + + except Exception: + return None + + +def _get_llm_request_attributes( + instance: "ChatModelBase", + kwargs: Dict[str, Any], +) -> Dict[str, Any]: + """Get LLM request attributes for OpenTelemetry tracing. + + Extracts request parameters from LLM model calls into GenAI attributes. + + Args: + instance (`ChatModelBase`): + The chat model instance making the request. + kwargs (`Dict[str, Any]`): + Keyword arguments including generation parameters such as + temperature, top_p, top_k, max_tokens, presence_penalty, + frequency_penalty, stop_sequences, seed, tools, and tool_choice. + + Returns: + `Dict[str, Any]`: + OpenTelemetry GenAI attributes with mixed-type values (``str``, + ``int``, ``float``, or ``list``), including operation name, + provider name, model name, generation parameters (e.g. + temperature, max_tokens, stop_sequences), and tool definitions. + """ + + attributes = { + # required attributes + SpanAttributes.GEN_AI_OPERATION_NAME: OperationNameValues.CHAT, + SpanAttributes.GEN_AI_PROVIDER_NAME: _get_provider_name(instance), + # conditionally required attributes + SpanAttributes.GEN_AI_REQUEST_MODEL: instance.model, + # recommended attributes + SpanAttributes.GEN_AI_REQUEST_TEMPERATURE: kwargs.get("temperature"), + SpanAttributes.GEN_AI_REQUEST_TOP_P: kwargs.get("p") + or kwargs.get("top_p"), + SpanAttributes.GEN_AI_REQUEST_TOP_K: kwargs.get("top_k"), + SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS: kwargs.get("max_tokens"), + SpanAttributes.GEN_AI_REQUEST_PRESENCE_PENALTY: kwargs.get( + "presence_penalty", + ), + SpanAttributes.GEN_AI_REQUEST_FREQUENCY_PENALTY: kwargs.get( + "frequency_penalty", + ), + SpanAttributes.GEN_AI_REQUEST_STOP_SEQUENCES: kwargs.get( + "stop_sequences", + ), + SpanAttributes.GEN_AI_REQUEST_SEED: kwargs.get("seed"), + } + + # Extract tool definitions if provided + tool_definitions = _get_tool_definitions( + tools=kwargs.get("tools"), + tool_choice=kwargs.get("tool_choice"), + ) + if tool_definitions: + attributes[SpanAttributes.GEN_AI_TOOL_DEFINITIONS] = tool_definitions + + return {k: v for k, v in attributes.items() if v is not None} + + +def _get_llm_span_name(attributes: Dict[str, str]) -> str: + """Generate span name for LLM operations. + + Args: + attributes (`Dict[str, str]`): + LLM request attributes dictionary containing operation name and + model name. + + Returns: + `str`: + Formatted span name in the format "{operation} {model}", + e.g., "chat gpt-4" or "chat qwen-plus". + """ + return ( + f"{attributes[SpanAttributes.GEN_AI_OPERATION_NAME]} " + f"{attributes[SpanAttributes.GEN_AI_REQUEST_MODEL]}" + ) + + +def _get_llm_output_messages( + chat_response: ChatResponse | None, +) -> list[dict[str, Any]]: + """Extract and format LLM output messages for tracing. + + Converts ChatResponse objects to standardized message format compatible + with OpenTelemetry GenAI specification. + + Args: + chat_response (` ChatResponse | None`): + Chat response object with content blocks. Should be a ChatResponse + instance containing content blocks (text, tool_use, etc.). + + Returns: + `list[dict[str, Any]]`: + List containing a single formatted message dictionary with role, + parts, and finish_reason. Returns the original response if it's + not a ChatResponse instance, or an error message format if + conversion fails. + """ + try: + if not isinstance(chat_response, ChatResponse): + return [ + { + "role": "assistant", + "parts": [ + { + "type": "text", + "content": str(chat_response), + }, + ], + "finish_reason": "unknown", + }, + ] + + parts = [] + finish_reason = "stop" # Default finish reason + + for block in chat_response.content: + part = _convert_block_to_part(block) + if part: + parts.append(part) + + output_message = { + "role": "assistant", + "parts": parts, + "finish_reason": finish_reason, + } + + return [output_message] + + except Exception: + return [ + { + "role": "assistant", + "parts": [ + { + "type": "text", + "content": "", + }, + ], + "finish_reason": "error", + }, + ] + + +def _get_llm_response_attributes( + chat_response: ChatResponse | None, +) -> Dict[str, Any]: + """Get LLM response attributes for OpenTelemetry tracing. + + Extracts response metadata and formats into GenAI attributes. + + Args: + chat_response (`ChatResponse | None`): + Chat response object with data and usage info. Should have + attributes like id, usage (with input_tokens and output_tokens), + and content blocks. + + Returns: + `Dict[str, Any]`: + OpenTelemetry GenAI response attributes including response ID, + finish reasons, token usage (input/output tokens), and output + messages. + """ + attributes = { + SpanAttributes.GEN_AI_RESPONSE_ID: getattr( + chat_response, + "id", + "unknown_id", + ), + # FIXME: finish reason should be capture in chat response + SpanAttributes.GEN_AI_RESPONSE_FINISH_REASONS: '["stop"]', + } + if hasattr(chat_response, "usage") and chat_response.usage: + usage = chat_response.usage + attributes[ + SpanAttributes.GEN_AI_USAGE_INPUT_TOKENS + ] = usage.input_tokens + attributes[ + SpanAttributes.GEN_AI_USAGE_OUTPUT_TOKENS + ] = usage.output_tokens + + cache_input = usage.cache_input_tokens + if cache_input: + attributes[ + SpanAttributes.AGENTSCOPE_CACHE_INPUT_TOKENS + ] = cache_input + + cache_creation = usage.cache_creation_input_tokens + if cache_creation: + attributes[ + SpanAttributes.AGENTSCOPE_CACHE_CREATION_INPUT_TOKENS + ] = cache_creation + + output_messages = _get_llm_output_messages(chat_response) + if output_messages: + attributes[SpanAttributes.GEN_AI_OUTPUT_MESSAGES] = _serialize_to_str( + output_messages, + ) + + return attributes + + +def _get_agent_messages( + msg: Msg | list[Msg], +) -> list[dict[str, Any]]: + """Convert AgentScope message(s) to standardized parts format. + + Transforms Msg objects into OpenTelemetry GenAI format. + + Args: + msg (`Msg | list[Msg]`): + AgentScope message object or list of message objects with + content blocks. + + Returns: + `list[dict[str, Any]]`: + List of formatted message dictionaries with role, parts, name, + and finish_reason. + """ + try: + if isinstance(msg, Msg): + msg = [msg] + + formatted_msgs = [] + for m in msg: + parts = [] + for block in m.get_content_blocks(): + part = _convert_block_to_part(block) + if part: + parts.append(part) + formatted_msg = { + "role": m.role, + "parts": parts, + "name": m.name, + "finish_reason": "stop", + } + formatted_msgs.append(formatted_msg) + + return formatted_msgs + except Exception: + # Fallback: try simple attribute access on the original object. + # If msg was already converted to a list or lacks role/name, return + # an empty list rather than raising a secondary exception. + try: + single = msg[0] if isinstance(msg, list) else msg + return [ + { + "role": single.role, + "parts": [ + { + "type": "text", + "content": ( + str(single.content) if single.content else "" + ), + }, + ], + "name": single.name, + "finish_reason": "stop", + }, + ] + except Exception: + return [] + + +def _get_agent_request_attributes( + instance: "Agent", + kwargs: Dict[str, Any], +) -> Dict[str, str]: + """Get agent request attributes for OpenTelemetry tracing. + + Extracts agent metadata and input data into GenAI attributes. + + Args: + instance (`Agent`): + The agent instance making the request. + kwargs (`Dict[str, Any]`): + Keyword arguments passed to the agent's reply method. + + Returns: + `Dict[str, str]`: + OpenTelemetry GenAI attributes including operation name, agent ID, + agent name, agent description, and input messages (if provided). + """ + attributes = { + SpanAttributes.GEN_AI_OPERATION_NAME: ( + OperationNameValues.INVOKE_AGENT + ), + SpanAttributes.GEN_AI_AGENT_NAME: instance.name, + SpanAttributes.GEN_AI_AGENT_DESCRIPTION: inspect.getdoc( + instance.__class__, + ) + or "No description available", + } + + inputs = kwargs.get("inputs") + if inputs is not None: + if isinstance(inputs, (Msg, list)): + input_messages = _get_agent_messages(inputs) + attributes[ + SpanAttributes.GEN_AI_INPUT_MESSAGES + ] = _serialize_to_str(input_messages) + elif isinstance(inputs, UserConfirmResultEvent): + attributes[ + SpanAttributes.AGENTSCOPE_INCOMING_EVENT_TYPE + ] = "user_confirm_result" + elif isinstance(inputs, ExternalExecutionResultEvent): + attributes[ + SpanAttributes.AGENTSCOPE_INCOMING_EVENT_TYPE + ] = "external_execution_result" + + return attributes + + +def _get_agent_span_name(attributes: Dict[str, str]) -> str: + """Generate span name for agent operations. + + Args: + attributes (`Dict[str, str]`): + Agent request attributes dictionary containing operation name and + agent name. + + Returns: + `str`: + Formatted span name in the format "{operation} {agent_name}", + e.g., "invoke_agent MyAgent". + """ + return ( + f"{attributes[SpanAttributes.GEN_AI_OPERATION_NAME]} " + f"{attributes[SpanAttributes.GEN_AI_AGENT_NAME]}" + ) + + +def _get_agent_response_attributes( + agent_response: Msg, +) -> Dict[str, str]: + """Get agent response attributes for OpenTelemetry tracing. + + Args: + agent_response (`Msg`): + Response message returned by agent, containing content blocks. + + Returns: + `Dict[str, str]`: + OpenTelemetry GenAI response attributes including output messages. + """ + attributes = { + SpanAttributes.GEN_AI_OUTPUT_MESSAGES: _serialize_to_str( + _get_agent_messages(agent_response), + ), + } + return attributes + + +def _get_tool_request_attributes( + instance: "Toolkit", + tool_call: ToolCallBlock, +) -> Dict[str, str]: + """Get tool request attributes for OpenTelemetry tracing. + + Extracts tool execution metadata into GenAI attributes. + + Args: + instance (`Toolkit`): + Toolkit instance with tool definitions. Used to extract tool + description from the tool's JSON schema. + tool_call (`ToolCallBlock`): + Tool use block with call information including id, name, and input + arguments. + + Returns: + `Dict[str, str]`: + OpenTelemetry GenAI tool attributes including operation name, tool + call ID, tool name, tool description (if available), and tool call + arguments. + """ + attributes = { + SpanAttributes.GEN_AI_OPERATION_NAME: ( + OperationNameValues.EXECUTE_TOOL + ), + } + + if tool_call: + tool_name = tool_call.name + attributes[SpanAttributes.GEN_AI_TOOL_CALL_ID] = tool_call.id + attributes[SpanAttributes.GEN_AI_TOOL_NAME] = tool_name + # tool_call.input is already a JSON string; pass it directly to avoid + # double-encoding (e.g. '{"city": "Beijing"}' → '"{\\"city\\"...}"') + attributes[SpanAttributes.GEN_AI_TOOL_CALL_ARGUMENTS] = tool_call.input + + if tool_name: + registered = getattr(instance, "tools", {}).get(tool_name) + if registered is not None: + tool_obj = getattr(registered, "tool", None) + description = getattr(tool_obj, "description", None) + if description: + attributes[ + SpanAttributes.GEN_AI_TOOL_DESCRIPTION + ] = description + + return attributes + + +def _get_tool_span_name(attributes: Dict[str, str]) -> str: + """Generate span name for tool operations. + + Args: + attributes (`Dict[str, str]`): + Tool request attributes dictionary containing operation name and + tool name. + + Returns: + `str`: + Formatted span name in the format "{operation} {tool_name}", + e.g., "execute_tool search". + """ + return ( + f"{attributes[SpanAttributes.GEN_AI_OPERATION_NAME]} " + f"{attributes[SpanAttributes.GEN_AI_TOOL_NAME]}" + ) + + +def _get_tool_response_attributes( + tool_response: Any, +) -> Dict[str, str]: + """Get tool response attributes for OpenTelemetry tracing. + + Args: + tool_response (`Any`): + Response object from tool execution. Can be any serializable object + returned by the tool function. + + Returns: + `Dict[str, str]`: + OpenTelemetry GenAI response attributes including tool call result. + """ + attributes = { + SpanAttributes.GEN_AI_TOOL_CALL_RESULT: _serialize_to_str( + tool_response, + ), + } + return attributes diff --git a/src/agentscope/middleware/_tracing/_setup.py b/src/agentscope/middleware/_tracing/_setup.py new file mode 100644 index 0000000000000000000000000000000000000000..b74abf81d4a7cc0f860f9718ebe19a6b56589dd7 --- /dev/null +++ b/src/agentscope/middleware/_tracing/_setup.py @@ -0,0 +1,19 @@ +# -*- coding: utf-8 -*- +"""The tracing interface class in agentscope.""" +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from opentelemetry.trace import Tracer +else: + Tracer = "Tracer" + + +def _get_tracer() -> Tracer: + """Get the tracer + Returns: + `Tracer`: The tracer with the name "agentscope" and version. + """ + from opentelemetry import trace + from ..._version import __version__ + + return trace.get_tracer("agentscope", __version__) diff --git a/src/agentscope/middleware/_tracing/_trace.py b/src/agentscope/middleware/_tracing/_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..48a9e87618f45b95cd2b94b658b5e12ae956d4ea --- /dev/null +++ b/src/agentscope/middleware/_tracing/_trace.py @@ -0,0 +1,348 @@ +# -*- coding: utf-8 -*- +"""TracingMiddleware and supporting utilities for OpenTelemetry tracing.""" +import json +from typing import ( + Any, + AsyncGenerator, + Callable, + Awaitable, + Union, + TypeVar, + TYPE_CHECKING, +) + +import aioitertools + +from opentelemetry import trace as otel_trace +from opentelemetry.trace import StatusCode + +from .._base import MiddlewareBase +from ...event import ( + ExternalExecutionResultEvent, + RequireExternalExecutionEvent, + RequireUserConfirmEvent, + ReplyStartEvent, +) +from ...message import Msg, ToolCallBlock +from ...model import ChatModelBase + +from ._attributes import SpanAttributes, OperationNameValues +from ._extractor import ( + _get_common_attributes, + _get_agent_request_attributes, + _get_agent_span_name, + _get_agent_response_attributes, + _get_llm_request_attributes, + _get_llm_span_name, + _get_llm_response_attributes, + _get_tool_request_attributes, + _get_tool_span_name, + _get_tool_response_attributes, +) +from ._setup import _get_tracer +from ._utils import _serialize_to_str + +if TYPE_CHECKING: + from opentelemetry.trace import Span + from ...agent import Agent + from ...model import ChatResponse + +T = TypeVar("T") + + +# --------------------------------------------------------------------------- +# Utility helpers +# --------------------------------------------------------------------------- + + +def _check_tracing_enabled() -> bool: + """Check if the OpenTelemetry tracer is initialised with a real SDK + TracerProvider (i.e. ``setup_tracing`` was called). Returns ``False`` + when only the default no-op proxy provider is active or when the SDK + package is not installed. + """ + try: + from opentelemetry.sdk.trace import TracerProvider + except ImportError: + return False + + return isinstance(otel_trace.get_tracer_provider(), TracerProvider) + + +def _set_span_success_status(span: "Span") -> None: + """Set the span status to OK and end it.""" + span.set_status(StatusCode.OK) + span.end() + + +def _set_span_error_status(span: "Span", e: BaseException) -> None: + """Set the span status to ERROR, record the exception and end it.""" + span.set_status(StatusCode.ERROR, str(e)) + span.record_exception(e) + span.end() + + +async def _trace_async_generator_wrapper( + res: AsyncGenerator[T, None], + span: "Span", +) -> AsyncGenerator[T, None]: + """Wrap an async generator so that response attributes are captured from + the last yielded chunk before the span is closed.""" + has_error = False + + try: + last_chunk = None + async for chunk in aioitertools.iter(res): + last_chunk = chunk + yield chunk + + except BaseException as e: + has_error = True + _set_span_error_status(span, e) + raise + + finally: + if not has_error: + response_attributes = _get_llm_response_attributes(last_chunk) + span.set_attributes(response_attributes) + _set_span_success_status(span) + + +# --------------------------------------------------------------------------- +# TracingMiddleware +# --------------------------------------------------------------------------- + + +class TracingMiddleware(MiddlewareBase): + """Agent middleware that adds OpenTelemetry tracing to the reply, + model-call and tool-execution lifecycles. + + When tracing has not been configured (``setup_tracing`` was not called), + every hook short-circuits to ``next_handler`` with near-zero overhead. + + Example:: + + from agentscope.middleware import TracingMiddleware + + agent = Agent( + ..., + middlewares=[TracingMiddleware()], + ) + """ + + # ------------------------------------------------------------------ + # on_reply + # ------------------------------------------------------------------ + async def on_reply( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + if not _check_tracing_enabled(): + async for item in next_handler(**input_kwargs): + yield item + return + + session_id = agent.state.session_id + common_attrs = _get_common_attributes(session_id) + + tracer = _get_tracer() + request_attributes = _get_agent_request_attributes( + agent, + input_kwargs, + ) + span_name = _get_agent_span_name(request_attributes) + + with tracer.start_as_current_span( + name=span_name, + attributes={ + **request_attributes, + **common_attrs, + }, + end_on_exit=False, + ) as span: + # Synthetic execute_tool spans for externally executed tools + event_arg = input_kwargs.get("inputs") + if isinstance(event_arg, ExternalExecutionResultEvent): + for result in event_arg.execution_results: + tool_attrs: dict[str, Any] = { + SpanAttributes.GEN_AI_OPERATION_NAME: ( + OperationNameValues.EXECUTE_TOOL + ), + SpanAttributes.GEN_AI_TOOL_CALL_ID: result.id, + SpanAttributes.GEN_AI_TOOL_NAME: result.name, + SpanAttributes.AGENTSCOPE_IS_EXTERNAL_EXECUTION: ( + True + ), + **common_attrs, + } + if result.output is not None: + tool_attrs[ + SpanAttributes.GEN_AI_TOOL_CALL_RESULT + ] = _serialize_to_str(result.output) + with tracer.start_as_current_span( + name=( + f"{OperationNameValues.EXECUTE_TOOL}" + f" {result.name}" + ), + attributes=tool_attrs, + ): + pass + + has_error = False + error_exc: BaseException | None = None + last_msg: Msg | None = None + hitl_pending: list[str] = [] + external_pending: list[str] = [] + observed_reply_id: str | None = None + + try: + async for item in next_handler(**input_kwargs): + if isinstance(item, ReplyStartEvent): + observed_reply_id = item.reply_id + elif isinstance(item, RequireUserConfirmEvent): + hitl_pending.extend(t.name for t in item.tool_calls) + elif isinstance(item, RequireExternalExecutionEvent): + external_pending.extend( + t.name for t in item.tool_calls + ) + if isinstance(item, Msg): + last_msg = item + yield item + except BaseException as e: + has_error = True + error_exc = e + raise + finally: + reply_id = observed_reply_id or agent.state.reply_id + if reply_id: + span.set_attribute( + SpanAttributes.AGENTSCOPE_REPLY_ID, + reply_id, + ) + if hitl_pending: + span.set_attribute( + SpanAttributes.AGENTSCOPE_HITL_PENDING_TOOLS, + json.dumps(hitl_pending, ensure_ascii=False), + ) + if external_pending: + span.set_attribute( + SpanAttributes.AGENTSCOPE_EXTERNAL_EXECUTION_PENDING_TOOLS, # noqa + json.dumps(external_pending, ensure_ascii=False), + ) + if has_error and error_exc is not None: + _set_span_error_status(span, error_exc) + else: + if last_msg is not None: + span.set_attributes( + _get_agent_response_attributes(last_msg), + ) + _set_span_success_status(span) + + # ------------------------------------------------------------------ + # on_model_call + # ------------------------------------------------------------------ + async def on_model_call( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[ + ..., + Awaitable["ChatResponse" | AsyncGenerator["ChatResponse", None]], + ], + ) -> Union["ChatResponse", AsyncGenerator["ChatResponse", None]]: + if not _check_tracing_enabled(): + return await next_handler(**input_kwargs) + + model = input_kwargs.get("current_model") + if not isinstance(model, ChatModelBase): + return await next_handler(**input_kwargs) + + tracer = _get_tracer() + + combined_kwargs = { + **getattr(model, "generate_kwargs", {}), + **input_kwargs, + } + request_attributes = _get_llm_request_attributes( + model, + combined_kwargs, + ) + span_name = _get_llm_span_name(request_attributes) + + with tracer.start_as_current_span( + name=span_name, + attributes={ + **request_attributes, + **_get_common_attributes(agent.state.session_id), + }, + end_on_exit=False, + ) as span: + try: + result = await next_handler(**input_kwargs) + + if isinstance(result, AsyncGenerator): + return _trace_async_generator_wrapper(result, span) + + span.set_attributes(_get_llm_response_attributes(result)) + _set_span_success_status(span) + return result + + except BaseException as e: + _set_span_error_status(span, e) + raise + + # ------------------------------------------------------------------ + # on_acting + # ------------------------------------------------------------------ + async def on_acting( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + if not _check_tracing_enabled(): + async for item in next_handler(**input_kwargs): + yield item + return + + tool_call = input_kwargs.get("tool_call") + if not isinstance(tool_call, ToolCallBlock): + async for item in next_handler(**input_kwargs): + yield item + return + + tracer = _get_tracer() + + request_attributes = _get_tool_request_attributes( + agent.toolkit, + tool_call, + ) + span_name = _get_tool_span_name(request_attributes) + + with tracer.start_as_current_span( + name=span_name, + attributes={ + **request_attributes, + **_get_common_attributes(agent.state.session_id), + }, + end_on_exit=False, + ) as span: + has_error = False + last_item = None + try: + async for item in next_handler(**input_kwargs): + last_item = item + yield item + except BaseException as e: + has_error = True + _set_span_error_status(span, e) + raise + finally: + if not has_error: + if last_item is not None: + span.set_attributes( + _get_tool_response_attributes(last_item), + ) + _set_span_success_status(span) diff --git a/src/agentscope/middleware/_tracing/_utils.py b/src/agentscope/middleware/_tracing/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d3f5be2a91bc8e38c5e9701617b9ed947ea99742 --- /dev/null +++ b/src/agentscope/middleware/_tracing/_utils.py @@ -0,0 +1,78 @@ +# -*- coding: utf-8 -*- +"""Serialize objects to JSON string.""" +import datetime +import enum +import inspect +import json +from dataclasses import is_dataclass +from typing import Any + +from pydantic import BaseModel + +from ...message import Msg + + +def _to_serializable( + obj: Any, +) -> Any: + """Convert an object to a JSON serializable type. + + Args: + obj (`Any`): + The object to be converted to JSON serializable. + + Returns: + `Any`: + The converted JSON serializable object + """ + + # Handle primitive types first + if isinstance(obj, (str, int, bool, float, type(None))): + res = obj + + elif isinstance(obj, (list, tuple, set, frozenset)): + res = [_to_serializable(x) for x in obj] + + elif isinstance(obj, dict): + res = {str(key): _to_serializable(val) for (key, val) in obj.items()} + + elif isinstance(obj, (Msg, BaseModel)) or is_dataclass(obj): + res = repr(obj) + + elif inspect.isclass(obj) and issubclass(obj, BaseModel): + res = repr(obj) + + elif isinstance(obj, (datetime.date, datetime.datetime, datetime.time)): + res = obj.isoformat() + + elif isinstance(obj, datetime.timedelta): + res = obj.total_seconds() + + elif isinstance(obj, enum.Enum): + res = _to_serializable(obj.value) + + else: + res = str(obj) + + return res + + +def _serialize_to_str(value: Any) -> str: + """Serialize input value to JSON string. + + Args: + value (`Any`): + The input value + + Returns: + `str`: + JSON serialized string of the input value + """ + try: + return json.dumps(value, ensure_ascii=False) + + except TypeError: + return json.dumps( + _to_serializable(value), + ensure_ascii=False, + ) diff --git a/src/agentscope/middleware/_tts_middleware.py b/src/agentscope/middleware/_tts_middleware.py new file mode 100644 index 0000000000000000000000000000000000000000..3194c15078f50ecc902c2294bdf3153dcbb48832 --- /dev/null +++ b/src/agentscope/middleware/_tts_middleware.py @@ -0,0 +1,179 @@ +# -*- coding: utf-8 -*- +"""Middleware that turns reasoning text into speech and injects it as +``DATA_BLOCK_*`` events into the agent's event stream.""" +from typing import TYPE_CHECKING, AsyncGenerator, Callable + +from ._base import MiddlewareBase +from .._utils._common import _generate_id +from ..event import ( + DataBlockDeltaEvent, + DataBlockEndEvent, + DataBlockStartEvent, + TextBlockDeltaEvent, + TextBlockEndEvent, +) +from ..tts import TTSModelBase, TTSResponse + +if TYPE_CHECKING: + from ..agent import Agent + + +class TTSMiddleware(MiddlewareBase): + """Synthesize speech for every text block produced during reasoning and + inject the audio as ``DATA_BLOCK_*`` events into the stream. + + - Non-realtime TTS (``realtime=False``): on each + ``TextBlockEndEvent`` the accumulated text is sent to + :meth:`TTSModelBase.synthesize`; the resulting audio chunks are + emitted as one ``DATA_BLOCK_START`` + N ``DATA_BLOCK_DELTA`` + + ``DATA_BLOCK_END``. + - Realtime TTS (``realtime=True``): each + ``TextBlockDeltaEvent`` is pushed into the model via + :meth:`TTSModelBase.push`; any audio produced is emitted immediately. + On ``TextBlockEndEvent`` :meth:`TTSModelBase.synthesize` is called + to drain remaining audio, and the data block is closed. + + Each ``DataBlockDeltaEvent.data`` carries an **incremental** base64 PCM + chunk; the full audio is the concatenation of every delta's decoded + bytes (the data block is keyed by ``block_id``). + """ + + def __init__(self, tts_model: TTSModelBase) -> None: + """Initialize the TTS middleware. + + Args: + tts_model (`TTSModelBase`): + The TTS model used to synthesize speech for assistant text + blocks produced during reasoning. + """ + self.tts = tts_model + + async def on_reply( + self, + agent: "Agent", + input_kwargs: dict, + next_handler: Callable[..., AsyncGenerator], + ) -> AsyncGenerator: + """Intercept the reply stream, synthesize speech for text blocks, + and inject ``DATA_BLOCK_*`` audio events into the output.""" + text_buffer: str = "" + audio_block_id: str | None = None + audio_media_type: str | None = None + + async with self.tts: + async for evt in next_handler(**input_kwargs): + yield evt + + if isinstance(evt, TextBlockDeltaEvent): + text_buffer += evt.delta + if self.tts.realtime and evt.delta: + tts_res = await self.tts.push(evt.delta) + async for audio_evt in self._emit_chunk( + agent, + tts_res, + audio_block_id, + audio_media_type, + ): + if isinstance(audio_evt, DataBlockStartEvent): + audio_block_id = audio_evt.block_id + audio_media_type = audio_evt.media_type + yield audio_evt + + elif isinstance(evt, TextBlockEndEvent): + text = text_buffer + text_buffer = "" + + if self.tts.realtime: + res = await self.tts.synthesize() + async for audio_evt in self._emit_synth_result( + agent, + res, + audio_block_id, + audio_media_type, + ): + if isinstance(audio_evt, DataBlockStartEvent): + audio_block_id = audio_evt.block_id + audio_media_type = audio_evt.media_type + yield audio_evt + elif text.strip(): + res = await self.tts.synthesize(text) + async for audio_evt in self._emit_synth_result( + agent, + res, + audio_block_id, + audio_media_type, + ): + if isinstance(audio_evt, DataBlockStartEvent): + audio_block_id = audio_evt.block_id + audio_media_type = audio_evt.media_type + yield audio_evt + + if audio_block_id is not None: + yield DataBlockEndEvent( + reply_id=agent.state.reply_id, + block_id=audio_block_id, + ) + audio_block_id = None + audio_media_type = None + + async def _emit_synth_result( + self, + agent: "Agent", + res: TTSResponse | AsyncGenerator[TTSResponse, None], + audio_block_id: str | None, + audio_media_type: str | None, + ) -> AsyncGenerator: + """Normalize ``synthesize()`` returns (single response or async + generator) into a stream of ``DATA_BLOCK_*`` events.""" + if isinstance(res, AsyncGenerator): + async for chunk in res: + async for ae in self._emit_chunk( + agent, + chunk, + audio_block_id, + audio_media_type, + ): + if isinstance(ae, DataBlockStartEvent): + audio_block_id = ae.block_id + audio_media_type = ae.media_type + yield ae + else: + async for ae in self._emit_chunk( + agent, + res, + audio_block_id, + audio_media_type, + ): + yield ae + + @staticmethod + async def _emit_chunk( + agent: "Agent", + tts_res: TTSResponse | None, + audio_block_id: str | None, + audio_media_type: str | None, + ) -> AsyncGenerator: + """Emit one TTSResponse chunk as ``DATA_BLOCK_START`` (if needed) + followed by ``DATA_BLOCK_DELTA``.""" + if tts_res is None or tts_res.content is None: + return + media_type = tts_res.content.source.media_type + data = tts_res.content.source.data + if not data: + return + + if audio_block_id is None: + audio_block_id = _generate_id() + audio_media_type = media_type + yield DataBlockStartEvent( + reply_id=agent.state.reply_id, + block_id=audio_block_id, + media_type=media_type, + ) + + yield DataBlockDeltaEvent( + reply_id=agent.state.reply_id, + block_id=audio_block_id, + data=data, + media_type=audio_media_type, + ) diff --git a/src/agentscope/model/__init__.py b/src/agentscope/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0522d21f71c5a091f2f8ad8dd52639bae02aad3e --- /dev/null +++ b/src/agentscope/model/__init__.py @@ -0,0 +1,33 @@ +# -*- coding: utf-8 -*- +"""The model module.""" + +from ._base import ChatModelBase +from ._model_card import ModelCard +from ._model_response import ChatResponse, StructuredResponse +from ._model_usage import ChatUsage +from ._anthropic import AnthropicChatModel +from ._dashscope import DashScopeChatModel +from ._deepseek import DeepSeekChatModel +from ._gemini import GeminiChatModel +from ._ollama import OllamaChatModel +from ._openai_chat import OpenAIChatModel +from ._xai import XAIChatModel +from ._moonshot import MoonshotChatModel +from ._openai_response import OpenAIResponseModel + +__all__ = [ + "ChatUsage", + "ChatModelBase", + "ChatResponse", + "ModelCard", + "StructuredResponse", + "AnthropicChatModel", + "DashScopeChatModel", + "DeepSeekChatModel", + "GeminiChatModel", + "OllamaChatModel", + "OpenAIChatModel", + "XAIChatModel", + "MoonshotChatModel", + "OpenAIResponseModel", +] diff --git a/src/agentscope/model/_anthropic/__init__.py b/src/agentscope/model/_anthropic/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2be48e8e49ce726b41d051553b8ef99a6c6fb60f --- /dev/null +++ b/src/agentscope/model/_anthropic/__init__.py @@ -0,0 +1,9 @@ +# -*- coding: utf-8 -*- +"""The Anthropic LLM API modules.""" + +from ._model import AnthropicCredential, AnthropicChatModel + +__all__ = [ + "AnthropicCredential", + "AnthropicChatModel", +] diff --git a/src/agentscope/model/_anthropic/_model.py b/src/agentscope/model/_anthropic/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..46a743e1eae039b4c2c0075df45b406bc8f64ae4 --- /dev/null +++ b/src/agentscope/model/_anthropic/_model.py @@ -0,0 +1,564 @@ +# -*- coding: utf-8 -*- +"""The Anthropic chat model implementation.""" +from collections import OrderedDict +from datetime import datetime +from typing import Literal, Any, AsyncGenerator, TYPE_CHECKING, List, Type + +from pydantic import BaseModel, Field + +from .._base import ChatModelBase, _TOOL_CHOICE_LITERAL_MODES +from .._model_response import ChatResponse, StructuredResponse +from .._model_usage import ChatUsage +from ...credential import AnthropicCredential +from ...formatter import FormatterBase, AnthropicChatFormatter +from ...message import Msg, ThinkingBlock, ToolCallBlock, TextBlock +from ...tool import ToolChoice + +if TYPE_CHECKING: + from anthropic.types.message import Message + from anthropic import AsyncStream +else: + Message = Any + AsyncStream = Any + + +class AnthropicChatModel(ChatModelBase): + """The Anthropic chat model.""" + + type: Literal["anthropic_chat"] = "anthropic_chat" + """The type of the chat model.""" + + class Parameters(BaseModel): + """The parameters for the Anthropic chat model.""" + + max_tokens: int | None = Field( + default=None, + title="Max Tokens", + description=( + "The maximum number of tokens to generate in the chat " + "completion." + ), + gt=0, + ) + + thinking_enable: bool = Field( + default=False, + title="Thinking", + description="The thinking enable for the LLM output.", + ) + + thinking_budget: int | None = Field( + default=None, + title="Thinking Budget", + description="The thinking budget for the LLM output.", + gt=0, + ) + + def __init__( + self, + credential: AnthropicCredential, + model: str, + parameters: "AnthropicChatModel.Parameters | None" = None, + stream: bool = True, + max_retries: int = 3, + retry_delay: float = 1.0, + context_size: int = 200000, + formatter: FormatterBase | None = None, + client_kwargs: dict[str, Any] | None = None, + ) -> None: + """Initialize the Anthropic chat model. + + Args: + credential (`AnthropicCredential`): + The Anthropic credential used to authenticate API calls. + model (`str`): + The Anthropic model name, e.g. ``claude-opus-4-7``. + parameters (`AnthropicChatModel.Parameters | None`, defaults to \ + `None`): + The Anthropic API parameters. When ``None``, the default + parameters will be used. + stream (`bool`, defaults to `True`): + Whether to enable streaming output. + max_retries (`int`, defaults to `3`): + The maximum number of retries for the Anthropic API. + retry_delay (`float`, defaults to `1.0`): + Seconds to sleep between retry attempts. + context_size (`int`, defaults to `200000`): + The model context size used for context compression. + formatter (`FormatterBase | None`, defaults to `None`): + The formatter that converts ``Msg`` objects to the format + required by the Anthropic API. When ``None``, an + ``AnthropicChatFormatter`` instance will be used. + client_kwargs (`dict[str, Any] | None`, defaults to `None`): + Extra keyword arguments forwarded to + ``anthropic.AsyncAnthropic`` (e.g. ``timeout``, + ``default_headers``, ``http_client``, ``auth_token``). + """ + super().__init__( + credential=credential, + model=model, + parameters=parameters or self.Parameters(), + stream=stream, + max_retries=max_retries, + retry_delay=retry_delay, + context_size=context_size, + ) + self.formatter = formatter or AnthropicChatFormatter() + self.client_kwargs = client_kwargs or {} + + @classmethod + def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]: + import anthropic + + return ( + anthropic.APIConnectionError, + anthropic.APITimeoutError, + anthropic.RateLimitError, + anthropic.InternalServerError, + ) + + async def _call_api( + self, + model_name: str, + messages: list[Msg], + tools: list[dict] | None = None, + tool_choice: ToolChoice | None = None, + **generate_kwargs: Any, + ) -> ChatResponse | AsyncGenerator[ChatResponse, None]: + """Get the response from Anthropic chat completions API by the given + arguments. + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list[dict]`): + A list of dictionaries, where `role` and `content` fields are + required, and `name` field is optional. + tools (`list[dict]`, default `None`): + The tools JSON schemas. + tool_choice (`ToolChoice | None`, optional): + Controls which (if any) tool is called by the model. + **generate_kwargs (`Any`): + The keyword arguments for Anthropic chat completions API. + + Returns: + `ChatResponse | AsyncGenerator[ChatResponse, None]`: + A ``ChatResponse`` when streaming is disabled, or an async + generator of ``ChatResponse`` objects when streaming is + enabled. + """ + + import anthropic + + client = anthropic.AsyncAnthropic( + **{ + "api_key": self.credential.api_key.get_secret_value(), + "base_url": self.credential.base_url, + **self.client_kwargs, + }, + ) + + # Anthropic requires max_tokens; fall back to a safe default when + # the user hasn't configured one explicitly. + max_tokens = self.parameters.max_tokens or 8192 + + kwargs: dict[str, Any] = { + "model": model_name, + "max_tokens": max_tokens, + "stream": self.stream, + **generate_kwargs, + } + + # Anthropic extended thinking — only set when explicitly enabled. + # Anthropic requires max_tokens > budget_tokens strictly. + if self.parameters.thinking_enable and "thinking" not in kwargs: + budget = self.parameters.thinking_budget or (max_tokens // 2) + if budget >= max_tokens: + # Auto-expand max_tokens to satisfy the strict inequality. + max_tokens = budget + 1024 + kwargs["max_tokens"] = max_tokens + kwargs["thinking"] = { + "type": "enabled", + "budget_tokens": budget, + } + + fmt_tools, fmt_tool_choice = self._format_tools(tools, tool_choice) + if fmt_tools: + kwargs["tools"] = fmt_tools + if fmt_tool_choice is not None: + kwargs["tool_choice"] = fmt_tool_choice + + formatted_messages = await self.formatter.format(messages) + + # Extract the system message + if formatted_messages and formatted_messages[0]["role"] == "system": + kwargs["system"] = formatted_messages[0]["content"] + formatted_messages = formatted_messages[1:] + + kwargs["messages"] = formatted_messages + + start_datetime = datetime.now() + + response = await client.messages.create(**kwargs) + + if self.stream: + return self._parse_anthropic_stream_completion_response( + start_datetime, + response, + ) + + # Non-streaming response + parsed_response = await self._parse_anthropic_completion_response( + start_datetime, + response, + ) + + return parsed_response + + async def _parse_anthropic_completion_response( + self, + start_datetime: datetime, + response: Message, + ) -> ChatResponse: + """Given an Anthropic Message object, extract the content blocks and + usages from it. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`Message`): + Anthropic Message object to parse. + + Returns: + `ChatResponse`: + A single ``ChatResponse`` with ``is_last=True`` containing + the extracted content blocks and usage. + """ + content_blocks: List[ThinkingBlock | TextBlock | ToolCallBlock] = [] + + if hasattr(response, "content") and response.content: + for content_block in response.content: + if ( + hasattr(content_block, "type") + and content_block.type == "thinking" + ): + thinking_block = ThinkingBlock( + thinking=content_block.thinking, + signature=getattr(content_block, "signature", "") + or "", + ) + content_blocks.append(thinking_block) + + elif ( + hasattr(content_block, "type") + and content_block.type == "text" + ): + content_blocks.append( + TextBlock(text=content_block.text), + ) + + elif ( + hasattr(content_block, "type") + and content_block.type == "tool_use" + ): + import json + + content_blocks.append( + ToolCallBlock( + id=content_block.id, + name=content_block.name, + input=json.dumps(content_block.input), + ), + ) + + usage = None + if response.usage: + u = response.usage + usage = ChatUsage( + input_tokens=u.input_tokens, + output_tokens=u.output_tokens, + time=(datetime.now() - start_datetime).total_seconds(), + cache_creation_input_tokens=getattr( + u, + "cache_creation_input_tokens", + 0, + ), + cache_input_tokens=getattr( + u, + "cache_read_input_tokens", + 0, + ), + ) + + resp_kwargs: dict[str, Any] = { + "content": content_blocks, + "is_last": True, + "usage": usage, + } + response_id = getattr(response, "id", None) + if response_id: + resp_kwargs["id"] = response_id + + return ChatResponse(**resp_kwargs) + + async def _parse_anthropic_stream_completion_response( + self, + start_datetime: datetime, + response: AsyncStream, + ) -> AsyncGenerator[ChatResponse, None]: + """Given an Anthropic streaming response, extract the content blocks + and usages from it and yield ChatResponse objects. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`AsyncStream`): + Anthropic AsyncStream object to parse. + + Yields: + `ChatResponse`: + Incremental ``ChatResponse`` objects with ``is_last=False`` + followed by a final one with ``is_last=True`` containing the + fully accumulated content blocks and usage. + """ + + usage = None + response_id: str | None = None + # All delta should have the same block identifier + acc_text = TextBlock(text="") + acc_thinking = ThinkingBlock(thinking="") + thinking_signature = "" + # index -> {id, name, input} + acc_tool_calls: OrderedDict = OrderedDict() + + async for event in response: + delta_content: list = [] + + if event.type == "message_start": + message = event.message + if response_id is None: + response_id = getattr(message, "id", None) + if message.usage: + u = message.usage + usage = ChatUsage( + input_tokens=u.input_tokens, + output_tokens=getattr(u, "output_tokens", 0), + time=(datetime.now() - start_datetime).total_seconds(), + cache_creation_input_tokens=getattr( + u, + "cache_creation_input_tokens", + 0, + ), + cache_input_tokens=getattr( + u, + "cache_read_input_tokens", + 0, + ), + ) + + elif event.type == "content_block_start": + if event.content_block.type == "tool_use": + block_index = event.index + tool_block = event.content_block + acc_tool_calls[block_index] = { + "id": tool_block.id, + "name": tool_block.name, + "input": "", + } + + elif event.type == "content_block_delta": + block_index = event.index + delta = event.delta + if delta.type == "text_delta": + acc_text.text += delta.text + delta_content.append( + TextBlock(id=acc_text.id, text=delta.text), + ) + elif delta.type == "thinking_delta": + acc_thinking.thinking += delta.thinking + delta_content.append( + ThinkingBlock( + id=acc_thinking.id, + thinking=delta.thinking, + ), + ) + elif delta.type == "signature_delta": + thinking_signature = delta.signature + elif ( + delta.type == "input_json_delta" + and block_index in acc_tool_calls + ): + fragment = delta.partial_json or "" + acc_tool_calls[block_index]["input"] += fragment + tc = acc_tool_calls[block_index] + delta_content.append( + ToolCallBlock( + id=tc["id"], + name=tc["name"], + input=fragment, + ), + ) + + elif event.type == "message_delta": + if event.usage and usage: + usage.output_tokens = event.usage.output_tokens + + if delta_content: + _kwargs: dict[str, Any] = { + "content": delta_content, + "is_last": False, + "usage": usage, + } + if response_id: + _kwargs["id"] = response_id + yield ChatResponse(**_kwargs) + + # Build final accumulated content + final_content: list = [] + if acc_thinking.thinking: + acc_thinking.signature = thinking_signature + final_content.append(acc_thinking) + if acc_text.text: + final_content.append(acc_text) + for tc in acc_tool_calls.values(): + input_str = tc["input"] + final_content.append( + ToolCallBlock( + id=tc["id"], + name=tc["name"], + input=input_str, + ), + ) + + _final_kwargs: dict[str, Any] = { + "content": final_content, + "is_last": True, + "usage": usage, + } + if response_id: + _final_kwargs["id"] = response_id + yield ChatResponse(**_final_kwargs) + + def _format_tools( + self, + tools: list[dict] | None, + tool_choice: ToolChoice | None, + ) -> tuple[list[dict] | None, dict | None]: + """Validate and format tools and tool_choice for Anthropic. + + Converts tool schemas to Anthropic's flat format and maps + tool_choice modes to Anthropic's type-based format. When + ``tool_choice.tools`` is specified the schemas list is filtered + to only those tools. When ``tool_choice.mode`` is a specific tool + name (str) the model is forced to call exactly that tool without + needing to filter the list, preserving prompt-cache efficiency. + + Args: + tools (`list[dict] | None`, optional): + The raw tool schemas. + tool_choice (`ToolChoice | None`, optional): + The tool choice configuration. + + Returns: + `tuple[list[dict] | None, dict | None]`: + A tuple of (formatted_tools, formatted_tool_choice). + """ + if tool_choice and tools: + self._validate_tool_choice(tool_choice, tools) + if tool_choice.tools: + allowed = set(tool_choice.tools) + tools = [t for t in tools if t["function"]["name"] in allowed] + + fmt_tools = None + if tools: + fmt_tools = [] + for schema in tools: + assert ( + "function" in schema + ), f"Invalid schema: {schema}, expect key 'function'." + assert "name" in schema["function"], ( + f"Invalid schema: {schema}, " + "expect key 'name' in 'function' field." + ) + fmt_tools.append( + { + "name": schema["function"]["name"], + "description": schema["function"].get( + "description", + "", + ), + "input_schema": schema["function"].get( + "parameters", + {}, + ), + }, + ) + + if not tool_choice: + return fmt_tools, None + + mode = tool_choice.mode + + if mode not in _TOOL_CHOICE_LITERAL_MODES: + # mode is a specific tool name — force call it + return fmt_tools, {"type": "tool", "name": mode} + + type_mapping = { + "auto": {"type": "auto"}, + "none": {"type": "none"}, + "required": {"type": "any"}, + } + return fmt_tools, type_mapping[mode] + + async def _call_api_with_structured_output( + self, + model_name: str, + messages: list[Msg], + structured_model: Type[BaseModel] | dict, + tool_choice: ToolChoice | None = None, + **kwargs: Any, + ) -> StructuredResponse: + """Anthropic-specific override for structured output. + + Anthropic's extended thinking mode only supports + ``tool_choice={"type": "auto"}`` or ``{"type": "none"}``; any + forcing form (``"any"`` or a specific tool) raises an API error. + When ``thinking_enable`` is on we default ``tool_choice`` to + ``"auto"`` and rely on the base class's injected system-reminder + prompt to guide the model. When thinking is disabled, this falls + through to the base implementation (force the structured-output + tool). + + See: + https://platform.claude.com/docs/en/build-with-claude/extended-thinking#extended-thinking-with-tool-use + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list[Msg]`): + The context for the LLM to generate the structured output. + structured_model (`Type[BaseModel] | dict`): + A Pydantic model class or a JSON schema dict describing the + required output structure. + tool_choice (`ToolChoice | None`, defaults to `None`): + The tool_choice forwarded to ``_call_api``. When ``None`` + and thinking mode is enabled, it is downgraded to + ``ToolChoice(mode="auto")``; otherwise the base default + (force the structured-output tool) is used. + **kwargs (`Any`): + Additional keyword arguments forwarded to ``_call_api``. + + Returns: + `StructuredResponse`: + The structured response whose ``content`` is the validated + output dict matching ``structured_model``. + """ + if tool_choice is None and self.parameters.thinking_enable: + tool_choice = ToolChoice(mode="auto") + return await super()._call_api_with_structured_output( + model_name=model_name, + messages=messages, + structured_model=structured_model, + tool_choice=tool_choice, + **kwargs, + ) diff --git a/src/agentscope/model/_anthropic/_models/claude-haiku-4-5.yaml b/src/agentscope/model/_anthropic/_models/claude-haiku-4-5.yaml new file mode 100644 index 0000000000000000000000000000000000000000..11a191d024575e5690b83af5742c0893bb93ce6f --- /dev/null +++ b/src/agentscope/model/_anthropic/_models/claude-haiku-4-5.yaml @@ -0,0 +1,21 @@ +name: claude-haiku-4-5 +label: Claude Haiku 4.5 +status: active + +input_types: + - text/plain + - application/x-thinking + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + - application/x-thinking + +context_size: 200000 +output_size: 65536 + +parameter_overrides: + max_tokens: {"maximum": 65536} diff --git a/src/agentscope/model/_anthropic/_models/claude-opus-4-5.yaml b/src/agentscope/model/_anthropic/_models/claude-opus-4-5.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d9247d3932d8bf04a452a7c5430bf8592bc14466 --- /dev/null +++ b/src/agentscope/model/_anthropic/_models/claude-opus-4-5.yaml @@ -0,0 +1,21 @@ +name: claude-opus-4-5 +label: Claude Opus 4.5 +status: active + +input_types: + - text/plain + - application/x-thinking + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + - application/x-thinking + +context_size: 200000 +output_size: 65536 + +parameter_overrides: + max_tokens: {"maximum": 65536} diff --git a/src/agentscope/model/_anthropic/_models/claude-opus-4-6.yaml b/src/agentscope/model/_anthropic/_models/claude-opus-4-6.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7c2a82ec6c7f6e09c521bdfcef6eef9816fb40f5 --- /dev/null +++ b/src/agentscope/model/_anthropic/_models/claude-opus-4-6.yaml @@ -0,0 +1,21 @@ +name: claude-opus-4-6 +label: Claude Opus 4.6 +status: active + +input_types: + - text/plain + - application/x-thinking + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + - application/x-thinking + +context_size: 1000000 +output_size: 131072 + +parameter_overrides: + max_tokens: {"maximum": 131072} diff --git a/src/agentscope/model/_anthropic/_models/claude-opus-4-7.yaml b/src/agentscope/model/_anthropic/_models/claude-opus-4-7.yaml new file mode 100644 index 0000000000000000000000000000000000000000..854999319dab6947b81171584c76f02f38d2b380 --- /dev/null +++ b/src/agentscope/model/_anthropic/_models/claude-opus-4-7.yaml @@ -0,0 +1,21 @@ +name: claude-opus-4-7 +label: Claude Opus 4.7 +status: active + +input_types: + - text/plain + - application/x-thinking + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + - application/x-thinking + +context_size: 1000000 +output_size: 128000 + +parameter_overrides: + max_tokens: {"maximum": 128000} diff --git a/src/agentscope/model/_anthropic/_models/claude-opus-4-8.yaml b/src/agentscope/model/_anthropic/_models/claude-opus-4-8.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4bf81ca142b0241f4bbb41249bcfd5f2e7b60ba6 --- /dev/null +++ b/src/agentscope/model/_anthropic/_models/claude-opus-4-8.yaml @@ -0,0 +1,21 @@ +name: claude-opus-4-8 +label: Claude Opus 4.8 +status: active + +input_types: + - text/plain + - application/x-thinking + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + - application/x-thinking + +context_size: 1000000 +output_size: 128000 + +parameter_overrides: + max_tokens: {"maximum": 128000} diff --git a/src/agentscope/model/_anthropic/_models/claude-sonnet-4-5.yaml b/src/agentscope/model/_anthropic/_models/claude-sonnet-4-5.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e005495f7ea05fb61e4b05577fe5126e52a2e571 --- /dev/null +++ b/src/agentscope/model/_anthropic/_models/claude-sonnet-4-5.yaml @@ -0,0 +1,21 @@ +name: claude-sonnet-4-5 +label: Claude Sonnet 4.5 +status: active + +input_types: + - text/plain + - application/x-thinking + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + - application/x-thinking + +context_size: 200000 +output_size: 65536 + +parameter_overrides: + max_tokens: {"maximum": 65536} diff --git a/src/agentscope/model/_anthropic/_models/claude-sonnet-4-6.yaml b/src/agentscope/model/_anthropic/_models/claude-sonnet-4-6.yaml new file mode 100644 index 0000000000000000000000000000000000000000..81bcb4da8637d8e873eb294020b8dc91dc3121c4 --- /dev/null +++ b/src/agentscope/model/_anthropic/_models/claude-sonnet-4-6.yaml @@ -0,0 +1,21 @@ +name: claude-sonnet-4-6 +label: Claude Sonnet 4.6 +status: active + +input_types: + - text/plain + - application/x-thinking + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + - application/x-thinking + +context_size: 1000000 +output_size: 65536 + +parameter_overrides: + max_tokens: {"maximum": 65536} diff --git a/src/agentscope/model/_base.py b/src/agentscope/model/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..19e3a7792dd746f986a565d1729993a205b45f4f --- /dev/null +++ b/src/agentscope/model/_base.py @@ -0,0 +1,585 @@ +# -*- coding: utf-8 -*- +"""The base class for the chat models.""" +import asyncio +import inspect +import json +from abc import abstractmethod +from copy import deepcopy +from pathlib import Path +from typing import Type, Any, AsyncGenerator + +import jsonschema +from pydantic import BaseModel + +from ._model_response import StructuredResponse, ChatResponse +from ._model_card import ModelCard +from .._logging import logger +from .._utils._common import _json_loads_with_repair +from ..credential import CredentialBase +from ..message import ( + Msg, + TextBlock, + UserMsg, + ToolCallBlock, + ThinkingBlock, + ToolResultBlock, + DataBlock, + URLSource, + Base64Source, + HintBlock, +) +from ..tool import ToolChoice + +_TOOL_CHOICE_LITERAL_MODES = {"auto", "none", "required"} + + +class ChatModelBase: + """The base class for chat models.""" + + class Parameters(BaseModel): + """Each subclass should implement this inner class to define its + parameters.""" + + credential: CredentialBase + """The API credential.""" + + model: str + """The model name.""" + + stream: bool + """The enable stream output for the LLM output.""" + + max_retries: int + """The maximum number of retries for the underlying API.""" + + retry_delay: float + """Seconds to sleep between retry attempts.""" + + context_size: int + """The model context size that will be used in the context compression.""" + + def __init__( + self, + credential: CredentialBase, + model: str, + parameters: BaseModel, + stream: bool = True, + max_retries: int = 3, + retry_delay: float = 1.0, + context_size: int = 32768, + ) -> None: + """Initialize the chat model base. + + Args: + credential (CredentialBase): + The API credential. + model (`str`): + The model name. + parameters (`BaseModel`): + The model parameters. + stream (`bool`, defaults to `True`): + Whether to enable streaming output for the LLM. + max_retries (`int`, defaults to `3`): + The maximum number of retries for API calls. Only exceptions + listed in ``_get_retryable_exceptions()`` count against this + budget; other exceptions are raised immediately. + retry_delay (`float`, defaults to `1.0`): + Seconds to sleep between retry attempts. + context_size (`int`, defaults to `32768`): + The model context size used for context compression. + """ + self.credential = credential + self.model = model + self.parameters = parameters + self.stream = stream + self.max_retries = max_retries + self.retry_delay = retry_delay + self.context_size = context_size + + @classmethod + def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]: + """Return the exception types that should trigger a retry. + + Defaults to an empty tuple (no retries). Subclasses can override to + declare provider-specific retryable exceptions. SDK exception types + should be imported lazily inside the override so the SDK stays an + optional dependency. + """ + return () + + @classmethod + def list_models( + cls, + custom_yaml_dir: str | None = None, + ) -> list[ModelCard]: + """List candidate models of the API. + + Args: + custom_yaml_dir (`str | None`): + The custom YAML directory. + + Returns: + `list[ModelCard]`: + A list of candidate models. + """ + + # Determine YAML directory + if custom_yaml_dir is None: + # Use the ``_models`` directory that sits next to the concrete + # subclass's source file (not this base file). + subclass_file = Path(inspect.getfile(cls)) + yaml_dir = subclass_file.parent / "_models" + else: + yaml_dir = Path(custom_yaml_dir) + + # Find all .yaml files + yaml_files = list(yaml_dir.glob("*.yaml")) + + # Load each YAML file and create ModelCard + model_cards = [] + for yaml_file in yaml_files: + try: + card = ModelCard.from_yaml( + yaml_path=str(yaml_file), + parameter_class=cls.Parameters, + ) + model_cards.append(card) + except Exception as e: + # Log error but continue with other files + logger.warning( + "Warning: Failed to load %s: %s", + yaml_file, + str(e), + ) + continue + + return model_cards + + async def __call__( + self, + messages: list[Msg], + tools: list[dict] | None = None, + tool_choice: ToolChoice | None = None, + **kwargs: Any, + ) -> ChatResponse | AsyncGenerator[ChatResponse, None]: + """Call the model with retry logic. + + Attempts to call the model up to ``max_retries + 1`` times. Only + exceptions listed in ``_get_retryable_exceptions()`` count against + this budget; other exceptions are raised immediately. + + Args: + messages (`list[Msg]`): + The messages to send to the model. + tools (`list[dict] | None`, optional): + The tools available to the model. + tool_choice (`ToolChoice | None`, optional): + The tool choice mode or function name. + **kwargs: + Additional keyword arguments passed to the underlying API. + """ + + retryable = tuple(self._get_retryable_exceptions()) + last_error: Exception | None = None + for attempt in range(self.max_retries + 1): + try: + return await self._call_api( + self.model, + messages=messages, + tools=tools, + tool_choice=tool_choice, + **kwargs, + ) + except Exception as e: + if not isinstance(e, retryable): + raise + last_error = e + if attempt < self.max_retries: + logger.warning( + "Attempt %d failed for model %s: %s. " + "Retrying in %.1fs...", + attempt + 1, + self.model, + str(e), + self.retry_delay, + ) + await asyncio.sleep(self.retry_delay) + else: + logger.warning( + "All %d attempt(s) failed for model %s.", + self.max_retries + 1, + self.model, + ) + if last_error is not None: + raise last_error + raise RuntimeError( + f"Failed to call model {self.model} after " + f"{self.max_retries + 1} retries.", + ) + + @abstractmethod + async def _call_api( + self, + model_name: str, + messages: list[Msg], + tools: list[dict] | None = None, + tool_choice: ToolChoice | None = None, + **kwargs: Any, + ) -> ChatResponse | AsyncGenerator[ChatResponse, None]: + """Call the underlying API. Subclasses must implement this method. + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list[Msg]`): + The messages to send to the model. + tools (`list[dict] | None`, optional): + The tools available to the model. + tool_choice (`ToolChoice | None`, optional): + The tool choice mode or function name. + **kwargs: + Additional keyword arguments for the underlying API. + """ + + def _validate_tool_choice( + self, + tool_choice: ToolChoice | None, + tools: list[dict] | None, + ) -> None: + """Validate tool_choice parameter. + + Args: + tool_choice (`ToolChoice | None`): + Tool choice with ``mode`` and optional ``tools`` fields. + tools (`list[dict] | None`): + Available tools list. + + Raises: + `ValueError`: + If mode or tool names are invalid. + """ + if tool_choice is None: + return + + mode = tool_choice.mode + available_functions = [ + tool["function"]["name"] for tool in (tools or []) + ] + + tool_names = tool_choice.tools + if tool_names is not None: + for name in tool_names: + if name not in available_functions: + raise ValueError( + f"Invalid tool name '{name}' in tool_choice.tools. " + f"Available tools: " + f"{', '.join(sorted(available_functions))}", + ) + + if mode not in _TOOL_CHOICE_LITERAL_MODES: + # mode is a specific tool name — validate it exists + # Fall back to all available tools when tool_names is empty or None + validation_scope = ( + tool_names if tool_names else available_functions + ) + if mode not in validation_scope: + raise ValueError( + f"Invalid tool name '{mode}' in tool_choice.mode. " + + ( + f"Available tools in tool_choice.tools: " + f"{', '.join(sorted(tool_names))}" + if tool_names is not None + else f"Available tools: " + f"{', '.join(sorted(available_functions))}" + ), + ) + + async def count_tokens( + self, + messages: list[Msg], + tools: list[dict] | None, + ) -> int: + """A quick and unified method to estimate the token count of the + model input by dividing the total input size in bytes by 4. + + Note a standard way to count the tokens is first formatting the input + messages into the API required format, then use the tokenizer of the + underlying API to count the tokens. + + Subclasses may override this method to provide a more accurate + implementation tailored to their specific tokenizer. + + Args: + messages (`list[Msg]`): + The messages to send to the model. + tools (`list[dict] | None`): + The tools available to the model. + + Returns: + `int`: + The number of tokens in the model. + """ + cnt = 0 + + acc_texts = [] + data_blocks = [] + for msg in messages: + for block in msg.get_content_blocks(): + if isinstance(block, TextBlock): + acc_texts.append(block.text) + + elif isinstance(block, ThinkingBlock): + acc_texts.append(block.thinking) + + elif isinstance(block, HintBlock): + # ``hint`` may be a plain string or a list of + # ``TextBlock`` / ``DataBlock`` for multimodal + # content; mirror the ``ToolResultBlock.output`` + # branching above. + if isinstance(block.hint, str): + acc_texts.append(block.hint) + else: + for item in block.hint: + if isinstance(item, TextBlock): + acc_texts.append(item.text) + elif isinstance(item, DataBlock): + data_blocks.append(item) + + elif isinstance(block, ToolCallBlock): + acc_texts.append(block.input) + + elif isinstance(block, ToolResultBlock): + if isinstance(block.output, str): + acc_texts.append(block.output) + elif isinstance(block.output, list): + for item in block.output: + if isinstance(item, TextBlock): + acc_texts.append(item.text) + elif isinstance(item, DataBlock): + data_blocks.append(item) + + elif isinstance(block, DataBlock): + data_blocks.append(block) + + else: + logger.warning( + "Unknown block type %s in token counting, skipping.", + type(block), + ) + + # Count the tokens of the tool JSON schemas + if tools: + acc_texts.append(json.dumps(tools, ensure_ascii=False)) + + # Add the multimodal tokens + for block in data_blocks: + if isinstance(block.source, URLSource): + # We don't download the content here to avoid blocking + acc_texts.append(str(block.source.url)) + elif isinstance(block.source, Base64Source): + cnt += len(block.source.data) // 4 + + # Count the text tokens + acc_text = "".join(acc_texts) + cnt += int(len(acc_text.encode("utf-8")) / 4 + 0.5) + + return cnt + + async def generate_structured_output( + self, + messages: list[Msg], + structured_model: Type[BaseModel] | dict, + **kwargs: Any, + ) -> StructuredResponse: + """Generate required structured output by the given model. + + Shares the same retry settings (``max_retries``, ``retry_delay``, and + ``_get_retryable_exceptions()``) as the ``__call__`` method. + + Args: + messages (`list[Msg]`): + The context for LLM to generate the structured output. + structured_model (`Type[BaseModel] | dict`): + A Pydantic model or a dict of JSON schemas. + + Returns: + `StructuredResponse`: + The structured response generated by the model. + """ + + if len(messages) == 0: + raise ValueError( + "The input messages cannot be empty for the " + "`generate_structured_output` method.", + ) + + retryable = tuple(self._get_retryable_exceptions()) + last_error: Exception | None = None + for attempt in range(self.max_retries + 1): + try: + return await self._call_api_with_structured_output( + self.model, + messages=messages, + structured_model=structured_model, + **kwargs, + ) + except Exception as e: + if not isinstance(e, retryable): + raise + last_error = e + if attempt < self.max_retries: + logger.warning( + "Attempt %d failed for model %s: %s. " + "Retrying in %.1fs...", + attempt + 1, + self.model, + str(e), + self.retry_delay, + ) + await asyncio.sleep(self.retry_delay) + else: + logger.warning( + "All %d attempt(s) failed for model %s.", + self.max_retries + 1, + self.model, + ) + if last_error is not None: + raise last_error + raise RuntimeError( + f"Failed to generate structured output after " + f"{self.max_retries + 1} retries.", + ) + + async def _call_api_with_structured_output( + self, + model_name: str, + messages: list[Msg], + structured_model: Type[BaseModel] | dict, + tool_choice: ToolChoice | None = None, + **kwargs: Any, + ) -> StructuredResponse: + """This function constructs a 'generate_structured_output' tool to + help LLM generate structured output as a compromise for LLM APIs that + don't support structured output. + + If your subclasses inherit from `ChatModelBase` and the underlying + API supports structured output, you can override this method to + provide a more accurate implementation. + + Note by default this method forces LLM to call the + 'generate_structured_output' tool via tool_choice, and adds + instructions into the input messages. Subclasses whose underlying + API rejects forced tool_choice in certain modes (e.g. DashScope in + thinking mode) can pass ``tool_choice=ToolChoice(mode="auto")`` and + rely solely on the injected system-reminder prompt. LLM APIs that + don't support "required" tool choice may still fail (e.g. generate + text output and ignore the tool call, or fail in validation). + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list[Msg]`): + The context for the LLM to generate the structured output. + structured_model (`Type[BaseModel] | dict`): + A Pydantic model class or a JSON schema dict describing the + required output structure. + tool_choice (`ToolChoice | None`, defaults to `None`): + The tool_choice forwarded to ``_call_api``. When ``None``, + defaults to forcing the ``generate_structured_output`` tool. + **kwargs (`Any`): + Additional keyword arguments forwarded to ``_call_api``. + """ + + if isinstance(structured_model, dict): + input_schema = structured_model + else: + input_schema = structured_model.model_json_schema() + + func_name = "generate_structured_output" + if tool_choice is None: + tool_choice = ToolChoice(mode=func_name) + instruction = ( + "Now you **MUST** call the tool named " + f"'{func_name}' to generate the structured output required " + "by the user. DON'T do anything else." + ) + + copied_messages = deepcopy(messages) + # Insert instruction to ensure llm is correctly guided + if copied_messages[-1].role == "user": + # Insert a user message to the last + copied_messages[-1].content = copied_messages[ + -1 + ].get_content_blocks() + [TextBlock(text=instruction)] + else: + copied_messages.append( + UserMsg(name="user", content=[TextBlock(text=instruction)]), + ) + + res = await self._call_api( + model_name=model_name, + messages=copied_messages, + tools=[ + { + "type": "function", + "function": { + "name": func_name, + "description": "Call this function to generate " + "structured output required by " + "the user.", + "parameters": input_schema, + }, + }, + ], + tool_choice=tool_choice, + **kwargs, + ) + + completed_response: ChatResponse | None = None + if self.stream: + async for chunk in res: + if chunk.is_last: + completed_response = chunk + else: + completed_response = res + + if completed_response is None: + raise RuntimeError( + f"Failed to get the completed response from model " + f"{model_name}.", + ) + + structured_output: dict[str, Any] | None = None + for _ in completed_response.content: + if isinstance(_, ToolCallBlock) and _.name == func_name: + structured_output = _json_loads_with_repair( + _.input, + input_schema, + ) + break + + if structured_output is None: + raise RuntimeError( + "Failed to generate structured output for model.", + ) + + # Validate the output + if isinstance(structured_model, dict): + jsonschema.validate(structured_output, structured_model) + + elif issubclass(structured_model, BaseModel): + structured_model.model_validate(structured_output) + + else: + raise ValueError( + "The structured_model is expected to be a subclass of " + "Pydantic.BaseModel or a dict, " + f"but got {type(structured_model)}.", + ) + + return StructuredResponse( + id=completed_response.id, + created_at=completed_response.created_at, + content=structured_output, + usage=completed_response.usage, + ) diff --git a/src/agentscope/model/_dashscope/__init__.py b/src/agentscope/model/_dashscope/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e51813c45eba3a372e2a1c6d6f5495c189da941a --- /dev/null +++ b/src/agentscope/model/_dashscope/__init__.py @@ -0,0 +1,9 @@ +# -*- coding: utf-8 -*- +"""The DashScope API modules.""" + +from ._model import DashScopeChatModel, DashScopeCredential + +__all__ = [ + "DashScopeChatModel", + "DashScopeCredential", +] diff --git a/src/agentscope/model/_dashscope/_model.py b/src/agentscope/model/_dashscope/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..34202fff38b58cf8c13e1eef14859ccc7016917d --- /dev/null +++ b/src/agentscope/model/_dashscope/_model.py @@ -0,0 +1,655 @@ +# -*- coding: utf-8 -*- +"""The DashScope chat model class (OpenAI-compatible implementation).""" +import base64 +import io +import warnings +import wave +from collections import OrderedDict +from datetime import datetime +from typing import Any, AsyncGenerator, List, Literal, Type, TYPE_CHECKING + +from pydantic import BaseModel, Field + +from ..._utils._audio import _build_streaming_wav_header +from ..._utils._common import _generate_id +from .._base import ChatModelBase, _TOOL_CHOICE_LITERAL_MODES +from .._model_response import ChatResponse, StructuredResponse +from .._model_usage import ChatUsage +from ...credential import DashScopeCredential +from ...formatter import FormatterBase, DashScopeChatFormatter +from ...message import ( + Msg, + TextBlock, + ThinkingBlock, + ToolCallBlock, + DataBlock, + Base64Source, +) +from ...tool import ToolChoice + +if TYPE_CHECKING: + from openai.types.chat import ChatCompletion + from openai import AsyncStream +else: + ChatCompletion = Any + AsyncStream = Any + + +class DashScopeChatModel(ChatModelBase): + """The DashScope chat model (OpenAI-compatible implementation). + + This implementation uses the OpenAI Python SDK to call DashScope's + OpenAI-compatible endpoint (``compatible-mode/v1``), which supports + both text-only and multimodal (image/video) inputs through the same + unified API. + """ + + class Parameters(BaseModel): + """The parameters for DashScope LLM API.""" + + max_tokens: int | None = Field( + default=None, + title="Max Tokens", + description="The maximum number of tokens for the LLM output.", + gt=0, + ) + + thinking_enable: bool = Field( + default=False, + title="Thinking", + description="The thinking enable for the LLM output.", + ) + + thinking_budget: int | None = Field( + default=None, + title="Thinking Budget", + description="The thinking budget for the LLM output.", + gt=0, + ) + + temperature: float | None = Field( + default=None, + title="Temperature", + description="The temperature for the LLM output.", + ge=0, + lt=2, + ) + + top_p: float | None = Field( + default=None, + title="Top P", + description="The top P value for the LLM output.", + gt=0, + le=1, + ) + + top_k: int | None = Field( + default=None, + title="Top K", + description="The top K value for the LLM output.", + gt=0, + le=100, + ) + + parallel_tool_calls: bool = Field( + default=True, + title="Parallel Tool Calls", + description="If enable parallel tool calls for the LLM output.", + ) + + voice: str | None = Field( + default=None, + title="Voice", + description=( + "Voice for audio output on omni-style models (e.g. " + "``qwen3.5-omni-plus``). Setting this implicitly asks the " + "model to speak its response — ``modalities`` is filled in " + "automatically. Supported voices vary by model — see the " + "model card's ``voice.suggestions``. Any value the API " + "accepts works — the suggestions are convenience-only. " + "Leave unset for text-only " + "responses." + ), + ) + + type: Literal["dashscope_chat"] = "dashscope_chat" + """The type of the chat model.""" + + def __init__( + self, + credential: DashScopeCredential, + model: str, + parameters: "DashScopeChatModel.Parameters | None" = None, + stream: bool = True, + max_retries: int = 3, + retry_delay: float = 1.0, + context_size: int = 131072, + formatter: FormatterBase | None = None, + client_kwargs: dict[str, Any] | None = None, + ) -> None: + """Initialize the DashScope chat model. + + Args: + credential (`DashScopeCredential`): + The DashScope credential used to authenticate API calls. + model (`str`): + The DashScope model name, e.g. ``qwen-plus``. + parameters (`DashScopeChatModel.Parameters | None`, defaults to \ + `None`): + The DashScope API parameters. When ``None``, the default + parameters will be used. + stream (`bool`, defaults to `True`): + Whether to enable streaming output. + max_retries (`int`, defaults to `3`): + The maximum number of retries for the DashScope API. + retry_delay (`float`, defaults to `1.0`): + Seconds to sleep between retry attempts. + context_size (`int`, defaults to `131072`): + The model context size used for context compression. + formatter (`FormatterBase | None`, defaults to `None`): + The formatter that converts ``Msg`` objects to the format + required by the DashScope API. When ``None``, a + ``DashScopeChatFormatter`` instance will be used. + client_kwargs (`dict[str, Any] | None`, defaults to `None`): + Extra keyword arguments forwarded to ``openai.AsyncClient`` + (e.g. ``timeout``, ``default_headers``, ``http_client``). + """ + super().__init__( + credential=credential, + model=model, + parameters=parameters or self.Parameters(), + stream=stream, + max_retries=max_retries, + retry_delay=retry_delay, + context_size=context_size, + ) + self.formatter = formatter or DashScopeChatFormatter() + self.client_kwargs = client_kwargs or {} + + @classmethod + def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]: + import openai + + return ( + openai.APIConnectionError, + openai.APITimeoutError, + openai.RateLimitError, + openai.InternalServerError, + ) + + async def _call_api( + self, + model_name: str, + messages: list[Msg], + tools: list[dict] | None = None, + tool_choice: ToolChoice | None = None, + **kwargs: Any, + ) -> ChatResponse | AsyncGenerator[ChatResponse, None]: + """Call the DashScope chat completions API via OpenAI-compatible + endpoint. + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list`): + The Msg objects that will be formatted and sent to the API. + tools (`list[dict] | None`, default `None`): + The tools JSON schemas that the model can use. + tool_choice (`ToolChoice | None`, default `None`): + Controls which (if any) tool is called by the model. + **kwargs (`Any`): + The keyword arguments for DashScope chat completions API, + e.g. ``temperature``, ``max_tokens``, ``top_p``, etc. + """ + import openai + + client = openai.AsyncClient( + **{ + "api_key": self.credential.api_key.get_secret_value(), + "base_url": self.credential.base_url, + **self.client_kwargs, + }, + ) + + formatted_messages = await self.formatter.format(messages) + + request_kwargs: dict[str, Any] = { + "model": model_name, + "messages": formatted_messages, + "stream": self.stream, + } + + if self.parameters.max_tokens is not None: + request_kwargs["max_tokens"] = self.parameters.max_tokens + + if self.parameters.temperature is not None: + request_kwargs["temperature"] = self.parameters.temperature + + if self.parameters.top_p is not None: + request_kwargs["top_p"] = self.parameters.top_p + + if self.parameters.voice is not None: + # Requesting audio output implies ``modalities`` must include + # ``"audio"``; set it automatically so callers don't have to. + # ``format`` is forced to ``pcm16``: omni streaming delivers raw + # PCM upstream regardless of the requested format, and we wrap + # it as WAV in ``_parse_stream_response`` before yielding. + request_kwargs["audio"] = { + "voice": self.parameters.voice, + "format": "pcm16", + } + request_kwargs["modalities"] = ["text", "audio"] + + request_kwargs.update(kwargs) + + fmt_tools, fmt_tool_choice = self._format_tools(tools, tool_choice) + if fmt_tools is not None: + request_kwargs["tools"] = fmt_tools + if not self.parameters.parallel_tool_calls: + request_kwargs["parallel_tool_calls"] = False + if fmt_tool_choice is not None: + request_kwargs["tool_choice"] = fmt_tool_choice + + extra_body: dict[str, Any] = {} + if self.parameters.thinking_enable is not None: + extra_body["enable_thinking"] = self.parameters.thinking_enable + if self.parameters.thinking_budget is not None: + extra_body["thinking_budget"] = self.parameters.thinking_budget + if self.parameters.top_k is not None: + extra_body["top_k"] = self.parameters.top_k + + if extra_body: + request_kwargs.setdefault("extra_body", {}) + request_kwargs["extra_body"].update(extra_body) + + if self.stream: + request_kwargs["stream_options"] = {"include_usage": True} + + start_datetime = datetime.now() + response = await client.chat.completions.create(**request_kwargs) + + if self.stream: + return self._parse_stream_response(start_datetime, response) + + return self._parse_completion_response(start_datetime, response) + + async def _parse_stream_response( + self, + start_datetime: datetime, + response: AsyncStream, + ) -> AsyncGenerator[ChatResponse, None]: + """Parse the DashScope streaming response (OpenAI-compatible format). + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`AsyncStream`): + The OpenAI-compatible async stream object. + + Yields: + `ChatResponse`: + Incremental ``ChatResponse`` objects with ``is_last=False`` + followed by a final one with ``is_last=True``. + """ + usage = None + response_id: str | None = None + acc_text = TextBlock(text="") + acc_thinking = ThinkingBlock(thinking="") + acc_tool_calls: OrderedDict = OrderedDict() + # Raw PCM bytes accumulated across chunks. Storing the decoded form + # (rather than concatenated base64 strings) avoids the risk of + # corrupting the byte stream when an intermediate chunk happens to + # carry base64 padding (``=``). + acc_audio_data: bytearray = bytearray() + audio_block_id: str | None = None + # ``True`` once the first audio chunk has been prefixed with a + # streaming WAV header and yielded. + audio_header_sent: bool = False + + async with response as stream: + async for chunk in stream: + if chunk.usage: + u = chunk.usage + ptd = getattr(u, "prompt_tokens_details", None) + if ptd and hasattr(ptd, "cached_tokens"): + cache_read = ptd.cached_tokens or 0 + else: + cache_read = 0 + usage = ChatUsage( + input_tokens=u.prompt_tokens or 0, + output_tokens=u.completion_tokens or 0, + time=(datetime.now() - start_datetime).total_seconds(), + cache_input_tokens=cache_read, + ) + + response_id = response_id or getattr(chunk, "id", None) + + if not chunk.choices: + continue + + choice = chunk.choices[0] + delta = choice.delta + + delta_thinking = ( + getattr(delta, "reasoning_content", None) or "" + ) + delta_text = getattr(delta, "content", None) or "" + + # Collect audio output from Omni models (delta.audio.data). + # Upstream sends raw PCM (24kHz, 16-bit mono); we prefix the + # first chunk with a streaming WAV header so the frontend + # can start playback immediately rather than waiting for + # end-of-stream. + delta_audio = getattr(delta, "audio", None) + delta_audio_block: DataBlock | None = None + if delta_audio is not None: + if isinstance(delta_audio, dict): + audio_chunk = delta_audio.get("data", "") + else: + audio_chunk = getattr(delta_audio, "data", "") or "" + if audio_chunk: + if audio_block_id is None: + audio_block_id = _generate_id() + pcm_bytes = base64.b64decode(audio_chunk) + acc_audio_data += pcm_bytes + if not audio_header_sent: + payload = _build_streaming_wav_header() + pcm_bytes + audio_header_sent = True + else: + payload = pcm_bytes + delta_audio_block = DataBlock( + id=audio_block_id, + source=Base64Source( + data=base64.b64encode(payload).decode( + "ascii", + ), + media_type="audio/wav", + ), + ) + + acc_thinking.thinking += delta_thinking + acc_text.text += delta_text + + delta_tool_call_blocks: List[ToolCallBlock] = [] + for tool_call in getattr(delta, "tool_calls", None) or []: + idx = tool_call.index + args = ( + tool_call.function.arguments + if tool_call.function + else "" + ) or "" + if idx in acc_tool_calls: + acc_tool_calls[idx]["input"] += args + else: + acc_tool_calls[idx] = { + "id": tool_call.id or "", + "name": ( + tool_call.function.name + if tool_call.function + else "" + ), + "input": args, + } + tc = acc_tool_calls[idx] + delta_tool_call_blocks.append( + ToolCallBlock( + id=tc["id"], + name=tc["name"], + input=args, + ), + ) + + delta_contents: List[ + TextBlock | ToolCallBlock | ThinkingBlock | DataBlock + ] = [] + if delta_thinking: + delta_contents.append( + ThinkingBlock( + id=acc_thinking.id, + thinking=delta_thinking, + ), + ) + if delta_text: + delta_contents.append( + TextBlock(id=acc_text.id, text=delta_text), + ) + delta_contents.extend(delta_tool_call_blocks) + if delta_audio_block is not None: + delta_contents.append(delta_audio_block) + + if delta_contents: + _kwargs: dict[str, Any] = { + "content": delta_contents, + "usage": usage, + "is_last": False, + } + if response_id: + _kwargs["id"] = response_id + yield ChatResponse(**_kwargs) + + final_contents: List[ + TextBlock | ToolCallBlock | ThinkingBlock | DataBlock + ] = [] + if acc_thinking.thinking: + final_contents.append(acc_thinking) + if acc_text.text: + final_contents.append(acc_text) + for tc in acc_tool_calls.values(): + final_contents.append( + ToolCallBlock( + id=tc["id"], + name=tc["name"], + input=tc["input"], + ), + ) + if acc_audio_data: + # PCM bytes were already streamed incrementally above (first + # chunk prefixed with a WAV header). Here we also assemble a + # standalone fixed-size WAV and attach it to the ``is_last`` + # chunk so callers that consume the model directly (i.e. + # without going through ``Agent``, which filters audio blocks + # out of context) get a self-contained audio block for + # downstream serialization / display. + buf = io.BytesIO() + with wave.open(buf, "wb") as wav: + wav.setnchannels(1) + wav.setsampwidth(2) + wav.setframerate(24000) + wav.writeframes(bytes(acc_audio_data)) + final_contents.append( + DataBlock( + id=audio_block_id, + source=Base64Source( + data=base64.b64encode(buf.getvalue()).decode("ascii"), + media_type="audio/wav", + ), + ), + ) + + _final_kwargs: dict[str, Any] = { + "content": final_contents, + "usage": usage, + "is_last": True, + } + if response_id: + _final_kwargs["id"] = response_id + yield ChatResponse(**_final_kwargs) + + def _parse_completion_response( + self, + start_datetime: datetime, + response: ChatCompletion, + ) -> ChatResponse: + """Parse the DashScope non-streaming response (OpenAI-compatible + format). + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`ChatCompletion`): + The OpenAI-compatible chat completion object. + + Returns: + `ChatResponse`: + A single ``ChatResponse`` with ``is_last=True``. + """ + content_blocks: List[TextBlock | ToolCallBlock | ThinkingBlock] = [] + + if response.choices: + choice = response.choices[0] + reasoning = getattr(choice.message, "reasoning_content", None) + if isinstance(reasoning, str) and reasoning: + content_blocks.append(ThinkingBlock(thinking=reasoning)) + + if choice.message.content: + content_blocks.append(TextBlock(text=choice.message.content)) + + for tool_call in choice.message.tool_calls or []: + content_blocks.append( + ToolCallBlock( + id=tool_call.id, + name=tool_call.function.name, + input=tool_call.function.arguments, + ), + ) + + usage = None + if response.usage: + u = response.usage + ptd = getattr(u, "prompt_tokens_details", None) + if ptd and hasattr(ptd, "cached_tokens"): + cache_read = ptd.cached_tokens or 0 + else: + cache_read = 0 + usage = ChatUsage( + input_tokens=u.prompt_tokens, + output_tokens=u.completion_tokens, + time=(datetime.now() - start_datetime).total_seconds(), + cache_input_tokens=cache_read, + ) + + resp_kwargs: dict[str, Any] = { + "content": content_blocks, + "is_last": True, + "usage": usage, + } + response_id = getattr(response, "id", None) + if response_id: + resp_kwargs["id"] = response_id + + return ChatResponse(**resp_kwargs) + + def _format_tools( + self, + tools: list[dict] | None, + tool_choice: ToolChoice | None, + ) -> tuple[list[dict] | None, str | dict | None]: + """Validate and format tools and tool_choice for DashScope. + + DashScope supports "auto", "none", and "required" modes in + OpenAI-compatible format. When ``tool_choice.tools`` is specified + the schemas list is filtered to only those tools. When + ``tool_choice.mode`` is a specific tool name (str) the model is + forced to call exactly that tool. + + Args: + tools (`list[dict] | None`, optional): + The raw tool schemas. + tool_choice (`ToolChoice | None`, optional): + The tool choice configuration. + + Returns: + `tuple[list[dict] | None, str | dict | None]`: + A tuple of (formatted_tools, formatted_tool_choice). + """ + if tool_choice and tools: + self._validate_tool_choice(tool_choice, tools) + if tool_choice.tools: + allowed = set(tool_choice.tools) + tools = [t for t in tools if t["function"]["name"] in allowed] + + fmt_tools = None + if tools: + for value in tools: + if ( + not isinstance(value, dict) + or "type" not in value + or value["type"] != "function" + or "function" not in value + ): + raise ValueError( + f"Each schema must be a dict with 'type' as " + f"'function' and 'function' key, got {value}", + ) + fmt_tools = tools + + if not tool_choice: + return fmt_tools, None + + mode = tool_choice.mode + + if mode not in _TOOL_CHOICE_LITERAL_MODES: + return fmt_tools, { + "type": "function", + "function": {"name": mode}, + } + + if mode == "required": + warnings.warn( + f"'{mode}' is not fully supported by DashScope API. " + "It will be converted to 'auto'.", + DeprecationWarning, + stacklevel=2, + ) + return fmt_tools, "auto" + + return fmt_tools, mode + + async def _call_api_with_structured_output( + self, + model_name: str, + messages: list[Msg], + structured_model: Type[BaseModel] | dict, + tool_choice: ToolChoice | None = None, + **kwargs: Any, + ) -> StructuredResponse: + """DashScope-specific override for structured output. + + DashScope rejects ``tool_choice="required"`` or an object-form + ``tool_choice`` when thinking mode is enabled. In that case we + default ``tool_choice`` to ``"auto"`` and rely on the base class's + injected system-reminder prompt to guide the model. When thinking + is disabled, this falls through to the base implementation. + + See: https://help.aliyun.com/en/model-studio/qwen-function-calling + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list[Msg]`): + The context for the LLM to generate the structured output. + structured_model (`Type[BaseModel] | dict`): + A Pydantic model class or a JSON schema dict describing the + required output structure. + tool_choice (`ToolChoice | None`, defaults to `None`): + The tool_choice forwarded to ``_call_api``. When ``None`` + and thinking mode is enabled, it is downgraded to + ``ToolChoice(mode="auto")``; otherwise the base default + (force the structured-output tool) is used. + **kwargs (`Any`): + Additional keyword arguments forwarded to ``_call_api``. + + Returns: + `StructuredResponse`: + The structured response whose ``content`` is the validated + output dict matching ``structured_model``. + """ + if tool_choice is None and self.parameters.thinking_enable: + tool_choice = ToolChoice(mode="auto") + return await super()._call_api_with_structured_output( + model_name=model_name, + messages=messages, + structured_model=structured_model, + tool_choice=tool_choice, + **kwargs, + ) diff --git a/src/agentscope/model/_dashscope/_models/qwen-long.yaml b/src/agentscope/model/_dashscope/_models/qwen-long.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0cb69744deaf47c4e27ca6235905bc281e4bca2a --- /dev/null +++ b/src/agentscope/model/_dashscope/_models/qwen-long.yaml @@ -0,0 +1,21 @@ +name: qwen-long +label: Qwen Long +status: active + +input_types: + - text/plain + +output_types: + - text/plain + +context_size: 10000000 +output_size: 8192 + +parameter_overrides: + max_tokens: {"maximum": 8192} + thinking_enable: + hidden: true + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_dashscope/_models/qwen-plus.yaml b/src/agentscope/model/_dashscope/_models/qwen-plus.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6be18321d8873cf22eccc1ea2c007ac8a0707528 --- /dev/null +++ b/src/agentscope/model/_dashscope/_models/qwen-plus.yaml @@ -0,0 +1,20 @@ +name: qwen-plus +label: Qwen Plus +status: active + +input_types: + - text/plain + +output_types: + - text/plain + - application/x-thinking + +context_size: 1000000 +output_size: 16384 + +parameter_overrides: + max_tokens: {"maximum": 16384} + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_dashscope/_models/qwen3.5-omni-plus.yaml b/src/agentscope/model/_dashscope/_models/qwen3.5-omni-plus.yaml new file mode 100644 index 0000000000000000000000000000000000000000..95ea70f9c8adbe91978574db31aef9e0257da56a --- /dev/null +++ b/src/agentscope/model/_dashscope/_models/qwen3.5-omni-plus.yaml @@ -0,0 +1,38 @@ +name: qwen3.5-omni-plus +label: Qwen3.5-Omni-Plus +status: active + +input_types: + - text/plain + - image/bmp + - image/jpeg + - image/png + - image/tiff + - image/webp + - image/heic + - audio/mp3 + - audio/wav + - audio/flac + - video/mp4 + +output_types: + - text/plain + - application/x-thinking + - audio/wav + +context_size: 262144 +output_size: 65536 + +parameter_overrides: + max_tokens: + maximum: 65536 + # Source: https://help.aliyun.com/zh/model-studio/omni-voice-list + voice: + default: Tina + enum: + - Tina + - Cherry + - Ethan + - Serena + - Chelsie + - Cindy diff --git a/src/agentscope/model/_dashscope/_models/qwen3.5-plus.yaml b/src/agentscope/model/_dashscope/_models/qwen3.5-plus.yaml new file mode 100644 index 0000000000000000000000000000000000000000..124bc789cfba0905fdeccf6bc69f1e1e326757d3 --- /dev/null +++ b/src/agentscope/model/_dashscope/_models/qwen3.5-plus.yaml @@ -0,0 +1,28 @@ +name: qwen3.5-plus +label: Qwen3.5-Plus +status: active + +input_types: + - text/plain + - application/x-thinking + - image/bmp + - image/jpeg + - image/png + - image/tiff + - image/webp + - image/heic + - video/mp4 + +output_types: + - text/plain + - application/x-thinking + +context_size: 1000000 +output_size: 65536 + +parameter_overrides: + max_tokens: {"maximum": 65536} + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_dashscope/_models/qwen3.6-max-preview.yaml b/src/agentscope/model/_dashscope/_models/qwen3.6-max-preview.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c017d8a7dbca8016b1e3e093721023194697b5f4 --- /dev/null +++ b/src/agentscope/model/_dashscope/_models/qwen3.6-max-preview.yaml @@ -0,0 +1,20 @@ +name: qwen3.6-max-preview +label: Qwen3.6-Max-Preview +status: active + +input_types: + - text/plain + +output_types: + - text/plain + - application/x-thinking + +context_size: 240000 +output_size: 64000 + +parameter_overrides: + max_tokens: {"maximum": 64000} + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_dashscope/_models/qwen3.6-plus.yaml b/src/agentscope/model/_dashscope/_models/qwen3.6-plus.yaml new file mode 100644 index 0000000000000000000000000000000000000000..947ffbde5e379fc50eb2fe61c41346a781bec3a6 --- /dev/null +++ b/src/agentscope/model/_dashscope/_models/qwen3.6-plus.yaml @@ -0,0 +1,28 @@ +name: qwen3.6-plus +label: Qwen3.6-Plus +status: active + +input_types: + - text/plain + - application/x-thinking + - image/bmp + - image/jpeg + - image/png + - image/tiff + - image/webp + - image/heic + - video/mp4 + +output_types: + - text/plain + - application/x-thinking + +context_size: 1000000 +output_size: 65536 + +parameter_overrides: + max_tokens: {"maximum": 65536} + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_dashscope/_models/qwen3.7-max.yaml b/src/agentscope/model/_dashscope/_models/qwen3.7-max.yaml new file mode 100644 index 0000000000000000000000000000000000000000..086dbdfacceeae57831b56e731373c88f1798987 --- /dev/null +++ b/src/agentscope/model/_dashscope/_models/qwen3.7-max.yaml @@ -0,0 +1,20 @@ +name: qwen3.7-max +label: Qwen Max 3.7 +status: active + +input_types: + - text/plain + +output_types: + - text/plain + - application/x-thinking + +context_size: 1000000 +output_size: 65536 + +parameter_overrides: + max_tokens: {"maximum": 65536} + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_deepseek/__init__.py b/src/agentscope/model/_deepseek/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f073872ebbaba3dccb1eac7902cb861993f1daec --- /dev/null +++ b/src/agentscope/model/_deepseek/__init__.py @@ -0,0 +1,9 @@ +# -*- coding: utf-8 -*- +"""The DeepSeek LLM API modules.""" + +from ._model import DeepSeekCredential, DeepSeekChatModel + +__all__ = [ + "DeepSeekCredential", + "DeepSeekChatModel", +] diff --git a/src/agentscope/model/_deepseek/_model.py b/src/agentscope/model/_deepseek/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..16e2baf1c8d5fab7b552fed9ede7244fe492e9fe --- /dev/null +++ b/src/agentscope/model/_deepseek/_model.py @@ -0,0 +1,499 @@ +# -*- coding: utf-8 -*- +"""The DeepSeek chat model implementation.""" +from collections import OrderedDict +from datetime import datetime +from typing import Literal, Any, AsyncGenerator, TYPE_CHECKING, List, Type + +from pydantic import BaseModel, Field + +from .._base import ChatModelBase, _TOOL_CHOICE_LITERAL_MODES +from .._model_response import ChatResponse, StructuredResponse +from .._model_usage import ChatUsage +from ...credential import DeepSeekCredential +from ...formatter import FormatterBase, DeepSeekChatFormatter +from ...message import Msg, ThinkingBlock, ToolCallBlock, TextBlock +from ...tool import ToolChoice + +if TYPE_CHECKING: + from openai.types.chat import ChatCompletion + from openai import AsyncStream +else: + ChatCompletion = Any + AsyncStream = Any + + +class DeepSeekChatModel(ChatModelBase): + """The DeepSeek chat model.""" + + class Parameters(BaseModel): + """The parameters for the DeepSeek chat model.""" + + max_tokens: int | None = Field( + default=None, + title="Max Tokens", + description="The maximum number of tokens for the LLM output.", + gt=0, + ) + + thinking_enable: bool = Field( + default=False, + title="Thinking", + description=( + "Whether to enable thinking mode. When enabled, the model " + "outputs a chain-of-thought reasoning before the final " + "answer via the reasoning_content field." + ), + ) + + reasoning_effort: Literal["high", "max"] | None = Field( + default=None, + title="Reasoning Effort", + description=( + "Controls the depth of reasoning in thinking mode. " + "Supported values: high (default), max. " + "For compatibility, low/medium map to high, " + "xhigh maps to max." + ), + ) + + temperature: float | None = Field( + default=None, + title="Temperature", + description="The temperature for the LLM output.", + ge=0, + le=2, + ) + + top_p: float | None = Field( + default=None, + title="Top P", + description="The top P value for the LLM output.", + gt=0, + le=1, + ) + + type: Literal["deepseek_chat"] = "deepseek_chat" + """The type of the chat model.""" + + def __init__( + self, + credential: DeepSeekCredential, + model: str, + parameters: "DeepSeekChatModel.Parameters | None" = None, + stream: bool = True, + max_retries: int = 3, + retry_delay: float = 1.0, + context_size: int = 65536, + formatter: FormatterBase | None = None, + client_kwargs: dict[str, Any] | None = None, + ) -> None: + """Initialize the DeepSeek chat model. + + Args: + credential (`DeepSeekCredential`): + The DeepSeek credential used to authenticate API calls. + model (`str`): + The DeepSeek model name, e.g. ``deepseek-chat``. + parameters (`DeepSeekChatModel.Parameters | None`, defaults to \ + `None`): + The DeepSeek API parameters. When ``None``, the default + parameters will be used. + stream (`bool`, defaults to `True`): + Whether to enable streaming output. + max_retries (`int`, defaults to `3`): + The maximum number of retries for the DeepSeek API. + retry_delay (`float`, defaults to `1.0`): + Seconds to sleep between retry attempts. + context_size (`int`, defaults to `65536`): + The model context size used for context compression. + formatter (`FormatterBase | None`, defaults to `None`): + The formatter that converts ``Msg`` objects to the format + required by the DeepSeek API. When ``None``, a + ``DeepSeekChatFormatter`` instance will be used. + client_kwargs (`dict[str, Any] | None`, defaults to `None`): + Extra keyword arguments forwarded to ``openai.AsyncClient`` + (e.g. ``timeout``, ``default_headers``, ``http_client``). + """ + super().__init__( + credential=credential, + model=model, + parameters=parameters or self.Parameters(), + stream=stream, + max_retries=max_retries, + retry_delay=retry_delay, + context_size=context_size, + ) + self.formatter = formatter or DeepSeekChatFormatter() + self.client_kwargs = client_kwargs or {} + + @classmethod + def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]: + import openai + + return ( + openai.APIConnectionError, + openai.APITimeoutError, + openai.RateLimitError, + openai.InternalServerError, + ) + + async def _call_api( + self, + model_name: str, + messages: list[Msg], + tools: list[dict] | None = None, + tool_choice: ToolChoice | None = None, + **generate_kwargs: Any, + ) -> ChatResponse | AsyncGenerator[ChatResponse, None]: + """Call the DeepSeek chat completions API. + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list`): + A list of message dicts with ``role`` and ``content`` keys. + tools (`list[dict]`, default `None`): + The tools JSON schemas. + tool_choice (`ToolChoice | None`, optional): + Controls which (if any) tool is called by the model. + **generate_kwargs (`Any`): + Extra keyword arguments forwarded to the API. + + Returns: + `ChatResponse | AsyncGenerator[ChatResponse, None]`: + A ``ChatResponse`` when streaming is disabled, or an async + generator of ``ChatResponse`` objects when streaming is + enabled. + """ + import openai + + client = openai.AsyncClient( + **{ + "api_key": self.credential.api_key.get_secret_value(), + "base_url": self.credential.base_url, + **self.client_kwargs, + }, + ) + + formatted_messages = await self.formatter.format(messages) + + kwargs: dict[str, Any] = { + "model": model_name, + "messages": formatted_messages, + "stream": self.stream, + } + + if self.parameters.max_tokens is not None: + kwargs["max_tokens"] = self.parameters.max_tokens + + if self.parameters.temperature is not None: + kwargs["temperature"] = self.parameters.temperature + + if self.parameters.top_p is not None: + kwargs["top_p"] = self.parameters.top_p + + if self.parameters.reasoning_effort is not None: + kwargs["reasoning_effort"] = self.parameters.reasoning_effort + + kwargs.update(generate_kwargs) + + thinking_type = ( + "enabled" if self.parameters.thinking_enable else "disabled" + ) + kwargs.setdefault("extra_body", {}) + kwargs["extra_body"].setdefault("thinking", {}) + kwargs["extra_body"]["thinking"].setdefault("type", thinking_type) + + fmt_tools, fmt_tool_choice = self._format_tools(tools, tool_choice) + + if fmt_tools: + kwargs["tools"] = fmt_tools + + if fmt_tool_choice is not None: + kwargs["tool_choice"] = fmt_tool_choice + + if self.stream: + kwargs["stream_options"] = {"include_usage": True} + + start_datetime = datetime.now() + response = await client.chat.completions.create(**kwargs) + + if self.stream: + return self._parse_stream_response(start_datetime, response) + + return self._parse_completion_response(start_datetime, response) + + async def _parse_stream_response( + self, + start_datetime: datetime, + response: AsyncStream, + ) -> AsyncGenerator[ChatResponse, None]: + """Parse the DeepSeek streaming response. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`AsyncStream`): + The OpenAI-compatible async stream object. + + Yields: + `ChatResponse`: + Incremental ``ChatResponse`` objects with ``is_last=False`` + followed by a final one with ``is_last=True``. + """ + usage = None + response_id: str | None = None + # All delta should have the same block identifier + acc_text = TextBlock(text="") + acc_thinking = ThinkingBlock(thinking="") + acc_tool_calls: OrderedDict = OrderedDict() + + async with response as stream: + async for chunk in stream: + if chunk.usage: + u = chunk.usage + usage = ChatUsage( + input_tokens=u.prompt_tokens, + output_tokens=u.completion_tokens, + time=(datetime.now() - start_datetime).total_seconds(), + cache_input_tokens=getattr( + u, + "prompt_cache_hit_tokens", + 0, + ), + ) + + # Capture response_id from the first chunk that carries it + response_id = response_id or getattr(chunk, "id", None) + + if not chunk.choices: + continue + + choice = chunk.choices[0] + delta = choice.delta + + delta_thinking = ( + getattr(delta, "reasoning_content", None) or "" + ) + delta_text = getattr(delta, "content", None) or "" + + acc_thinking.thinking += delta_thinking + acc_text.text += delta_text + + delta_tool_call_blocks: List[ToolCallBlock] = [] + for tool_call in getattr(delta, "tool_calls", None) or []: + idx = tool_call.index + args = tool_call.function.arguments or "" + if idx in acc_tool_calls: + acc_tool_calls[idx]["input"] += args + else: + acc_tool_calls[idx] = { + "id": tool_call.id, + "name": tool_call.function.name, + "input": args, + } + tc = acc_tool_calls[idx] + delta_tool_call_blocks.append( + ToolCallBlock( + id=tc["id"], + name=tc["name"], + input=args, + ), + ) + + delta_contents: List[ + TextBlock | ToolCallBlock | ThinkingBlock + ] = [] + if delta_thinking: + delta_contents.append( + ThinkingBlock( + id=acc_thinking.id, + thinking=delta_thinking, + ), + ) + if delta_text: + delta_contents.append( + TextBlock(id=acc_text.id, text=delta_text), + ) + delta_contents.extend(delta_tool_call_blocks) + + if delta_contents: + _kwargs: dict[str, Any] = { + "content": delta_contents, + "usage": usage, + "is_last": False, + } + if response_id: + _kwargs["id"] = response_id + yield ChatResponse(**_kwargs) + + final_contents: List[TextBlock | ToolCallBlock | ThinkingBlock] = [] + if acc_thinking.thinking: + final_contents.append(acc_thinking) + if acc_text.text: + final_contents.append(acc_text) + for tc in acc_tool_calls.values(): + final_contents.append( + ToolCallBlock( + id=tc["id"], + name=tc["name"], + input=tc["input"], + ), + ) + + _final_kwargs: dict[str, Any] = { + "content": final_contents, + "usage": usage, + "is_last": True, + } + if response_id: + _final_kwargs["id"] = response_id + yield ChatResponse(**_final_kwargs) + + def _parse_completion_response( + self, + start_datetime: datetime, + response: ChatCompletion, + ) -> ChatResponse: + """Parse the DeepSeek non-streaming response. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`ChatCompletion`): + The OpenAI-compatible chat completion object. + + Returns: + `ChatResponse`: + A single ``ChatResponse`` with ``is_last=True``. + """ + content_blocks: List[TextBlock | ToolCallBlock | ThinkingBlock] = [] + + if response.choices: + choice = response.choices[0] + reasoning = getattr(choice.message, "reasoning_content", None) + if isinstance(reasoning, str) and reasoning: + content_blocks.append(ThinkingBlock(thinking=reasoning)) + + if choice.message.content: + content_blocks.append(TextBlock(text=choice.message.content)) + + for tool_call in choice.message.tool_calls or []: + content_blocks.append( + ToolCallBlock( + id=tool_call.id, + name=tool_call.function.name, + input=tool_call.function.arguments, + ), + ) + + usage = None + if response.usage: + u = response.usage + usage = ChatUsage( + input_tokens=u.prompt_tokens, + output_tokens=u.completion_tokens, + time=(datetime.now() - start_datetime).total_seconds(), + cache_input_tokens=getattr( + u, + "prompt_cache_hit_tokens", + 0, + ), + ) + + resp_kwargs: dict[str, Any] = { + "content": content_blocks, + "is_last": True, + "usage": usage, + } + response_id = getattr(response, "id", None) + if response_id: + resp_kwargs["id"] = response_id + + return ChatResponse(**resp_kwargs) + + def _format_tools( + self, + tools: list[dict] | None, + tool_choice: ToolChoice | None, + ) -> tuple[list[dict] | None, str | dict | None]: + """Validate, filter, and format tools and tool_choice for the DeepSeek + API. + + When ``tool_choice.tools`` is specified the schemas list is filtered + to only those tools. When ``tool_choice.mode`` is a specific tool name + (str) the model is forced to call exactly that tool without needing to + filter the list, preserving prompt-cache efficiency. + + Args: + tools (`list[dict] | None`, optional): + The raw tool schemas. + tool_choice (`ToolChoice | None`, optional): + The tool choice configuration. + + Returns: + `tuple[list[dict] | None, str | dict | None]`: + A tuple of (formatted_tools, formatted_tool_choice). + """ + if tool_choice and tools: + self._validate_tool_choice(tool_choice, tools) + if tool_choice.tools: + allowed = set(tool_choice.tools) + tools = [t for t in tools if t["function"]["name"] in allowed] + + if not tool_choice: + return tools, None + + mode = tool_choice.mode + + if mode not in _TOOL_CHOICE_LITERAL_MODES: + return tools, {"type": "function", "function": {"name": mode}} + + return tools, mode + + async def _call_api_with_structured_output( + self, + model_name: str, + messages: list[Msg], + structured_model: Type[BaseModel] | dict, + tool_choice: ToolChoice | None = None, + **kwargs: Any, + ) -> StructuredResponse: + """DeepSeek-specific override for structured output. + + DeepSeek rejects ``tool_choice="required"`` or an object-form + ``tool_choice`` when thinking mode is enabled. In that case we + default ``tool_choice`` to ``"auto"`` and rely on the base class's + injected system-reminder prompt to guide the model. When thinking + is disabled, this falls through to the base implementation. + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list[Msg]`): + The context for the LLM to generate the structured output. + structured_model (`Type[BaseModel] | dict`): + A Pydantic model class or a JSON schema dict describing the + required output structure. + tool_choice (`ToolChoice | None`, defaults to `None`): + The tool_choice forwarded to ``_call_api``. When ``None`` + and thinking mode is enabled, it is downgraded to + ``ToolChoice(mode="auto")``; otherwise the base default + (force the structured-output tool) is used. + **kwargs (`Any`): + Additional keyword arguments forwarded to ``_call_api``. + + Returns: + `StructuredResponse`: + The structured response whose ``content`` is the validated + output dict matching ``structured_model``. + """ + if tool_choice is None and self.parameters.thinking_enable: + tool_choice = ToolChoice(mode="auto") + return await super()._call_api_with_structured_output( + model_name=model_name, + messages=messages, + structured_model=structured_model, + tool_choice=tool_choice, + **kwargs, + ) diff --git a/src/agentscope/model/_deepseek/_models/deepseek-chat.yaml b/src/agentscope/model/_deepseek/_models/deepseek-chat.yaml new file mode 100644 index 0000000000000000000000000000000000000000..02dc5aef0fa911f91e3e26c7969942721ea9afcf --- /dev/null +++ b/src/agentscope/model/_deepseek/_models/deepseek-chat.yaml @@ -0,0 +1,17 @@ +name: deepseek-chat +label: DeepSeek Chat +status: sunset +deprecated_at: "2026-07-24T00:00:00" + +input_types: + - text/plain + +output_types: + - text/plain + +context_size: 1000000 +output_size: 384000 + +parameter_overrides: + max_tokens: + maximum: 384000 diff --git a/src/agentscope/model/_deepseek/_models/deepseek-reasoner.yaml b/src/agentscope/model/_deepseek/_models/deepseek-reasoner.yaml new file mode 100644 index 0000000000000000000000000000000000000000..337c0484902721d24ff4f7c2b2a1e939243e26bd --- /dev/null +++ b/src/agentscope/model/_deepseek/_models/deepseek-reasoner.yaml @@ -0,0 +1,18 @@ +name: deepseek-reasoner +label: DeepSeek Reasoner +status: sunset +deprecated_at: "2026-07-24T00:00:00" + +input_types: + - text/plain + +output_types: + - text/plain + - application/x-thinking + +context_size: 1000000 +output_size: 384000 + +parameter_overrides: + max_tokens: + maximum: 384000 diff --git a/src/agentscope/model/_deepseek/_models/deepseek-v4-flash.yaml b/src/agentscope/model/_deepseek/_models/deepseek-v4-flash.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8133187c5e06a2bee99b5b580190841f80f299e8 --- /dev/null +++ b/src/agentscope/model/_deepseek/_models/deepseek-v4-flash.yaml @@ -0,0 +1,18 @@ +name: deepseek-v4-flash +label: DeepSeek V4 Flash +status: active + +input_types: + - text/plain + - application/x-thinking + +output_types: + - text/plain + - application/x-thinking + +context_size: 1000000 +output_size: 384000 + +parameter_overrides: + max_tokens: + maximum: 384000 diff --git a/src/agentscope/model/_deepseek/_models/deepseek-v4-pro.yaml b/src/agentscope/model/_deepseek/_models/deepseek-v4-pro.yaml new file mode 100644 index 0000000000000000000000000000000000000000..62f78c9f7e6e3e33759539a095e5aa628f311da0 --- /dev/null +++ b/src/agentscope/model/_deepseek/_models/deepseek-v4-pro.yaml @@ -0,0 +1,18 @@ +name: deepseek-v4-pro +label: DeepSeek V4 Pro +status: active + +input_types: + - text/plain + - application/x-thinking + +output_types: + - text/plain + - application/x-thinking + +context_size: 1000000 +output_size: 384000 + +parameter_overrides: + max_tokens: + maximum: 384000 diff --git a/src/agentscope/model/_gemini/__init__.py b/src/agentscope/model/_gemini/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c82b032e4c3fe28ef8facfd2e54fddbe6b8ceda0 --- /dev/null +++ b/src/agentscope/model/_gemini/__init__.py @@ -0,0 +1,9 @@ +# -*- coding: utf-8 -*- +"""The Google Gemini LLM API modules.""" + +from ._model import GeminiCredential, GeminiChatModel + +__all__ = [ + "GeminiCredential", + "GeminiChatModel", +] diff --git a/src/agentscope/model/_gemini/_model.py b/src/agentscope/model/_gemini/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..4d6e2435f9b76755a923d11e2ba98e4296b29866 --- /dev/null +++ b/src/agentscope/model/_gemini/_model.py @@ -0,0 +1,574 @@ +# -*- coding: utf-8 -*- +"""The Google Gemini chat model implementation.""" +import base64 +import json +from datetime import datetime +from typing import Literal, Any, AsyncGenerator, TYPE_CHECKING, List, Type + +from pydantic import BaseModel, Field + +from ..._utils._common import _generate_id, _flatten_json_schema +from .._base import ChatModelBase, _TOOL_CHOICE_LITERAL_MODES +from .._model_response import ChatResponse +from .._model_usage import ChatUsage +from ...credential import GeminiCredential +from ...formatter import FormatterBase, GeminiChatFormatter +from ...message import Msg, ThinkingBlock, ToolCallBlock, TextBlock +from ...tool import ToolChoice + +if TYPE_CHECKING: + from google.genai.types import GenerateContentResponse +else: + GenerateContentResponse = Any + + +def _sanitize_schema_for_gemini(schema: Any) -> Any: + """Sanitize a JSON schema to be compatible with the Gemini API. + + Gemini API does not support certain JSON Schema constructs. This + function removes or rewrites the following: + + - ``additionalProperties``: removed entirely. + - ``anyOf`` containing a ``{"type": "null"}`` entry: simplified to + the single non-null type. If there is exactly one non-null + alternative it is inlined directly; otherwise the ``anyOf`` is + kept but the null entry is dropped. + - All nested sub-schemas (``properties``, ``items``, ``$defs``, + etc.) are processed recursively. + + Args: + schema (`Any`): + The JSON schema to sanitize. Non-dict values are returned + unchanged; lists are recursively sanitized element-wise. + + Returns: + `Any`: + A sanitized copy of the schema, or the original value if it + is not a dict or list. + """ + if not isinstance(schema, dict): + if isinstance(schema, list): + return [_sanitize_schema_for_gemini(v) for v in schema] + return schema + + schema = dict(schema) + + # Remove additionalProperties — not supported by Gemini + schema.pop("additionalProperties", None) + + # Simplify anyOf that only differs by a null type, e.g. Optional[X] + if "anyOf" in schema and isinstance(schema["anyOf"], list): + any_of = schema["anyOf"] + non_null = [v for v in any_of if v != {"type": "null"}] + if len(non_null) < len(any_of): # at least one null entry removed + if len(non_null) == 1: + # Inline the single non-null type, preserving outer keys + merged = dict(_sanitize_schema_for_gemini(non_null[0])) + for k, v in schema.items(): + if k != "anyOf": + merged.setdefault(k, v) + return merged + elif non_null: + schema["anyOf"] = [ + _sanitize_schema_for_gemini(v) for v in non_null + ] + else: + del schema["anyOf"] + + # Recursively process nested object schemas + for key in ["properties", "patternProperties", "$defs"]: + if key in schema and isinstance(schema[key], dict): + schema[key] = { + k: _sanitize_schema_for_gemini(v) + for k, v in schema[key].items() + } + + for key in ["items", "not", "if", "then", "else"]: + if key in schema: + schema[key] = _sanitize_schema_for_gemini(schema[key]) + + for key in ["allOf", "oneOf", "anyOf"]: + if key in schema and isinstance(schema[key], list): + schema[key] = [_sanitize_schema_for_gemini(v) for v in schema[key]] + + return schema + + +class GeminiChatModel(ChatModelBase): + """The Google Gemini chat model.""" + + class Parameters(BaseModel): + """The parameters for the Gemini chat model.""" + + max_tokens: int | None = Field( + default=None, + title="Max Tokens", + description="The maximum number of tokens for the LLM output.", + gt=0, + ) + + thinking_enable: bool = Field( + default=False, + title="Thinking", + description="Whether to enable thinking output.", + ) + + thinking_budget: int | None = Field( + default=None, + title="Thinking Budget", + description="The thinking budget in tokens.", + gt=0, + ) + + temperature: float | None = Field( + default=None, + title="Temperature", + description="The temperature for the LLM output.", + ge=0, + le=2, + ) + + top_p: float | None = Field( + default=None, + title="Top P", + description="The top P value for the LLM output.", + gt=0, + le=1, + ) + + type: Literal["gemini_chat"] = "gemini_chat" + """The type of the chat model.""" + + def __init__( + self, + credential: GeminiCredential, + model: str, + parameters: "GeminiChatModel.Parameters | None" = None, + stream: bool = True, + max_retries: int = 3, + retry_delay: float = 1.0, + context_size: int = 1048576, + formatter: FormatterBase | None = None, + client_kwargs: dict[str, Any] | None = None, + ) -> None: + """Initialize the Gemini chat model. + + Args: + credential (`GeminiCredential`): + The Google Gemini credential used to authenticate API calls. + model (`str`): + The Gemini model name, e.g. ``gemini-2.0-flash-exp``. + parameters (`GeminiChatModel.Parameters | None`, defaults to \ + `None`): + The Gemini API parameters. When ``None``, the default + parameters will be used. + stream (`bool`, defaults to `True`): + Whether to enable streaming output. + max_retries (`int`, defaults to `3`): + The maximum number of retries for the Gemini API. + retry_delay (`float`, defaults to `1.0`): + Seconds to sleep between retry attempts. + context_size (`int`, defaults to `1048576`): + The model context size used for context compression. + formatter (`FormatterBase | None`, defaults to `None`): + The formatter that converts ``Msg`` objects to the format + required by the Gemini API. When ``None``, a + ``GeminiChatFormatter`` instance will be used. + client_kwargs (`dict[str, Any] | None`, defaults to `None`): + Extra keyword arguments forwarded to ``google.genai.Client`` + (e.g. ``vertexai``, ``project``, ``location``, + ``credentials``, ``http_options``). + """ + super().__init__( + credential=credential, + model=model, + parameters=parameters or self.Parameters(), + stream=stream, + max_retries=max_retries, + retry_delay=retry_delay, + context_size=context_size, + ) + self.formatter = formatter or GeminiChatFormatter() + self.client_kwargs = client_kwargs or {} + + @classmethod + def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]: + from google.genai import errors + + # APIError is the common parent of ClientError (4xx) and ServerError + # (5xx). The google-genai SDK does not expose a dedicated rate-limit + # subclass, and 429 surfaces as ClientError — so we accept the wider + # set to make sure 429s are retried, at the cost of also retrying + # rare 4xx like auth/bad-request a few times. + return (errors.APIError,) + + async def _call_api( + self, + model_name: str, + messages: list[Msg], + tools: list[dict] | None = None, + tool_choice: ToolChoice | None = None, + **config_kwargs: Any, + ) -> ChatResponse | AsyncGenerator[ChatResponse, None]: + """Call the Gemini chat API. + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list`): + A list of message objects for Gemini API. + tools (`list[dict]`, default `None`): + The tools JSON schemas. + tool_choice (`ToolChoice | None`, optional): + Controls which (if any) tool is called by the model. + **config_kwargs (`Any`): + Extra keyword arguments for the Gemini config. + + Returns: + `ChatResponse | AsyncGenerator[ChatResponse, None]`: + A ``ChatResponse`` when streaming is disabled, or an async + generator of ``ChatResponse`` objects when streaming is + enabled. + """ + from google import genai + + client = genai.Client( + **{ + "api_key": self.credential.api_key.get_secret_value(), + **self.client_kwargs, + }, + ) + + formatted_messages = await self.formatter.format(messages) + + config: dict[str, Any] = {**config_kwargs} + + if self.parameters.max_tokens is not None: + config["max_output_tokens"] = self.parameters.max_tokens + + if self.parameters.temperature is not None: + config["temperature"] = self.parameters.temperature + + if self.parameters.top_p is not None: + config["top_p"] = self.parameters.top_p + + if self.parameters.thinking_enable: + config["thinking_config"] = { + "include_thoughts": True, + "thinking_budget": self.parameters.thinking_budget or 1024, + } + else: + config["thinking_config"] = { + "include_thoughts": False, + "thinking_budget": 0, + } + + fmt_tools, fmt_tool_choice = self._format_tools(tools, tool_choice) + + if fmt_tools is not None: + config["tools"] = fmt_tools + + if fmt_tool_choice is not None: + config["tool_config"] = fmt_tool_choice + + kwargs: dict[str, Any] = { + "model": model_name, + "contents": formatted_messages, + "config": config, + } + + start_datetime = datetime.now() + + if self.stream: + response = await client.aio.models.generate_content_stream( + **kwargs, + ) + # Pass client to the generator so the aiohttp session it owns + # stays alive until the stream is fully consumed. + return self._parse_stream_response( + start_datetime, + response, + client, + ) + + response = await client.aio.models.generate_content(**kwargs) + return self._parse_completion_response(start_datetime, response) + + async def _parse_stream_response( + self, + start_datetime: datetime, + response: Any, + _client: Any = None, + ) -> AsyncGenerator[ChatResponse, None]: + """Parse the Gemini streaming response. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`Any`): + The Gemini async stream object from + ``client.aio.models.generate_content_stream``. + _client (`Any`, optional): + The ``genai.Client`` that produced the stream. Held here so + its aiohttp session is not garbage-collected before the + stream is fully consumed. + + Yields: + `ChatResponse`: + Incremental ``ChatResponse`` objects with ``is_last=False`` + followed by a final one with ``is_last=True``. + """ + # All delta should have the same block identifier + # Use the API's response_id when available (it arrives at the first + # chunk); otherwise generate a UUID to ensure all chunks share a + # stable id. + response_id: str | None = None + acc_text = TextBlock(text="") + acc_thinking = ThinkingBlock(thinking="") + acc_tool_calls: dict = {} + usage = None + + async for chunk in response: + # Capture response_id from the first chunk that carries it + if response_id is None: + response_id = ( + getattr(chunk, "response_id", None) or _generate_id() + ) + + delta_content: list = [] + + if ( + chunk.candidates + and chunk.candidates[0].content + and chunk.candidates[0].content.parts + ): + for part in chunk.candidates[0].content.parts: + if part.text: + if part.thought: + acc_thinking.thinking += part.text + delta_content.append( + ThinkingBlock( + id=acc_thinking.id, + thinking=part.text, + ), + ) + else: + acc_text.text += part.text + delta_content.append( + TextBlock(id=acc_text.id, text=part.text), + ) + + if part.function_call: + keyword_args = part.function_call.args or {} + if part.thought_signature: + call_id = base64.b64encode( + part.thought_signature, + ).decode("utf-8") + else: + call_id = part.function_call.id or _generate_id() + input_str = json.dumps( + keyword_args, + ensure_ascii=False, + ) + acc_tool_calls[call_id] = { + "name": part.function_call.name, + "input": input_str, + } + delta_content.append( + ToolCallBlock( + id=call_id, + name=part.function_call.name, + input=input_str, + ), + ) + + usage = self._extract_usage(chunk.usage_metadata, start_datetime) + + if delta_content: + yield ChatResponse( + id=response_id, + content=delta_content, + is_last=False, + usage=usage, + ) + + final_content: list = [] + if acc_thinking.thinking: + final_content.append(acc_thinking) + if acc_text.text: + final_content.append(acc_text) + for call_id, tc in acc_tool_calls.items(): + final_content.append( + ToolCallBlock(id=call_id, name=tc["name"], input=tc["input"]), + ) + + yield ChatResponse( + id=response_id or _generate_id(), + content=final_content, + is_last=True, + usage=usage, + ) + + def _parse_completion_response( + self, + start_datetime: datetime, + response: GenerateContentResponse, + ) -> ChatResponse: + """Parse the Gemini non-streaming response. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`GenerateContentResponse`): + The Gemini generate content response object. + + Returns: + `ChatResponse`: + A single ``ChatResponse`` with ``is_last=True``. + """ + content_blocks: List[TextBlock | ToolCallBlock | ThinkingBlock] = [] + + if ( + response.candidates + and response.candidates[0].content + and response.candidates[0].content.parts + ): + for part in response.candidates[0].content.parts: + if part.text: + if part.thought: + content_blocks.append( + ThinkingBlock(thinking=part.text), + ) + else: + content_blocks.append(TextBlock(text=part.text)) + + if part.function_call: + keyword_args = part.function_call.args or {} + if part.thought_signature: + call_id = base64.b64encode( + part.thought_signature, + ).decode("utf-8") + else: + call_id = part.function_call.id or _generate_id() + content_blocks.append( + ToolCallBlock( + id=call_id, + name=part.function_call.name, + input=json.dumps(keyword_args, ensure_ascii=False), + ), + ) + + usage = self._extract_usage(response.usage_metadata, start_datetime) + + return ChatResponse( + id=getattr(response, "response_id", None) or _generate_id(), + content=content_blocks, + is_last=True, + usage=usage, + ) + + def _extract_usage( + self, + usage_metadata: Any, + start_datetime: datetime, + ) -> ChatUsage | None: + """Extract ChatUsage from usage_metadata. + + Args: + usage_metadata (`Any`): + The usage metadata object from a Gemini response. + start_datetime (`datetime`): + The start datetime of the response generation. + + Returns: + `ChatUsage | None`: + A ``ChatUsage`` object if usage data is available, otherwise + ``None``. + """ + if not usage_metadata: + return None + prompt_tokens = usage_metadata.prompt_token_count + total_tokens = usage_metadata.total_token_count + if prompt_tokens is not None and total_tokens is not None: + return ChatUsage( + input_tokens=prompt_tokens, + output_tokens=total_tokens - prompt_tokens, + time=(datetime.now() - start_datetime).total_seconds(), + cache_input_tokens=getattr( + usage_metadata, + "cached_content_token_count", + 0, + ), + ) + return None + + def _format_tools( + self, + tools: list[dict] | None, + tool_choice: ToolChoice | None, + ) -> tuple[list[dict] | None, dict | None]: + """Validate and format tools and tool_choice for Gemini. + + Converts tool schemas to Gemini's ``function_declarations`` + format (resolving ``$ref`` references) and maps tool_choice + modes to Gemini's ``function_calling_config``. When + ``tool_choice.tools`` is specified the schemas list is filtered + to only those tools. When ``tool_choice.mode`` is a specific + tool name (str) the model is restricted via + ``allowed_function_names`` without needing to filter the list, + preserving prompt-cache efficiency. + + Args: + tools (`list[dict] | None`, optional): + The raw tool schemas. + tool_choice (`ToolChoice | None`, optional): + The tool choice configuration. + + Returns: + `tuple[list[dict] | None, dict | None]`: + A tuple of (formatted_tools, formatted_tool_config). + """ + if tool_choice and tools: + self._validate_tool_choice(tool_choice, tools) + if tool_choice.tools: + allowed = set(tool_choice.tools) + tools = [t for t in tools if t["function"]["name"] in allowed] + + fmt_tools = None + if tools: + function_declarations = [] + for schema in tools: + if "function" not in schema: + continue + func = schema["function"].copy() + if "parameters" in func: + func["parameters"] = _sanitize_schema_for_gemini( + _flatten_json_schema(func["parameters"]), + ) + function_declarations.append(func) + fmt_tools = [{"function_declarations": function_declarations}] + + if not tool_choice: + return fmt_tools, None + + mode = tool_choice.mode + + if mode not in _TOOL_CHOICE_LITERAL_MODES: + # mode is a specific tool name — restrict to that single tool + fmt_choice: dict = { + "function_calling_config": { + "mode": "ANY", + "allowed_function_names": [mode], + }, + } + return fmt_tools, fmt_choice + + mode_mapping = { + "auto": "AUTO", + "none": "NONE", + "required": "ANY", + } + fmt_choice = { + "function_calling_config": {"mode": mode_mapping[mode]}, + } + return fmt_tools, fmt_choice diff --git a/src/agentscope/model/_gemini/_models/gemini-2.5-flash.yaml b/src/agentscope/model/_gemini/_models/gemini-2.5-flash.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d64bc3ca9d1a0a1ffe36a3de295d824fe0455d77 --- /dev/null +++ b/src/agentscope/model/_gemini/_models/gemini-2.5-flash.yaml @@ -0,0 +1,25 @@ +name: gemini-2.5-flash +label: Gemini 2.5 Flash +status: active + +input_types: + - text/plain + - application/x-thinking + - image/jpeg + - image/png + - image/gif + - image/webp + - audio/mp3 + - audio/wav + - video/mp4 + +output_types: + - text/plain + - application/x-thinking + +context_size: 1048576 +output_size: 65536 + +parameter_overrides: + max_tokens: + maximum: 65536 diff --git a/src/agentscope/model/_gemini/_models/gemini-2.5-pro.yaml b/src/agentscope/model/_gemini/_models/gemini-2.5-pro.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4b224b3ba18592104a4c40bff6de2014c133fd7d --- /dev/null +++ b/src/agentscope/model/_gemini/_models/gemini-2.5-pro.yaml @@ -0,0 +1,25 @@ +name: gemini-2.5-pro +label: Gemini 2.5 Pro +status: active + +input_types: + - text/plain + - application/x-thinking + - image/jpeg + - image/png + - image/gif + - image/webp + - audio/mp3 + - audio/wav + - video/mp4 + +output_types: + - text/plain + - application/x-thinking + +context_size: 1048576 +output_size: 65536 + +parameter_overrides: + max_tokens: + maximum: 65536 diff --git a/src/agentscope/model/_gemini/_models/gemini-3-flash-preview.yaml b/src/agentscope/model/_gemini/_models/gemini-3-flash-preview.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a736e015b193c174e30102197cb782d52c4ff0d6 --- /dev/null +++ b/src/agentscope/model/_gemini/_models/gemini-3-flash-preview.yaml @@ -0,0 +1,25 @@ +name: gemini-3-flash-preview +label: Gemini 3 Flash Preview +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + - audio/mp3 + - audio/wav + - video/mp4 + - application/pdf + +output_types: + - text/plain + - application/x-thinking + +context_size: 1048576 +output_size: 65536 + +parameter_overrides: + max_tokens: + maximum: 65536 diff --git a/src/agentscope/model/_gemini/_models/gemini-3.1-pro-preview.yaml b/src/agentscope/model/_gemini/_models/gemini-3.1-pro-preview.yaml new file mode 100644 index 0000000000000000000000000000000000000000..57bfbb4dbe33a4336470bd1e6499fd56a8b34d2d --- /dev/null +++ b/src/agentscope/model/_gemini/_models/gemini-3.1-pro-preview.yaml @@ -0,0 +1,25 @@ +name: gemini-3.1-pro-preview +label: Gemini 3.1 Pro Preview +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + - audio/mp3 + - audio/wav + - video/mp4 + - application/pdf + +output_types: + - text/plain + - application/x-thinking + +context_size: 1048576 +output_size: 65536 + +parameter_overrides: + max_tokens: + maximum: 65536 diff --git a/src/agentscope/model/_model_card.py b/src/agentscope/model/_model_card.py new file mode 100644 index 0000000000000000000000000000000000000000..b0e2546ef5c25dd9561d8dc822c21b70d70be2ce --- /dev/null +++ b/src/agentscope/model/_model_card.py @@ -0,0 +1,159 @@ +# -*- coding: utf-8 -*- +"""The model card class.""" +import copy +from datetime import datetime +from typing import Literal, Self, Type + +import yaml +from pydantic import BaseModel, Field + + +class ModelCard(BaseModel): + """The model card class.""" + + type: Literal["chat_model"] = "chat_model" + """The model card type.""" + + name: str = Field(description="The name of the model") + """The model name.""" + + label: str = Field(description="The model label.") + """The model label used for frontend rendering.""" + + status: Literal["active", "deprecated", "sunset"] = Field( + title="Status", + description="The model status", + ) + """The model status.""" + + deprecated_at: datetime | None = Field( + default=None, + description="The model deprecation date and time.", + title="Deprecation date", + ) + """The model deprecated at.""" + + input_types: list[str] = Field( + description="The supported model input types.", + title="Input types", + default=["text/plain"], + ) + """The model supported input types.""" + + output_types: list[str] = Field( + description="The supported model output types.", + title="Output types", + default=["text/plain"], + ) + """The model supported output types.""" + + context_size: int = Field( + title="Context size", + description="The context size.", + gt=0, + ) + """The model context size.""" + + output_size: int = Field( + title="Max output tokens", + description="The maximum number of tokens.", + gt=0, + ) + """The model max output tokens.""" + + parameter_schema: dict + """The parameters schema, which will be combined with the schema from the + DashScopeChatParameter class.""" + + parameters_overrides: dict[str, dict] + """The parameter overrides, which will be merged into the parameter schema. + """ + + @classmethod + def from_yaml( + cls, + yaml_path: str, + parameter_class: Type[BaseModel], + ) -> Self: + """Read a model card from a YAML file, and merge the parameter schema + with the override parameter schema in the yaml file. + + Args: + yaml_path (`str`): + Path to the YAML file + parameter_class (`Type[BaseModel]`): + The parameter class (e.g., DashScopeChatParameters) + + Returns: + `list[ModelCard]`: + ModelCard instance with merged parameter schema + """ + + # Load YAML config + with open(yaml_path, "r", encoding="utf-8") as file: + config = yaml.safe_load(file) + + # Get base schema from parameter class + base_schema = parameter_class.model_json_schema() + properties = copy.deepcopy(base_schema.get("properties", {})) + + # Auto-filter: remove thinking parameters if not supported + output_types = config.get("output_types", []) + if "application/x-thinking" not in output_types: + properties.pop("thinking_enable", None) + properties.pop("thinking_budget", None) + + # Auto-filter: only omni-style models that declare an ``audio/*`` + # output type expose the ``voice`` parameter to the frontend + # popover. Defense in depth — a model yaml can also drop ``voice`` + # explicitly via ``parameter_overrides: { voice: { hidden: true } }``. + if not any( + isinstance(t, str) and t.startswith("audio/") for t in output_types + ): + properties.pop("voice", None) + + # Auto-inject: set max_tokens maximum from output_size + if "max_tokens" in properties and "output_size" in config: + properties["max_tokens"]["maximum"] = config["output_size"] + + # Apply parameter_overrides with simple dict merge + overrides = config.get("parameter_overrides", {}) + for param_name, override in overrides.items(): + if override is None: + # null means remove + properties.pop(param_name, None) + continue + + if isinstance(override, dict): + # Check for hidden flag + if override.get("hidden"): + properties.pop(param_name, None) + continue + + # Simple dict merge: {**base, **override} + if param_name in properties: + properties[param_name] = { + **properties[param_name], + **override, + } + + # Build final parameter schema + final_schema = { + "type": "object", + "properties": properties, + "required": base_schema.get("required", []), + } + + # Create ModelCard instance + return cls( + name=config["name"], + label=config["label"], + status=config.get("status", "active"), + deprecated_at=config.get("deprecated_at"), + input_types=config.get("input_types", ["text/plain"]), + output_types=config.get("output_types", ["text/plain"]), + context_size=config["context_size"], + output_size=config["output_size"], + parameter_schema=final_schema, + parameters_overrides=config.get("parameter_overrides", {}), + ) diff --git a/src/agentscope/model/_model_response.py b/src/agentscope/model/_model_response.py new file mode 100644 index 0000000000000000000000000000000000000000..555fa7b9ffaef76a0cd00b8872fadbeea1600421 --- /dev/null +++ b/src/agentscope/model/_model_response.py @@ -0,0 +1,75 @@ +# -*- coding: utf-8 -*- +"""The model response module.""" +from dataclasses import dataclass, field +from datetime import datetime +from typing import Literal, Sequence + +from ._model_usage import ChatUsage +from .._utils._common import _generate_id +from .._utils._mixin import DictMixin +from ..message import ( + TextBlock, + ToolCallBlock, + ThinkingBlock, + DataBlock, +) +from ..types import JSONSerializableObject + + +@dataclass +class ChatResponse(DictMixin): + """The response of chat models.""" + + content: Sequence[TextBlock | ToolCallBlock | ThinkingBlock | DataBlock] + """The content of the chat response, which can include text blocks, + tool use blocks, or thinking blocks.""" + + is_last: bool + """Whether this response is the last response, if `Ture`, the content will + be the complete response, otherwise the content is a partial response""" + + id: str = field(default_factory=_generate_id) + """The unique identifier.""" + + created_at: str = field(default_factory=lambda: datetime.now().isoformat()) + """When the response was created""" + + type: Literal["chat_response"] = field( + default_factory=lambda: "chat_response", + ) + """The type of the response, which is always 'chat_response'.""" + + usage: ChatUsage | None = field(default_factory=lambda: None) + """The usage information of the chat response, if available.""" + + metadata: dict[str, JSONSerializableObject] = field( + default_factory=lambda: {}, + ) + """The metadata of the chat response""" + + +@dataclass +class StructuredResponse: + """The structured response of chat models.""" + + content: dict + """The structured output of the model.""" + + id: str = field(default_factory=_generate_id) + """The unique identifier.""" + + created_at: str = field(default_factory=lambda: datetime.now().isoformat()) + """When the response was created""" + + type: Literal["structured_response"] = field( + default_factory=lambda: "structured_response", + ) + """The type of the response, which is always 'structured_response'.""" + + usage: ChatUsage | None = field(default_factory=lambda: None) + """The usage information of the chat response, if available.""" + + metadata: dict[str, JSONSerializableObject] = field( + default_factory=lambda: {}, + ) + """The metadata of the chat response""" diff --git a/src/agentscope/model/_model_usage.py b/src/agentscope/model/_model_usage.py new file mode 100644 index 0000000000000000000000000000000000000000..037cf32517f33b62ccbf2e0eeee56fa5076d15fc --- /dev/null +++ b/src/agentscope/model/_model_usage.py @@ -0,0 +1,32 @@ +# -*- coding: utf-8 -*- +"""The model usage class in agentscope.""" +from dataclasses import dataclass, field +from typing import Any, Literal + +from .._utils._mixin import DictMixin + + +@dataclass +class ChatUsage(DictMixin): + """The usage of a chat model API invocation.""" + + input_tokens: int + """The number of input tokens.""" + + output_tokens: int + """The number of output tokens.""" + + time: float + """The time used in seconds.""" + + cache_creation_input_tokens: int = field(default_factory=lambda: 0) + """The number of input tokens used to create the prompt cache.""" + + cache_input_tokens: int = field(default_factory=lambda: 0) + """The number of input tokens read from the prompt cache.""" + + type: Literal["chat"] = field(default_factory=lambda: "chat") + """The type of the usage, must be `chat`.""" + + metadata: dict[str, Any] | None = field(default_factory=lambda: None) + """Optional metadata associated with the usage.""" diff --git a/src/agentscope/model/_moonshot/__init__.py b/src/agentscope/model/_moonshot/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e2972eea7378ad5c05df2b54ccc9a3baa8d10059 --- /dev/null +++ b/src/agentscope/model/_moonshot/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- +"""The Moonshot AI LLM API modules.""" + +from ._model import MoonshotChatModel + +__all__ = [ + "MoonshotChatModel", +] diff --git a/src/agentscope/model/_moonshot/_model.py b/src/agentscope/model/_moonshot/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..a228b200d821b7301c4f265d3d39e6ae53b9ca46 --- /dev/null +++ b/src/agentscope/model/_moonshot/_model.py @@ -0,0 +1,485 @@ +# -*- coding: utf-8 -*- +"""The Moonshot AI chat model implementation.""" +from collections import OrderedDict +from datetime import datetime +from typing import Literal, Any, AsyncGenerator, TYPE_CHECKING, List, Type + +from pydantic import BaseModel, Field + +from .._base import ChatModelBase, _TOOL_CHOICE_LITERAL_MODES +from .._model_response import ChatResponse, StructuredResponse +from .._model_usage import ChatUsage +from ...credential import MoonshotCredential +from ...formatter import FormatterBase, MoonshotChatFormatter +from ...message import Msg, ThinkingBlock, ToolCallBlock, TextBlock +from ...tool import ToolChoice + +if TYPE_CHECKING: + from openai.types.chat import ChatCompletion + from openai import AsyncStream +else: + ChatCompletion = Any + AsyncStream = Any + + +class MoonshotChatModel(ChatModelBase): + """The Moonshot AI chat model.""" + + class Parameters(BaseModel): + """The parameters for the Moonshot AI chat model.""" + + max_tokens: int | None = Field( + default=None, + title="Max Tokens", + description="The maximum number of tokens for the LLM output.", + gt=0, + ) + + thinking_enable: bool = Field( + default=False, + title="Thinking", + description=( + "Whether to enable thinking mode. For kimi-k2-thinking, " + "thinking is always enabled. For kimi-k2.6, thinking is " + "enabled by default but can be disabled." + ), + ) + + temperature: float | None = Field( + default=None, + title="Temperature", + description="The temperature for the LLM output.", + ge=0, + le=1, + ) + + top_p: float | None = Field( + default=None, + title="Top P", + description="The top P value for the LLM output.", + gt=0, + le=1, + ) + + type: Literal["moonshot_chat"] = "moonshot_chat" + """The type of the chat model.""" + + def __init__( + self, + credential: MoonshotCredential, + model: str, + parameters: "MoonshotChatModel.Parameters | None" = None, + stream: bool = True, + max_retries: int = 3, + retry_delay: float = 1.0, + context_size: int = 131072, + formatter: FormatterBase | None = None, + client_kwargs: dict[str, Any] | None = None, + ) -> None: + """Initialize the Moonshot AI chat model. + + Args: + credential (`MoonshotCredential`): + The Moonshot AI credential used to authenticate API calls. + model (`str`): + The model name, e.g. ``moonshot-v1-8k``, ``kimi-k2.6``. + parameters (`MoonshotChatModel.Parameters | None`, defaults to \ + `None`): + The API parameters. When ``None``, the default parameters + will be used. + stream (`bool`, defaults to `True`): + Whether to enable streaming output. + max_retries (`int`, defaults to `3`): + The maximum number of retries for the API. + retry_delay (`float`, defaults to `1.0`): + Seconds to sleep between retry attempts. + context_size (`int`, defaults to `131072`): + The model context size used for context compression. + formatter (`FormatterBase | None`, defaults to `None`): + The formatter that converts ``Msg`` objects to the format + required by the API. When ``None``, a + ``MoonshotChatFormatter`` instance will be used. + client_kwargs (`dict[str, Any] | None`, defaults to `None`): + Extra keyword arguments forwarded to ``openai.AsyncClient`` + (e.g. ``timeout``, ``default_headers``, ``http_client``). + """ + super().__init__( + credential=credential, + model=model, + parameters=parameters or self.Parameters(), + stream=stream, + max_retries=max_retries, + retry_delay=retry_delay, + context_size=context_size, + ) + self.formatter = formatter or MoonshotChatFormatter() + self.client_kwargs = client_kwargs or {} + + @classmethod + def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]: + import openai + + return ( + openai.APIConnectionError, + openai.APITimeoutError, + openai.RateLimitError, + openai.InternalServerError, + ) + + async def _call_api( + self, + model_name: str, + messages: list[Msg], + tools: list[dict] | None = None, + tool_choice: ToolChoice | None = None, + **generate_kwargs: Any, + ) -> ChatResponse | AsyncGenerator[ChatResponse, None]: + """Call the Moonshot AI chat API (OpenAI-compatible). + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list`): + A list of message dicts with ``role`` and ``content`` keys. + tools (`list[dict]`, default `None`): + The tools JSON schemas. + tool_choice (`ToolChoice | None`, optional): + Controls which (if any) tool is called by the model. + **generate_kwargs (`Any`): + Extra keyword arguments forwarded to the API. + + Returns: + `ChatResponse | AsyncGenerator[ChatResponse, None]`: + A ``ChatResponse`` when streaming is disabled, or an async + generator of ``ChatResponse`` objects when streaming is + enabled. + """ + import openai + + client = openai.AsyncClient( + **{ + "api_key": self.credential.api_key.get_secret_value(), + "base_url": self.credential.base_url, + **self.client_kwargs, + }, + ) + + formatted_messages = await self.formatter.format(messages) + + kwargs: dict[str, Any] = { + "model": model_name, + "messages": formatted_messages, + "stream": self.stream, + } + + if self.parameters.max_tokens is not None: + kwargs["max_tokens"] = self.parameters.max_tokens + + if self.parameters.temperature is not None: + kwargs["temperature"] = self.parameters.temperature + + if self.parameters.top_p is not None: + kwargs["top_p"] = self.parameters.top_p + + kwargs.update(generate_kwargs) + + thinking_type = ( + "enabled" if self.parameters.thinking_enable else "disabled" + ) + kwargs.setdefault("extra_body", {}) + kwargs["extra_body"].setdefault("thinking", {}) + kwargs["extra_body"]["thinking"].setdefault("type", thinking_type) + + fmt_tools, fmt_tool_choice = self._format_tools(tools, tool_choice) + + if fmt_tools: + kwargs["tools"] = fmt_tools + + if fmt_tool_choice is not None: + kwargs["tool_choice"] = fmt_tool_choice + + if self.stream: + kwargs["stream_options"] = {"include_usage": True} + + start_datetime = datetime.now() + response = await client.chat.completions.create(**kwargs) + + if self.stream: + return self._parse_stream_response(start_datetime, response) + + return self._parse_completion_response(start_datetime, response) + + async def _parse_stream_response( + self, + start_datetime: datetime, + response: AsyncStream, + ) -> AsyncGenerator[ChatResponse, None]: + """Parse the Moonshot AI streaming response. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`AsyncStream`): + The OpenAI-compatible async stream object. + + Yields: + `ChatResponse`: + Incremental ``ChatResponse`` objects with ``is_last=False`` + followed by a final one with ``is_last=True``. + """ + usage = None + response_id: str | None = None + # All delta should have the same block identifier + acc_text = TextBlock(text="") + acc_thinking = ThinkingBlock(thinking="") + acc_tool_calls: OrderedDict = OrderedDict() + + async with response as stream: + async for chunk in stream: + if chunk.usage: + u = chunk.usage + usage = ChatUsage( + input_tokens=u.prompt_tokens, + output_tokens=u.completion_tokens, + time=(datetime.now() - start_datetime).total_seconds(), + cache_input_tokens=getattr( + u, + "cached_tokens", + 0, + ), + ) + + # Capture response_id from the first chunk that carries it + response_id = response_id or getattr(chunk, "id", None) + + if not chunk.choices: + continue + + choice = chunk.choices[0] + delta = choice.delta + + # Thinking models (kimi-k2.6, kimi-k2-thinking) return + # reasoning_content before content in the stream. + delta_thinking = ( + getattr(delta, "reasoning_content", None) or "" + ) + if delta_thinking: + acc_thinking.thinking += delta_thinking + _thinking_kwargs: dict[str, Any] = { + "content": [ + ThinkingBlock( + id=acc_thinking.id, + thinking=delta_thinking, + ), + ], + "usage": usage, + "is_last": False, + } + if response_id: + _thinking_kwargs["id"] = response_id + yield ChatResponse(**_thinking_kwargs) + continue + + delta_text = getattr(delta, "content", None) or "" + acc_text.text += delta_text + + delta_tool_call_blocks: List[ToolCallBlock] = [] + for tool_call in getattr(delta, "tool_calls", None) or []: + idx = tool_call.index + args = tool_call.function.arguments or "" + if idx in acc_tool_calls: + acc_tool_calls[idx]["input"] += args + else: + acc_tool_calls[idx] = { + "id": tool_call.id, + "name": tool_call.function.name, + "input": args, + } + tc = acc_tool_calls[idx] + delta_tool_call_blocks.append( + ToolCallBlock( + id=tc["id"], + name=tc["name"], + input=args, + ), + ) + + delta_contents: List[TextBlock | ToolCallBlock] = [] + if delta_text: + delta_contents.append( + TextBlock(id=acc_text.id, text=delta_text), + ) + delta_contents.extend(delta_tool_call_blocks) + + if delta_contents: + _text_kwargs: dict[str, Any] = { + "content": delta_contents, + "usage": usage, + "is_last": False, + } + if response_id: + _text_kwargs["id"] = response_id + yield ChatResponse(**_text_kwargs) + + final_contents: List[ThinkingBlock | TextBlock | ToolCallBlock] = [] + if acc_thinking.thinking: + final_contents.append(acc_thinking) + if acc_text.text: + final_contents.append(acc_text) + for tc in acc_tool_calls.values(): + final_contents.append( + ToolCallBlock(id=tc["id"], name=tc["name"], input=tc["input"]), + ) + + _final_kwargs: dict[str, Any] = { + "content": final_contents, + "usage": usage, + "is_last": True, + } + if response_id: + _final_kwargs["id"] = response_id + yield ChatResponse(**_final_kwargs) + + def _parse_completion_response( + self, + start_datetime: datetime, + response: ChatCompletion, + ) -> ChatResponse: + """Parse the Moonshot AI non-streaming response. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`ChatCompletion`): + The OpenAI-compatible chat completion object. + + Returns: + `ChatResponse`: + A single ``ChatResponse`` with ``is_last=True``. + """ + content_blocks: List[ThinkingBlock | TextBlock | ToolCallBlock] = [] + + if response.choices: + choice = response.choices[0] + reasoning = getattr(choice.message, "reasoning_content", None) + if reasoning: + content_blocks.append(ThinkingBlock(thinking=reasoning)) + + if choice.message.content: + content_blocks.append(TextBlock(text=choice.message.content)) + + for tool_call in choice.message.tool_calls or []: + content_blocks.append( + ToolCallBlock( + id=tool_call.id, + name=tool_call.function.name, + input=tool_call.function.arguments, + ), + ) + + usage = None + if response.usage: + u = response.usage + usage = ChatUsage( + input_tokens=u.prompt_tokens, + output_tokens=u.completion_tokens, + time=(datetime.now() - start_datetime).total_seconds(), + cache_input_tokens=getattr(u, "cached_tokens", 0), + ) + + resp_kwargs: dict[str, Any] = { + "content": content_blocks, + "is_last": True, + "usage": usage, + } + response_id = getattr(response, "id", None) + if response_id: + resp_kwargs["id"] = response_id + + return ChatResponse(**resp_kwargs) + + def _format_tools( + self, + tools: list[dict] | None, + tool_choice: ToolChoice | None, + ) -> tuple[list[dict] | None, str | dict | None]: + """Validate, filter, and format tools and tool_choice for the + Moonshot AI API. + + When ``tool_choice.tools`` is specified the schemas list is filtered + to only those tools. When ``tool_choice.mode`` is a specific tool name + (str) the model is forced to call exactly that tool without needing to + filter the list, preserving prompt-cache efficiency. + + Args: + tools (`list[dict] | None`, optional): + The raw tool schemas. + tool_choice (`ToolChoice | None`, optional): + The tool choice configuration. + + Returns: + `tuple[list[dict] | None, str | dict | None]`: + A tuple of (formatted_tools, formatted_tool_choice). + """ + if tool_choice and tools: + self._validate_tool_choice(tool_choice, tools) + if tool_choice.tools: + allowed = set(tool_choice.tools) + tools = [t for t in tools if t["function"]["name"] in allowed] + + if not tool_choice: + return tools, None + + mode = tool_choice.mode + + if mode not in _TOOL_CHOICE_LITERAL_MODES: + return tools, {"type": "function", "function": {"name": mode}} + + return tools, mode + + async def _call_api_with_structured_output( + self, + model_name: str, + messages: list[Msg], + structured_model: Type[BaseModel] | dict, + tool_choice: ToolChoice | None = None, + **kwargs: Any, + ) -> StructuredResponse: + """Moonshot-specific override for structured output. + + Moonshot rejects ``tool_choice="required"`` or an object-form + ``tool_choice`` when thinking mode is enabled. In that case we + default ``tool_choice`` to ``"auto"`` and rely on the base class's + injected system-reminder prompt to guide the model. When thinking + is disabled, this falls through to the base implementation. + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list[Msg]`): + The context for the LLM to generate the structured output. + structured_model (`Type[BaseModel] | dict`): + A Pydantic model class or a JSON schema dict describing the + required output structure. + tool_choice (`ToolChoice | None`, defaults to `None`): + The tool_choice forwarded to ``_call_api``. When ``None`` + and thinking mode is enabled, it is downgraded to + ``ToolChoice(mode="auto")``; otherwise the base default + (force the structured-output tool) is used. + **kwargs (`Any`): + Additional keyword arguments forwarded to ``_call_api``. + + Returns: + `StructuredResponse`: + The structured response whose ``content`` is the validated + output dict matching ``structured_model``. + """ + if tool_choice is None and self.parameters.thinking_enable: + tool_choice = ToolChoice(mode="auto") + return await super()._call_api_with_structured_output( + model_name=model_name, + messages=messages, + structured_model=structured_model, + tool_choice=tool_choice, + **kwargs, + ) diff --git a/src/agentscope/model/_moonshot/_models/kimi-k2.5.yaml b/src/agentscope/model/_moonshot/_models/kimi-k2.5.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2f40c9256a0281a51935f9a36686a1732d61cee9 --- /dev/null +++ b/src/agentscope/model/_moonshot/_models/kimi-k2.5.yaml @@ -0,0 +1,22 @@ +name: kimi-k2.5 +label: Kimi K2.5 +status: active + +input_types: + - text/plain + - application/x-thinking + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + - application/x-thinking + +context_size: 262144 +output_size: 65536 + +parameter_overrides: + max_tokens: + maximum: 65536 diff --git a/src/agentscope/model/_moonshot/_models/kimi-k2.6.yaml b/src/agentscope/model/_moonshot/_models/kimi-k2.6.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5854a0b97a246ce96b2fba6261793f830379fb80 --- /dev/null +++ b/src/agentscope/model/_moonshot/_models/kimi-k2.6.yaml @@ -0,0 +1,27 @@ +name: kimi-k2.6 +label: Kimi K2.6 +status: active + +input_types: + - text/plain + - application/x-thinking + - image/jpeg + - image/png + - image/gif + - image/webp + - video/mp4 + - video/mpeg + - video/quicktime + - video/avi + - video/webm + +output_types: + - text/plain + - application/x-thinking + +context_size: 262144 +output_size: 32768 + +parameter_overrides: + max_tokens: + maximum: 32768 diff --git a/src/agentscope/model/_moonshot/_models/moonshot-v1-128k.yaml b/src/agentscope/model/_moonshot/_models/moonshot-v1-128k.yaml new file mode 100644 index 0000000000000000000000000000000000000000..90bef71a104431054362cc2c4991e2da59717242 --- /dev/null +++ b/src/agentscope/model/_moonshot/_models/moonshot-v1-128k.yaml @@ -0,0 +1,16 @@ +name: moonshot-v1-128k +label: Kimi Moonshot v1 128K +status: active + +input_types: + - text/plain + +output_types: + - text/plain + +context_size: 131072 +output_size: 131072 + +parameter_overrides: + max_tokens: + maximum: 131072 diff --git a/src/agentscope/model/_moonshot/_models/moonshot-v1-32k.yaml b/src/agentscope/model/_moonshot/_models/moonshot-v1-32k.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b65310164e878038ac0a21007194d412685761b8 --- /dev/null +++ b/src/agentscope/model/_moonshot/_models/moonshot-v1-32k.yaml @@ -0,0 +1,16 @@ +name: moonshot-v1-32k +label: Kimi Moonshot v1 32K +status: active + +input_types: + - text/plain + +output_types: + - text/plain + +context_size: 32768 +output_size: 32768 + +parameter_overrides: + max_tokens: + maximum: 32768 diff --git a/src/agentscope/model/_moonshot/_models/moonshot-v1-8k.yaml b/src/agentscope/model/_moonshot/_models/moonshot-v1-8k.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7784c6a53b65ab91744e793254925d11448bf6c1 --- /dev/null +++ b/src/agentscope/model/_moonshot/_models/moonshot-v1-8k.yaml @@ -0,0 +1,16 @@ +name: moonshot-v1-8k +label: Kimi Moonshot v1 8K +status: active + +input_types: + - text/plain + +output_types: + - text/plain + +context_size: 8192 +output_size: 8192 + +parameter_overrides: + max_tokens: + maximum: 8192 diff --git a/src/agentscope/model/_ollama/__init__.py b/src/agentscope/model/_ollama/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d02971352c9e3db25863a15e2901bddd8439d7cf --- /dev/null +++ b/src/agentscope/model/_ollama/__init__.py @@ -0,0 +1,9 @@ +# -*- coding: utf-8 -*- +"""The Ollama LLM API modules.""" + +from ._model import OllamaCredential, OllamaChatModel + +__all__ = [ + "OllamaCredential", + "OllamaChatModel", +] diff --git a/src/agentscope/model/_ollama/_model.py b/src/agentscope/model/_ollama/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..9768ef2fbaeef5c59e474704636f3b7021aa7eae --- /dev/null +++ b/src/agentscope/model/_ollama/_model.py @@ -0,0 +1,374 @@ +# -*- coding: utf-8 -*- +"""The Ollama chat model implementation.""" +import json +from datetime import datetime +from typing import Literal, Any, AsyncGenerator, TYPE_CHECKING, List, Type + +from pydantic import BaseModel, Field + +from ..._utils._common import _generate_id +from .._base import ChatModelBase +from .._model_response import ChatResponse +from .._model_usage import ChatUsage +from ...credential import OllamaCredential +from ...formatter import FormatterBase, OllamaChatFormatter +from ...message import Msg, ThinkingBlock, ToolCallBlock, TextBlock +from ...tool import ToolChoice +from ..._logging import logger + +if TYPE_CHECKING: + from ollama._types import ChatResponse as OllamaChatResponse +else: + OllamaChatResponse = Any + + +class OllamaChatModel(ChatModelBase): + """The Ollama chat model.""" + + class Parameters(BaseModel): + """The parameters for the Ollama chat model.""" + + max_tokens: int | None = Field( + default=None, + title="Max Tokens", + description="The maximum number of tokens for the LLM output.", + gt=0, + ) + + thinking_enable: bool = Field( + default=False, + title="Thinking", + description="Whether to enable thinking" + " (for models like qwen3, deepseek-r1).", + ) + + temperature: float | None = Field( + default=None, + title="Temperature", + description="The temperature for the LLM output.", + ge=0, + le=2, + ) + + type: Literal["ollama_chat"] = "ollama_chat" + """The type of the chat model.""" + + def __init__( + self, + credential: OllamaCredential | None = None, + model: str = "", + parameters: "OllamaChatModel.Parameters | None" = None, + stream: bool = True, + max_retries: int = 3, + retry_delay: float = 1.0, + context_size: int = 32768, + formatter: FormatterBase | None = None, + client_kwargs: dict[str, Any] | None = None, + ) -> None: + """Initialize the Ollama chat model. + + Args: + credential (`OllamaCredential | None`, defaults to `None`): + The Ollama connection settings. When ``None``, a default + ``OllamaCredential`` (localhost) will be used. + model (`str`): + The Ollama model name, e.g. ``llama3.3`` or ``qwen3:14b``. + parameters (`OllamaChatModel.Parameters | None`, defaults to \ + `None`): + The Ollama API parameters. When ``None``, the default + parameters will be used. + stream (`bool`, defaults to `True`): + Whether to enable streaming output. + max_retries (`int`, defaults to `3`): + The maximum number of retries for the Ollama API. + retry_delay (`float`, defaults to `1.0`): + Seconds to sleep between retry attempts. + context_size (`int`, defaults to `32768`): + The model context size used for context compression. + formatter (`FormatterBase | None`, defaults to `None`): + The formatter that converts ``Msg`` objects to the format + required by the Ollama API. When ``None``, an + ``OllamaChatFormatter`` instance will be used. + client_kwargs (`dict[str, Any] | None`, defaults to `None`): + Extra keyword arguments forwarded to ``ollama.AsyncClient`` + and onward to the underlying ``httpx.AsyncClient`` + (e.g. ``timeout``, ``headers``, ``verify``). + """ + resolved_credential = credential or OllamaCredential() + + super().__init__( + credential=resolved_credential, + model=model, + parameters=parameters or self.Parameters(), + stream=stream, + max_retries=max_retries, + retry_delay=retry_delay, + context_size=context_size, + ) + + self.formatter = formatter or OllamaChatFormatter() + self.client_kwargs = client_kwargs or {} + + @classmethod + def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]: + import httpx + + # Local service: retry transient transport-layer failures only. + # ollama.ResponseError wraps server-side errors regardless of cause + # (incl. 4xx like "model not found"), so we don't retry on it. + return ( + httpx.ConnectError, + httpx.ReadTimeout, + httpx.RemoteProtocolError, + ) + + async def _call_api( + self, + model_name: str, + messages: list[Msg], + tools: list[dict] | None = None, + tool_choice: ToolChoice | None = None, + **generate_kwargs: Any, + ) -> ChatResponse | AsyncGenerator[ChatResponse, None]: + """Call the Ollama chat API. + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list`): + A list of message dicts with ``role`` and ``content`` keys. + tools (`list[dict]`, default `None`): + The tools JSON schemas. + tool_choice (`ToolChoice | None`, optional): + Not supported by Ollama yet (ignored with warning). + **generate_kwargs (`Any`): + Extra keyword arguments forwarded to the Ollama API. + + Returns: + `ChatResponse | AsyncGenerator[ChatResponse, None]`: + A ``ChatResponse`` when streaming is disabled, or an async + generator of ``ChatResponse`` objects when streaming is + enabled. + """ + import ollama + + client = ollama.AsyncClient( + **{ + "host": self.credential.host, + **self.client_kwargs, + }, + ) + + formatted_messages = await self.formatter.format(messages) + + kwargs: dict[str, Any] = { + "model": model_name, + "messages": formatted_messages, + "stream": self.stream, + } + + options: dict[str, Any] = {} + if self.parameters.max_tokens is not None: + options["num_predict"] = self.parameters.max_tokens + if self.parameters.temperature is not None: + options["temperature"] = self.parameters.temperature + if options: + kwargs["options"] = options + + kwargs["think"] = self.parameters.thinking_enable + + kwargs.update(generate_kwargs) + + fmt_tools, _ = self._format_tools(tools, tool_choice) + + if fmt_tools: + kwargs["tools"] = fmt_tools + + start_datetime = datetime.now() + response = await client.chat(**kwargs) + + if self.stream: + return self._parse_stream_response(start_datetime, response) + + return await self._parse_completion_response(start_datetime, response) + + def _format_tools( + self, + tools: list[dict] | None, + tool_choice: ToolChoice | None, + ) -> tuple[list[dict] | None, None]: + """Validate, filter tools, and warn if tool_choice is set. + + Ollama does not support ``tool_choice`` natively. When + ``tool_choice.tools`` is specified the schemas list is filtered to + only those tools. Any ``tool_choice.mode`` value is ignored with a + warning. + + Args: + tools (`list[dict] | None`, optional): + The raw tool schemas. + tool_choice (`ToolChoice | None`, optional): + The tool choice configuration. + + Returns: + `tuple[list[dict] | None, None]`: + A tuple of (filtered_tools, None) — tool_choice is always + ``None`` since Ollama does not support it. + """ + if tool_choice and tools: + self._validate_tool_choice(tool_choice, tools) + if tool_choice.tools: + allowed = set(tool_choice.tools) + tools = [t for t in tools if t["function"]["name"] in allowed] + + if tool_choice: + logger.warning( + "Ollama ignores tool_choice.mode; " + "tool_choice.tools is still applied to filter tool schemas.", + ) + + return tools, None + + async def _parse_stream_response( + self, + start_datetime: datetime, + response: Any, + ) -> AsyncGenerator[ChatResponse, None]: + """Parse the Ollama streaming response. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`Any`): + The Ollama async stream object. + + Yields: + `ChatResponse`: + Incremental ``ChatResponse`` objects with ``is_last=False`` + followed by a final one with ``is_last=True``. + """ + # All delta should have the same block identifier. + # Ollama does not return a request id, so we generate one upfront + # to keep it stable. + response_id = getattr(response, "id", None) or _generate_id() + acc_text = TextBlock(text="") + acc_thinking = ThinkingBlock(thinking="") + acc_tool_calls: dict = {} + usage = None + + async for chunk in response: + delta_content: list = [] + msg = chunk.message + + chunk_thinking = getattr(msg, "thinking", None) + if chunk_thinking: + acc_thinking.thinking += chunk_thinking + delta_content.append( + ThinkingBlock(id=acc_thinking.id, thinking=chunk_thinking), + ) + + if msg.content: + acc_text.text += msg.content + delta_content.append( + TextBlock(id=acc_text.id, text=msg.content), + ) + + for idx, tool_call in enumerate(msg.tool_calls or []): + function = tool_call.function + tool_id = f"{idx}_{function.name}" + input_str = json.dumps(function.arguments) + acc_tool_calls[tool_id] = { + "name": function.name, + "input": input_str, + } + delta_content.append( + ToolCallBlock( + id=tool_id, + name=function.name, + input=input_str, + ), + ) + + current_time = (datetime.now() - start_datetime).total_seconds() + usage = ChatUsage( + input_tokens=getattr(chunk, "prompt_eval_count", 0) or 0, + output_tokens=getattr(chunk, "eval_count", 0) or 0, + time=current_time, + ) + + if delta_content: + yield ChatResponse( + id=response_id, + content=delta_content, + is_last=False, + usage=usage, + ) + + final_content: list = [] + if acc_thinking.thinking: + final_content.append(acc_thinking) + if acc_text.text: + final_content.append(acc_text) + for tool_id, tc in acc_tool_calls.items(): + final_content.append( + ToolCallBlock(id=tool_id, name=tc["name"], input=tc["input"]), + ) + + yield ChatResponse( + id=response_id, + content=final_content, + is_last=True, + usage=usage, + ) + + async def _parse_completion_response( + self, + start_datetime: datetime, + response: OllamaChatResponse, + ) -> ChatResponse: + """Parse the Ollama non-streaming response. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`OllamaChatResponse`): + The Ollama chat response object. + + Returns: + `ChatResponse`: + A single ``ChatResponse`` with ``is_last=True``. + """ + content_blocks: List[TextBlock | ToolCallBlock | ThinkingBlock] = [] + + message_thinking = getattr(response.message, "thinking", None) + if message_thinking: + content_blocks.append(ThinkingBlock(thinking=message_thinking)) + + if response.message.content: + content_blocks.append(TextBlock(text=response.message.content)) + + for idx, tool_call in enumerate(response.message.tool_calls or []): + content_blocks.append( + ToolCallBlock( + id=f"{idx}_{tool_call.function.name}", + name=tool_call.function.name, + input=json.dumps(tool_call.function.arguments), + ), + ) + + usage = None + prompt_eval = getattr(response, "prompt_eval_count", None) + eval_count = getattr(response, "eval_count", None) + if prompt_eval is not None and eval_count is not None: + usage = ChatUsage( + input_tokens=prompt_eval, + output_tokens=eval_count, + time=(datetime.now() - start_datetime).total_seconds(), + ) + + return ChatResponse( + id=getattr(response, "id", None) or _generate_id(), + content=content_blocks, + is_last=True, + usage=usage, + ) diff --git a/src/agentscope/model/_ollama/_models/deepseek-r1-14b.yaml b/src/agentscope/model/_ollama/_models/deepseek-r1-14b.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2a18ddbc49817927b9d1cc40b3c46e35c7ceb9d3 --- /dev/null +++ b/src/agentscope/model/_ollama/_models/deepseek-r1-14b.yaml @@ -0,0 +1,17 @@ +name: deepseek-r1:14b +label: DeepSeek R1 14B +status: active + +input_types: + - text/plain + +output_types: + - text/plain + - application/x-thinking + +context_size: 131072 +output_size: 8192 + +parameter_overrides: + max_tokens: + maximum: 8192 diff --git a/src/agentscope/model/_ollama/_models/llama4.yaml b/src/agentscope/model/_ollama/_models/llama4.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a9543236143b198eb1a9a084771412d09c06a514 --- /dev/null +++ b/src/agentscope/model/_ollama/_models/llama4.yaml @@ -0,0 +1,22 @@ +name: llama4 +label: Llama 4 Scout +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + +context_size: 10485760 +output_size: 8192 + +parameter_overrides: + max_tokens: + maximum: 8192 + thinking_enable: + hidden: true diff --git a/src/agentscope/model/_ollama/_models/qwen3-14b.yaml b/src/agentscope/model/_ollama/_models/qwen3-14b.yaml new file mode 100644 index 0000000000000000000000000000000000000000..20de34b0a6ac048915e5a69ba444d69c06ef90e1 --- /dev/null +++ b/src/agentscope/model/_ollama/_models/qwen3-14b.yaml @@ -0,0 +1,17 @@ +name: qwen3:14b +label: Qwen3 14B +status: active + +input_types: + - text/plain + +output_types: + - text/plain + - application/x-thinking + +context_size: 40960 +output_size: 8192 + +parameter_overrides: + max_tokens: + maximum: 8192 diff --git a/src/agentscope/model/_ollama/_models/qwen3.5-9b.yaml b/src/agentscope/model/_ollama/_models/qwen3.5-9b.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ec8bff21325fa11bc8a9f18074da8b4ab69ee949 --- /dev/null +++ b/src/agentscope/model/_ollama/_models/qwen3.5-9b.yaml @@ -0,0 +1,17 @@ +name: qwen3.5 +label: Qwen3.5 9B +status: active + +input_types: + - text/plain + +output_types: + - text/plain + - application/x-thinking + +context_size: 262144 +output_size: 8192 + +parameter_overrides: + max_tokens: + maximum: 8192 diff --git a/src/agentscope/model/_openai_chat/__init__.py b/src/agentscope/model/_openai_chat/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1fdaaff55da9c67f9eff623cb2e424274832f0d4 --- /dev/null +++ b/src/agentscope/model/_openai_chat/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- +"""The OpenAI Chat Completions API modules.""" + +from ._model import OpenAIChatModel + +__all__ = [ + "OpenAIChatModel", +] diff --git a/src/agentscope/model/_openai_chat/_model.py b/src/agentscope/model/_openai_chat/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..f6763dae7b137096408a846d7ffb525f48d9b5a4 --- /dev/null +++ b/src/agentscope/model/_openai_chat/_model.py @@ -0,0 +1,726 @@ +# -*- coding: utf-8 -*- +"""The OpenAI Chat Completions model implementation.""" +import warnings +import base64 +import io +import wave +from collections import OrderedDict +from datetime import datetime +from typing import Literal, Any, AsyncGenerator, TYPE_CHECKING, List, Type + +from pydantic import BaseModel, Field + +from ..._utils._audio import _build_streaming_wav_header +from ..._utils._common import _generate_id, _flatten_json_schema +from .._base import ChatModelBase, _TOOL_CHOICE_LITERAL_MODES +from .._model_response import ChatResponse, StructuredResponse +from .._model_usage import ChatUsage +from ...credential import OpenAICredential +from ...formatter import FormatterBase, OpenAIChatFormatter +from ...message import ( + Msg, + ThinkingBlock, + ToolCallBlock, + TextBlock, + DataBlock, + Base64Source, +) +from ...tool import ToolChoice + +if TYPE_CHECKING: + from openai.types.chat import ChatCompletion + from openai import AsyncStream +else: + ChatCompletion = Any + AsyncStream = Any + + +class OpenAIChatModel(ChatModelBase): + """The OpenAI Chat Completions model.""" + + class Parameters(BaseModel): + """The parameters for the OpenAI Chat model.""" + + max_tokens: int | None = Field( + default=None, + title="Max Tokens", + description="The maximum number of tokens for the LLM output.", + gt=0, + ) + + thinking_enable: bool = Field( + default=False, + title="Thinking", + description=( + "Whether to enable reasoning for reasoning models " + "(e.g. o3, o4-mini, gpt-5.5). Use reasoning_effort to " + "control the depth of reasoning." + ), + ) + + reasoning_effort: ( + Literal["none", "minimal", "low", "medium", "high", "xhigh"] | None + ) = Field( + default=None, + title="Reasoning Effort", + description=( + "Controls the depth of reasoning for reasoning models " + "(e.g. o3, o4-mini, gpt-5.5). Supported values are " + "model-dependent and may include: none, minimal, low, " + "medium, high, xhigh." + ), + ) + + temperature: float | None = Field( + default=None, + title="Temperature", + description="The temperature for the LLM output.", + ge=0, + le=2, + ) + + top_p: float | None = Field( + default=None, + title="Top P", + description="The top P value for the LLM output.", + gt=0, + le=1, + ) + + parallel_tool_calls: bool = Field( + default=True, + title="Parallel Tool Calls", + description="Whether to enable parallel tool calls.", + ) + + voice: str | None = Field( + default=None, + title="Voice", + description=( + "Voice for audio output on omni-style models (e.g. " + "``gpt-audio-mini``). Setting this implicitly asks the " + "model to speak its response — ``modalities`` is filled in " + "automatically. Supported voices vary by model — see the " + "model card's ``voice.suggestions``. Any value the API " + "accepts works — the suggestions are convenience-only. " + "Leave unset for text-only " + "responses." + ), + ) + + type: Literal["openai_chat"] = "openai_chat" + """The type of the chat model.""" + + def __init__( + self, + credential: OpenAICredential, + model: str, + parameters: "OpenAIChatModel.Parameters | None" = None, + stream: bool = True, + max_retries: int = 3, + retry_delay: float = 1.0, + context_size: int = 128000, + formatter: FormatterBase | None = None, + client_kwargs: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, + ) -> None: + """Initialize the OpenAI chat model. + + Args: + credential (`OpenAICredential`): + The OpenAI credential used to authenticate API calls. + model (`str`): + The OpenAI model name, e.g. ``gpt-4.1``. + parameters (`OpenAIChatModel.Parameters | None`, defaults to \ + `None`): + The OpenAI Chat API parameters. When ``None``, the default + parameters will be used. + stream (`bool`, defaults to `True`): + Whether to enable streaming output. + max_retries (`int`, defaults to `3`): + The maximum number of retries for the OpenAI API. + retry_delay (`float`, defaults to `1.0`): + Seconds to sleep between retry attempts. + context_size (`int`, defaults to `128000`): + The model context size used for context compression. + formatter (`FormatterBase | None`, defaults to `None`): + The formatter that converts ``Msg`` objects to the format + required by the OpenAI API. When ``None``, an + ``OpenAIChatFormatter`` instance will be used. + client_kwargs (`dict[str, Any] | None`, defaults to `None`): + Extra keyword arguments forwarded to ``openai.AsyncClient`` + (e.g. ``timeout``, ``default_headers``, ``http_client``). + extra_body (`dict[str, Any] | None`, defaults to `None`): + Additional request body fields forwarded to + OpenAI-compatible APIs. + """ + super().__init__( + credential=credential, + model=model, + parameters=parameters or self.Parameters(), + stream=stream, + max_retries=max_retries, + retry_delay=retry_delay, + context_size=context_size, + ) + self.formatter = formatter or OpenAIChatFormatter() + self.client_kwargs = client_kwargs or {} + self.extra_body = dict(extra_body) if extra_body is not None else None + + @classmethod + def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]: + import openai + + return ( + openai.APIConnectionError, + openai.APITimeoutError, + openai.RateLimitError, + openai.InternalServerError, + ) + + async def _call_api( + self, + model_name: str, + messages: list[Msg], + tools: list[dict] | None = None, + tool_choice: ToolChoice | None = None, + **generate_kwargs: Any, + ) -> ChatResponse | AsyncGenerator[ChatResponse, None]: + """Call the OpenAI Chat Completions API. + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list`): + A list of message dicts with ``role`` and ``content`` keys. + tools (`list[dict]`, default `None`): + The tools JSON schemas. + tool_choice (`ToolChoice | None`, optional): + Controls which (if any) tool is called by the model. + **generate_kwargs (`Any`): + Extra keyword arguments forwarded to the API. + + Returns: + `ChatResponse | AsyncGenerator[ChatResponse, None]`: + A ``ChatResponse`` when streaming is disabled, or an async + generator of ``ChatResponse`` objects when streaming is + enabled. + """ + import openai + + client = openai.AsyncClient( + **{ + "api_key": self.credential.api_key.get_secret_value(), + "organization": self.credential.organization, + "base_url": self.credential.base_url, + **self.client_kwargs, + }, + ) + + formatted_messages = await self.formatter.format(messages) + + kwargs: dict[str, Any] = { + "model": model_name, + "messages": formatted_messages, + "stream": self.stream, + } + + if self.parameters.max_tokens is not None: + kwargs["max_tokens"] = self.parameters.max_tokens + + if self.parameters.temperature is not None: + kwargs["temperature"] = self.parameters.temperature + + if self.parameters.top_p is not None: + kwargs["top_p"] = self.parameters.top_p + + if ( + self.parameters.thinking_enable + and self.parameters.reasoning_effort + ): + kwargs["reasoning_effort"] = self.parameters.reasoning_effort + + if self.parameters.voice is not None: + # Requesting audio output implies ``modalities`` must include + # ``"audio"``; set it automatically so callers don't have to. + # ``format`` is forced to ``pcm16``: OpenAI streaming only + # supports ``pcm16`` (other formats raise 400), and we re-wrap + # as WAV downstream so the frontend receives a playable block. + kwargs["audio"] = { + "voice": self.parameters.voice, + "format": "pcm16", + } + kwargs["modalities"] = ["text", "audio"] + + if self.extra_body is not None: + kwargs["extra_body"] = dict(self.extra_body) + + kwargs.update(generate_kwargs) + + fmt_tools, fmt_tool_choice = self._format_tools(tools, tool_choice) + + if fmt_tools: + kwargs["tools"] = fmt_tools + if not self.parameters.parallel_tool_calls: + kwargs["parallel_tool_calls"] = False + + if fmt_tool_choice is not None: + kwargs["tool_choice"] = fmt_tool_choice + + if self.stream: + kwargs["stream_options"] = {"include_usage": True} + + start_datetime = datetime.now() + response = await client.chat.completions.create(**kwargs) + + audio_cfg = kwargs.get("audio") + audio_fmt = ( + audio_cfg.get("format", "wav") + if isinstance(audio_cfg, dict) + else "wav" + ) + + if self.stream: + # Streaming wire format is always ``pcm16`` (forced above) and we + # re-wrap it as WAV before yielding, so downstream sees ``wav``. + return self._parse_stream_response( + start_datetime, + response, + ) + + return self._parse_completion_response( + start_datetime, + response, + audio_fmt, + ) + + async def _parse_stream_response( + self, + start_datetime: datetime, + response: AsyncStream, + ) -> AsyncGenerator[ChatResponse, None]: + """Parse the OpenAI Chat streaming response. + + Upstream sends raw PCM16 (24kHz, 16-bit mono — OpenAI's only + streaming-supported audio format). We prefix the first audio + chunk with a streaming WAV header so the frontend can start + playback immediately, and assemble a fixed-size WAV on the + final chunk for non-streaming consumers. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`AsyncStream`): + The OpenAI async stream object. + + Yields: + `ChatResponse`: + Incremental ``ChatResponse`` objects with ``is_last=False`` + followed by a final one with ``is_last=True``. + """ + usage = None + response_id: str | None = None + # All delta should have the same block identifier + acc_text = TextBlock(text="") + acc_thinking = ThinkingBlock(thinking="") + acc_tool_calls: OrderedDict = OrderedDict() + # Raw audio bytes accumulated across chunks. Storing the decoded + # form (rather than concatenated base64 strings) avoids corrupting + # the byte stream when an intermediate chunk happens to carry + # base64 padding (``=``). + acc_audio_data: bytearray = bytearray() + audio_block_id: str | None = None + # ``True`` once the first audio chunk has been prefixed with a + # streaming WAV header and yielded. + audio_header_sent: bool = False + + async with response as stream: + async for chunk in stream: + if chunk.usage: + u = chunk.usage + details = getattr(u, "prompt_tokens_details", None) + usage = ChatUsage( + input_tokens=u.prompt_tokens, + output_tokens=u.completion_tokens, + time=(datetime.now() - start_datetime).total_seconds(), + cache_input_tokens=getattr( + details, + "cached_tokens", + 0, + ) + if details + else 0, + ) + + # Capture response_id from the first chunk that carries it + response_id = response_id or getattr(chunk, "id", None) + + if not chunk.choices: + continue + + choice = chunk.choices[0] + delta = choice.delta + + delta_thinking = getattr(delta, "reasoning_content", None) + if not isinstance(delta_thinking, str): + delta_thinking = getattr(delta, "reasoning", None) + if not isinstance(delta_thinking, str): + delta_thinking = "" + + delta_text = getattr(delta, "content", None) or "" + + # Collect audio output (delta.audio.data / + # delta.audio.transcript) + delta_audio_block: DataBlock | None = None + transcript_chunk: str = "" + delta_audio = getattr(delta, "audio", None) + if delta_audio is not None: + if isinstance(delta_audio, dict): + audio_chunk = delta_audio.get("data", "") + transcript_chunk = delta_audio.get("transcript", "") + else: + audio_chunk = getattr(delta_audio, "data", "") or "" + transcript_chunk = ( + getattr(delta_audio, "transcript", "") or "" + ) + if audio_chunk: + if audio_block_id is None: + audio_block_id = _generate_id() + pcm_bytes = base64.b64decode(audio_chunk) + acc_audio_data += pcm_bytes + if not audio_header_sent: + payload = _build_streaming_wav_header() + pcm_bytes + audio_header_sent = True + else: + payload = pcm_bytes + delta_audio_block = DataBlock( + id=audio_block_id, + source=Base64Source( + data=base64.b64encode(payload).decode( + "ascii", + ), + media_type="audio/wav", + ), + ) + # Omni models deliver text via ``delta.audio.transcript`` + # (not ``delta.content``); fold it into ``delta_text`` so + # the agent's streaming pipeline emits ``TextBlockDelta`` + # events alongside the audio chunks. + if transcript_chunk: + delta_text += transcript_chunk + + acc_thinking.thinking += delta_thinking + acc_text.text += delta_text + + delta_tool_call_blocks: List[ToolCallBlock] = [] + for tool_call in getattr(delta, "tool_calls", None) or []: + idx = tool_call.index + args = tool_call.function.arguments or "" + if idx in acc_tool_calls: + acc_tool_calls[idx]["input"] += args + else: + acc_tool_calls[idx] = { + "id": tool_call.id, + "name": tool_call.function.name, + "input": args, + } + tc = acc_tool_calls[idx] + delta_tool_call_blocks.append( + ToolCallBlock( + id=tc["id"], + name=tc["name"], + input=args, + ), + ) + + delta_contents: List[ + TextBlock | ToolCallBlock | ThinkingBlock | DataBlock + ] = [] + if delta_thinking: + delta_contents.append( + ThinkingBlock( + id=acc_thinking.id, + thinking=delta_thinking, + ), + ) + if delta_text: + delta_contents.append( + TextBlock(id=acc_text.id, text=delta_text), + ) + delta_contents.extend(delta_tool_call_blocks) + if delta_audio_block is not None: + delta_contents.append(delta_audio_block) + + if delta_contents: + _kwargs: dict[str, Any] = { + "content": delta_contents, + "usage": usage, + "is_last": False, + } + if response_id: + _kwargs["id"] = response_id + yield ChatResponse(**_kwargs) + + final_contents: List[ + TextBlock | ToolCallBlock | ThinkingBlock | DataBlock + ] = [] + if acc_thinking.thinking: + final_contents.append(acc_thinking) + if acc_text.text: + final_contents.append(acc_text) + for tc in acc_tool_calls.values(): + final_contents.append( + ToolCallBlock(id=tc["id"], name=tc["name"], input=tc["input"]), + ) + if acc_audio_data: + # PCM bytes were already streamed incrementally above (first + # chunk prefixed with a WAV header). Here we also assemble a + # standalone fixed-size WAV and attach it to the ``is_last`` + # chunk so callers that consume the model directly (i.e. + # without going through ``Agent``, which filters audio blocks + # out of context) get a self-contained audio block for + # downstream serialization / display. + buf = io.BytesIO() + with wave.open(buf, "wb") as wav: + wav.setnchannels(1) + wav.setsampwidth(2) + wav.setframerate(24000) + wav.writeframes(bytes(acc_audio_data)) + final_contents.append( + DataBlock( + id=audio_block_id, + source=Base64Source( + data=base64.b64encode(buf.getvalue()).decode("ascii"), + media_type="audio/wav", + ), + ), + ) + + _final_kwargs: dict[str, Any] = { + "content": final_contents, + "usage": usage, + "is_last": True, + } + if response_id: + _final_kwargs["id"] = response_id + yield ChatResponse(**_final_kwargs) + + def _parse_completion_response( + self, + start_datetime: datetime, + response: ChatCompletion, + audio_format: str = "wav", + ) -> ChatResponse: + """Parse the OpenAI Chat non-streaming response. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`ChatCompletion`): + The OpenAI chat completion object. + audio_format (`str`, defaults to ``"wav"``): + The audio format requested (used to set the media type on + the output ``DataBlock``). + + Returns: + `ChatResponse`: + A single ``ChatResponse`` with ``is_last=True``. + """ + content_blocks: List[ + TextBlock | ToolCallBlock | ThinkingBlock | DataBlock + ] = [] + + if response.choices: + choice = response.choices[0] + reasoning = getattr(choice.message, "reasoning_content", None) + if not isinstance(reasoning, str): + reasoning = getattr(choice.message, "reasoning", None) + if isinstance(reasoning, str) and reasoning: + content_blocks.append(ThinkingBlock(thinking=reasoning)) + + if choice.message.content: + content_blocks.append(TextBlock(text=choice.message.content)) + + # Extract audio output (message.audio.data / + # message.audio.transcript) + audio_obj = getattr(choice.message, "audio", None) + if audio_obj is not None: + if isinstance(audio_obj, dict): + audio_data = audio_obj.get("data", "") + audio_transcript = audio_obj.get("transcript", "") + else: + audio_data = getattr(audio_obj, "data", "") or "" + audio_transcript = ( + getattr(audio_obj, "transcript", "") or "" + ) + if not choice.message.content and audio_transcript: + content_blocks.append(TextBlock(text=audio_transcript)) + if audio_data: + content_blocks.append( + DataBlock( + source=Base64Source( + data=audio_data, + media_type=f"audio/{audio_format}", + ), + ), + ) + + for tool_call in choice.message.tool_calls or []: + content_blocks.append( + ToolCallBlock( + id=tool_call.id, + name=tool_call.function.name, + input=tool_call.function.arguments, + ), + ) + + usage = None + if response.usage: + u = response.usage + details = getattr(u, "prompt_tokens_details", None) + usage = ChatUsage( + input_tokens=u.prompt_tokens, + output_tokens=u.completion_tokens, + time=(datetime.now() - start_datetime).total_seconds(), + cache_input_tokens=getattr( + details, + "cached_tokens", + 0, + ) + if details + else 0, + ) + + resp_kwargs: dict[str, Any] = { + "content": content_blocks, + "is_last": True, + "usage": usage, + } + response_id = getattr(response, "id", None) + if response_id: + resp_kwargs["id"] = response_id + + return ChatResponse(**resp_kwargs) + + async def _call_api_with_structured_output( + self, + model_name: str, + messages: list[Msg], + structured_model: Type[BaseModel] | dict, + tool_choice: ToolChoice | None = None, + **kwargs: Any, + ) -> StructuredResponse: + """OpenAI-compatible override for structured output. + + Some third-party providers that expose an OpenAI-compatible API (e.g. + DeepSeek via DashScope) reject forced ``tool_choice`` when thinking + mode is active. When such a ``BadRequestError`` is encountered, this + method automatically retries with ``tool_choice="auto"``. + """ + import openai + + try: + return await super()._call_api_with_structured_output( + model_name=model_name, + messages=messages, + structured_model=structured_model, + tool_choice=tool_choice, + **kwargs, + ) + except openai.BadRequestError as e: + if "tool_choice" not in str(e): + raise + # Thinking mode rejects forced tool_choice; fall back to auto + warnings.warn( + f"Forced tool_choice rejected by provider ({e}), " + "retrying with tool_choice='auto'.", + stacklevel=2, + ) + return await super()._call_api_with_structured_output( + model_name=model_name, + messages=messages, + structured_model=structured_model, + tool_choice=ToolChoice(mode="auto"), + **kwargs, + ) + + def _format_tools( + self, + tools: list[dict] | None, + tool_choice: ToolChoice | None, + ) -> tuple[list[dict] | None, str | dict | None]: + """Validate, filter, and format tools and tool_choice for the OpenAI + Chat Completions API. + + When ``tool_choice.tools`` is specified the schemas list is filtered + to only those tools. When ``tool_choice.mode`` is a specific tool name + (str) the model is forced to call exactly that tool without needing to + filter the list, preserving prompt-cache efficiency. + + Tool parameter schemas are flattened (``$ref`` / ``$defs`` resolved + inline) so that providers which do not support JSON Schema references + (e.g. GLM-5.x via OpenCode Go) receive a self-contained schema. + + Args: + tools (`list[dict] | None`, optional): + The raw tool schemas. + tool_choice (`ToolChoice | None`, optional): + The tool choice configuration. + + Returns: + `tuple[list[dict] | None, str | dict | None]`: + A tuple of (formatted_tools, formatted_tool_choice). + """ + if tool_choice and tools: + self._validate_tool_choice(tool_choice, tools) + if tool_choice.tools: + allowed = set(tool_choice.tools) + tools = [t for t in tools if t["function"]["name"] in allowed] + + if tools: + tools = self._flatten_tool_schemas(tools) + + if not tool_choice: + return tools, None + + mode = tool_choice.mode + + if mode not in _TOOL_CHOICE_LITERAL_MODES: + return tools, {"type": "function", "function": {"name": mode}} + + return tools, mode + + @staticmethod + def _flatten_tool_schemas( + tools: list[dict], + ) -> list[dict]: + """Inline ``$ref`` / ``$defs`` in each tool's parameter schema. + + Args: + tools (`list[dict]`): + The list of tool dicts, each with a ``"function"`` key + containing the tool name, description and ``"parameters"`` + JSON schema. + + Returns: + `list[dict]`: + A new list where each tool's ``parameters`` schema has all + local ``$ref`` / ``$defs`` resolved inline. Tools whose + schema contained no references are returned unchanged (same + object identity). + """ + result = [] + for tool in tools: + func = tool.get("function") + if not isinstance(func, dict): + result.append(tool) + continue + params = func.get("parameters") + if not isinstance(params, dict): + result.append(tool) + continue + flat = _flatten_json_schema(params) + if flat is not params: + tool = {**tool, "function": {**func, "parameters": flat}} + result.append(tool) + return result diff --git a/src/agentscope/model/_openai_chat/_models/gpt-4.1-mini.yaml b/src/agentscope/model/_openai_chat/_models/gpt-4.1-mini.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7fd164ffa6c5108b8821eb0630e82db88efa8af7 --- /dev/null +++ b/src/agentscope/model/_openai_chat/_models/gpt-4.1-mini.yaml @@ -0,0 +1,28 @@ +name: gpt-4.1-mini +label: GPT-4.1 Mini +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + +context_size: 1047576 +output_size: 32768 + +parameter_overrides: + max_tokens: + maximum: 32768 + thinking_enable: + hidden: true + reasoning_effort: + hidden: true + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_openai_chat/_models/gpt-4.1-nano.yaml b/src/agentscope/model/_openai_chat/_models/gpt-4.1-nano.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a2c10fd77666d59c69efb5ae6cf4690a3d74ed2b --- /dev/null +++ b/src/agentscope/model/_openai_chat/_models/gpt-4.1-nano.yaml @@ -0,0 +1,28 @@ +name: gpt-4.1-nano +label: GPT-4.1 Nano +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + +context_size: 1047576 +output_size: 32768 + +parameter_overrides: + max_tokens: + maximum: 32768 + thinking_enable: + hidden: true + reasoning_effort: + hidden: true + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_openai_chat/_models/gpt-4.1.yaml b/src/agentscope/model/_openai_chat/_models/gpt-4.1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4fde2ecb0f10fae72f114a658b58b8393ae6d951 --- /dev/null +++ b/src/agentscope/model/_openai_chat/_models/gpt-4.1.yaml @@ -0,0 +1,26 @@ +name: gpt-4.1 +label: GPT-4.1 +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + +context_size: 1047576 +output_size: 32768 + +parameter_overrides: + max_tokens: + maximum: 32768 + thinking_enable: + hidden: true + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_openai_chat/_models/gpt-4o-mini.yaml b/src/agentscope/model/_openai_chat/_models/gpt-4o-mini.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f75e57d093ed805d914fa49ab9777ded0bf3447f --- /dev/null +++ b/src/agentscope/model/_openai_chat/_models/gpt-4o-mini.yaml @@ -0,0 +1,28 @@ +name: gpt-4o-mini +label: GPT-4o Mini +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + +context_size: 128000 +output_size: 16384 + +parameter_overrides: + max_tokens: + maximum: 16384 + thinking_enable: + hidden: true + reasoning_effort: + hidden: true + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_openai_chat/_models/gpt-4o.yaml b/src/agentscope/model/_openai_chat/_models/gpt-4o.yaml new file mode 100644 index 0000000000000000000000000000000000000000..432188f72ac5a7e8a435d6d02408d5e344df929e --- /dev/null +++ b/src/agentscope/model/_openai_chat/_models/gpt-4o.yaml @@ -0,0 +1,28 @@ +name: gpt-4o +label: GPT-4o +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + +context_size: 128000 +output_size: 16384 + +parameter_overrides: + max_tokens: + maximum: 16384 + thinking_enable: + hidden: true + reasoning_effort: + hidden: true + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_openai_chat/_models/gpt-5.4.yaml b/src/agentscope/model/_openai_chat/_models/gpt-5.4.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6fae4fbb5f5189ef537bb9c0e94eacceeb4945d3 --- /dev/null +++ b/src/agentscope/model/_openai_chat/_models/gpt-5.4.yaml @@ -0,0 +1,26 @@ +name: gpt-5.4 +label: GPT-5.4 +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + +context_size: 1050000 +output_size: 128000 + +parameter_overrides: + max_tokens: + maximum: 128000 + thinking_enable: + hidden: true + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_openai_chat/_models/gpt-5.5.yaml b/src/agentscope/model/_openai_chat/_models/gpt-5.5.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bece1b575446272b92c7f2e0f5b76662452b6057 --- /dev/null +++ b/src/agentscope/model/_openai_chat/_models/gpt-5.5.yaml @@ -0,0 +1,30 @@ +name: gpt-5.5 +label: GPT-5.5 +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + - audio/mp3 + - audio/wav + +output_types: + - text/plain + +context_size: 1050000 +output_size: 128000 + +parameter_overrides: + max_tokens: + maximum: 128000 + thinking_enable: + hidden: true + reasoning_effort: + hidden: true + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_openai_chat/_models/gpt-audio-mini.yaml b/src/agentscope/model/_openai_chat/_models/gpt-audio-mini.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e743bcf0640e9dab7c16f68c6a7ad9e25ae4da4f --- /dev/null +++ b/src/agentscope/model/_openai_chat/_models/gpt-audio-mini.yaml @@ -0,0 +1,31 @@ +name: gpt-audio-mini +label: GPT Audio Mini +status: active + +input_types: + - text/plain + - audio/wav + - audio/mp3 + +output_types: + - text/plain + - audio/wav + +context_size: 128000 +output_size: 16384 + +parameter_overrides: + max_tokens: + maximum: 16384 + thinking_enable: + hidden: true + # Source: https://platform.openai.com/docs/guides/text-to-speech + voice: + default: alloy + enum: + - alloy + - echo + - nova + - shimmer + - sage + - verse diff --git a/src/agentscope/model/_openai_chat/_models/o3.yaml b/src/agentscope/model/_openai_chat/_models/o3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6aa5fcf91c16ec1b3dace48bf0a7f84259be1066 --- /dev/null +++ b/src/agentscope/model/_openai_chat/_models/o3.yaml @@ -0,0 +1,25 @@ +name: o3 +label: o3 +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + - application/x-thinking + +context_size: 200000 +output_size: 100000 + +parameter_overrides: + max_tokens: + maximum: 100000 + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_openai_chat/_models/o4-mini.yaml b/src/agentscope/model/_openai_chat/_models/o4-mini.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c9254fe1d84ada376256d58d18c6347ea0c76870 --- /dev/null +++ b/src/agentscope/model/_openai_chat/_models/o4-mini.yaml @@ -0,0 +1,25 @@ +name: o4-mini +label: o4-mini +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + - application/x-thinking + +context_size: 200000 +output_size: 100000 + +parameter_overrides: + max_tokens: + maximum: 100000 + # Voice only applies to omni-style models; hide the field for this + # non-omni model so the frontend popover doesn't render it. + voice: + hidden: true diff --git a/src/agentscope/model/_openai_response/__init__.py b/src/agentscope/model/_openai_response/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1a744576c3da73de311e1635b318acd39e341350 --- /dev/null +++ b/src/agentscope/model/_openai_response/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- +"""The OpenAI Responses API modules.""" + +from ._model import OpenAIResponseModel + +__all__ = [ + "OpenAIResponseModel", +] diff --git a/src/agentscope/model/_openai_response/_model.py b/src/agentscope/model/_openai_response/_model.py new file mode 100644 index 0000000000000000000000000000000000000000..7014b83ee27932f8e6a566046c4add50092533b0 --- /dev/null +++ b/src/agentscope/model/_openai_response/_model.py @@ -0,0 +1,528 @@ +# -*- coding: utf-8 -*- +"""The OpenAI Responses API chat model implementation.""" +from datetime import datetime +from typing import Literal, Any, AsyncGenerator, List, TYPE_CHECKING, Type + +from pydantic import BaseModel, Field + +from .._base import ChatModelBase, _TOOL_CHOICE_LITERAL_MODES +from .._model_response import ChatResponse +from .._model_usage import ChatUsage +from ...credential import OpenAICredential +from ...formatter import FormatterBase, OpenAIResponseFormatter +from ...message import Msg, ThinkingBlock, ToolCallBlock, TextBlock +from ...tool import ToolChoice + +if TYPE_CHECKING: + from openai.types.responses import Response + from openai.types.responses import ResponseStreamEvent + from openai import AsyncStream +else: + Response = Any + ResponseStreamEvent = Any + AsyncStream = Any + +# kwargs accepted by Chat Completions but NOT by the Responses API. +_RESPONSES_UNSUPPORTED_KWARGS = frozenset({"modalities", "audio"}) + + +class OpenAIResponseModel(ChatModelBase): + """The OpenAI Responses API chat model. + + Compared with the Chat Completions API, the Responses API provides + first-class streaming events for reasoning / thinking, text output + and function-call arguments, which makes it a natural fit for models + that expose chain-of-thought reasoning (e.g. ``o3``, ``o4-mini``). + """ + + class Parameters(BaseModel): + """The parameters for the OpenAI Response API model.""" + + max_tokens: int | None = Field( + default=None, + title="Max Tokens", + description="The maximum number of output tokens.", + gt=0, + ) + + thinking_enable: bool = Field( + default=False, + title="Thinking", + description=( + "Whether to enable reasoning for reasoning models " + "(e.g. o3, o4-mini, gpt-5.5). Use reasoning_effort to " + "control the depth of reasoning." + ), + ) + + reasoning_effort: ( + Literal["none", "minimal", "low", "medium", "high", "xhigh"] | None + ) = Field( + default=None, + title="Reasoning Effort", + description=( + "Controls the depth of reasoning for reasoning models " + "(e.g. o3, o4-mini, gpt-5.5). Supported values are " + "model-dependent and may include: none, minimal, low, " + "medium, high, xhigh." + ), + ) + + temperature: float | None = Field( + default=None, + title="Temperature", + description="The temperature for the LLM output.", + ge=0, + le=2, + ) + + type: Literal["openai_response"] = "openai_response" + """The type of the chat model.""" + + def __init__( + self, + credential: OpenAICredential, + model: str, + parameters: "OpenAIResponseModel.Parameters | None" = None, + stream: bool = True, + max_retries: int = 3, + retry_delay: float = 1.0, + context_size: int = 200000, + formatter: FormatterBase | None = None, + client_kwargs: dict[str, Any] | None = None, + ) -> None: + """Initialize the OpenAI Responses API chat model. + + Args: + credential (`OpenAICredential`): + The OpenAI credential used to authenticate API calls. + model (`str`): + The OpenAI model name, e.g. ``o3`` or ``o4-mini``. + parameters (`OpenAIResponseModel.Parameters | None`, defaults \ + to `None`): + The OpenAI Responses API parameters. When ``None``, the + default parameters will be used. + stream (`bool`, defaults to `True`): + Whether to enable streaming output. + max_retries (`int`, defaults to `3`): + The maximum number of retries for the OpenAI Responses API. + retry_delay (`float`, defaults to `1.0`): + Seconds to sleep between retry attempts. + context_size (`int`, defaults to `200000`): + The model context size used for context compression. + formatter (`FormatterBase | None`, defaults to `None`): + The formatter that converts ``Msg`` objects to the format + required by the OpenAI Responses API. When ``None``, an + ``OpenAIResponseFormatter`` instance will be used. + client_kwargs (`dict[str, Any] | None`, defaults to `None`): + Extra keyword arguments forwarded to ``openai.AsyncClient`` + (e.g. ``timeout``, ``default_headers``, ``http_client``). + """ + super().__init__( + credential=credential, + model=model, + parameters=parameters or self.Parameters(), + stream=stream, + max_retries=max_retries, + retry_delay=retry_delay, + context_size=context_size, + ) + self.formatter = formatter or OpenAIResponseFormatter() + self.client_kwargs = client_kwargs or {} + + @classmethod + def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]: + import openai + + return ( + openai.APIConnectionError, + openai.APITimeoutError, + openai.RateLimitError, + openai.InternalServerError, + ) + + async def _call_api( + self, + model_name: str, + messages: list[Msg], + tools: list[dict] | None = None, + tool_choice: ToolChoice | None = None, + **generate_kwargs: Any, + ) -> ChatResponse | AsyncGenerator[ChatResponse, None]: + """Call the OpenAI Responses API. + + Args: + model_name (`str`): + The model name to use for this call. + messages (`list`): + A list of input items for the Responses API. + tools (`list[dict]`, default `None`): + The tools JSON schemas. + tool_choice (`ToolChoice | None`, optional): + Controls which (if any) tool is called by the model. + **generate_kwargs (`Any`): + Extra keyword arguments forwarded to the API. + + Returns: + `ChatResponse | AsyncGenerator[ChatResponse, None]`: + A ``ChatResponse`` when streaming is disabled, or an async + generator of ``ChatResponse`` objects when streaming is + enabled. + """ + import openai + + client = openai.AsyncClient( + **{ + "api_key": self.credential.api_key.get_secret_value(), + "organization": self.credential.organization, + "base_url": self.credential.base_url, + **self.client_kwargs, + }, + ) + + formatted_messages = await self.formatter.format(messages) + + api_kwargs: dict[str, Any] = { + "model": model_name, + "input": formatted_messages, + "stream": self.stream, + } + + if self.parameters.max_tokens is not None: + api_kwargs["max_output_tokens"] = self.parameters.max_tokens + + if self.parameters.temperature is not None: + api_kwargs["temperature"] = self.parameters.temperature + + if ( + self.parameters.thinking_enable + and self.parameters.reasoning_effort + ): + api_kwargs["reasoning"] = { + "effort": self.parameters.reasoning_effort, + } + + # The Responses API does not yet support audio output + # (modalities / audio params). Strip them so callers that + # mistakenly pass Chat-Completions-style audio kwargs don't + # trigger a TypeError. + # https://developers.openai.com/api/docs/guides/migrate-to-responses + api_kwargs.update( + { + k: v + for k, v in generate_kwargs.items() + if k not in _RESPONSES_UNSUPPORTED_KWARGS + }, + ) + + fmt_tools, fmt_tool_choice = self._format_tools(tools, tool_choice) + if fmt_tools is not None: + api_kwargs["tools"] = fmt_tools + if fmt_tool_choice is not None: + api_kwargs["tool_choice"] = fmt_tool_choice + + start_datetime = datetime.now() + response = await client.responses.create(**api_kwargs) + + if self.stream: + return self._parse_stream_response(start_datetime, response) + + return self._parse_completion_response(start_datetime, response) + + async def _parse_stream_response( + self, + start_datetime: datetime, + response: "AsyncStream[ResponseStreamEvent]", + ) -> AsyncGenerator[ChatResponse, None]: + """Parse the OpenAI Responses API streaming response. + + Each event yields only the delta content produced by that event so + that callers see a true incremental stream, consistent with other + model implementations. The final ``response.completed`` event emits + an ``is_last=True`` response with the full accumulated state. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`AsyncStream[ResponseStreamEvent]`): + The OpenAI Responses API async stream object. + + Yields: + `ChatResponse`: + Incremental ``ChatResponse`` objects with ``is_last=False`` + followed by a final one with ``is_last=True``. + """ + usage: ChatUsage | None = None + response_id: str | None = None + # All delta should have the same block identifier + acc_text = TextBlock(text="") + acc_thinking = ThinkingBlock(thinking="") + tool_calls: dict[str, dict[str, Any]] = {} + + async for event in response: + event_type = event.type + + if response_id is None: + resp_obj = getattr(event, "response", None) + if resp_obj is not None: + response_id = getattr(resp_obj, "id", None) + + delta_contents: List[ + TextBlock | ToolCallBlock | ThinkingBlock + ] = [] + + if event_type == "response.reasoning_summary_text.delta": + # Reasoning summary text is NOT emitted by all models. + # As of 2026-05, o1 and o4-mini do not stream reasoning + # summary deltas. This handler exists for forward + # compatibility with models that do expose it. + delta = event.delta + acc_thinking.thinking += delta + delta_contents.append( + ThinkingBlock(id=acc_thinking.id, thinking=delta), + ) + + elif event_type == "response.output_text.delta": + delta = event.delta + acc_text.text += delta + delta_contents.append(TextBlock(id=acc_text.id, text=delta)) + + elif event_type == "response.output_item.added": + item = event.item + if getattr(item, "type", None) == "function_call": + # item.id → fc_xxx (item identifier, needed for + # function_call.id in multi-turn history) + # item.call_id → call_xxx (needed for + # function_call_output.call_id) + tool_calls[item.id] = { + "id": item.id, + "call_id": getattr(item, "call_id", None), + "name": getattr(item, "name", ""), + "input": "", + } + + elif event_type == "response.function_call_arguments.delta": + item_id = event.item_id + if item_id in tool_calls: + tool_calls[item_id]["input"] += event.delta + tc = tool_calls[item_id] + delta_contents.append( + ToolCallBlock( + id=tc["id"], + call_id=tc.get("call_id"), + name=tc["name"], + input=event.delta, + ), + ) + + elif event_type == "response.completed": + resp = event.response + if response_id is None: + response_id = getattr(resp, "id", None) + if resp.usage: + u = resp.usage + details = getattr(u, "input_tokens_details", None) + usage = ChatUsage( + input_tokens=u.input_tokens, + output_tokens=u.output_tokens, + time=(datetime.now() - start_datetime).total_seconds(), + cache_input_tokens=getattr( + details, + "cached_tokens", + 0, + ) + if details + else 0, + ) + # Attach reasoning item IDs from the completed response so the + # formatter can echo them back in multi-turn history. + # The Responses API requires every function_call item to be + # accompanied by its preceding reasoning item (see the + # function-calling guide). The reasoning item may have an + # empty summary when the model does not expose it (e.g. + # o1/o4-mini as of 2026-05). + for output_item in getattr(resp, "output", []): + if getattr(output_item, "type", None) == "reasoning": + reasoning_item_id = getattr(output_item, "id", None) + if reasoning_item_id: + acc_thinking = ThinkingBlock( + id=acc_thinking.id, + thinking=acc_thinking.thinking, + reasoning_item_id=reasoning_item_id, + ) + # Emit the full accumulated state as the final response + final_contents: List[ + TextBlock | ToolCallBlock | ThinkingBlock + ] = [] + if acc_thinking.thinking or getattr( + acc_thinking, + "reasoning_item_id", + None, + ): + final_contents.append(acc_thinking) + if acc_text.text: + final_contents.append(acc_text) + for tc in tool_calls.values(): + final_contents.append( + ToolCallBlock( + id=tc["id"], + call_id=tc.get("call_id"), + name=tc["name"], + input=tc["input"] or "{}", + ), + ) + final_kwargs: dict[str, Any] = { + "content": final_contents, + "is_last": True, + "usage": usage, + } + if response_id: + final_kwargs["id"] = response_id + yield ChatResponse(**final_kwargs) + return + + # Yield incremental delta for non-terminal events + if delta_contents: + chat_resp_kwargs: dict[str, Any] = { + "content": delta_contents, + "is_last": False, + "usage": usage, + } + if response_id: + chat_resp_kwargs["id"] = response_id + yield ChatResponse(**chat_resp_kwargs) + + def _parse_completion_response( + self, + start_datetime: datetime, + response: "Response", + ) -> ChatResponse: + """Parse the OpenAI Responses API non-streaming response. + + Args: + start_datetime (`datetime`): + The start datetime of the response generation. + response (`Response`): + The OpenAI Responses API response object. + + Returns: + `ChatResponse`: + A single ``ChatResponse`` with ``is_last=True``. + """ + content_blocks: List[TextBlock | ToolCallBlock | ThinkingBlock] = [] + + for item in response.output: + item_type = getattr(item, "type", None) + + if item_type == "reasoning": + reasoning_item_id = getattr(item, "id", None) + combined_summary = " ".join( + getattr(s, "text", "") + for s in getattr(item, "summary", []) + if getattr(s, "text", "") + ) + if combined_summary: + content_blocks.append( + ThinkingBlock( + type="thinking", + thinking=combined_summary, + reasoning_item_id=reasoning_item_id, + ), + ) + + elif item_type == "message": + for part in getattr(item, "content", []): + if getattr(part, "type", None) == "output_text": + content_blocks.append( + TextBlock(type="text", text=part.text), + ) + + elif item_type == "function_call": + content_blocks.append( + ToolCallBlock( + id=getattr(item, "id", ""), + call_id=getattr(item, "call_id", None), + name=item.name, + input=getattr(item, "arguments", "") or "{}", + ), + ) + + usage = None + if response.usage: + u = response.usage + details = getattr(u, "input_tokens_details", None) + usage = ChatUsage( + input_tokens=u.input_tokens, + output_tokens=u.output_tokens, + time=(datetime.now() - start_datetime).total_seconds(), + cache_input_tokens=getattr( + details, + "cached_tokens", + 0, + ) + if details + else 0, + ) + + resp_kwargs: dict[str, Any] = { + "content": content_blocks, + "is_last": True, + "usage": usage, + } + response_id = getattr(response, "id", None) + if response_id: + resp_kwargs["id"] = response_id + + return ChatResponse(**resp_kwargs) + + def _format_tools( + self, + tools: list[dict] | None, + tool_choice: ToolChoice | None, + ) -> tuple[list[dict] | None, str | dict | None]: + """Validate and format tools and tool_choice for the Responses API. + + The full ``tools`` list is always sent unchanged to maximise prompt + cache hits. When ``tool_choice.tools`` restricts the callable + subset, the ``allowed_tools`` tool_choice format is used instead of + filtering the schemas list. + + Args: + tools (`list[dict] | None`, optional): + The raw tool schemas. + tool_choice (`ToolChoice | None`, optional): + The tool choice configuration. + + Returns: + `tuple[list[dict] | None, str | dict | None]`: + A tuple of ``(formatted_tools, formatted_tool_choice)`` + ready for the Responses API. + """ + if tool_choice and tools: + self._validate_tool_choice(tool_choice, tools) + + fmt_tools = None + if tools: + fmt_tools = [ + {"type": "function", **tool["function"]} for tool in tools + ] + + if not tool_choice: + return fmt_tools, None + + mode = tool_choice.mode + + if mode not in _TOOL_CHOICE_LITERAL_MODES: + return fmt_tools, {"type": "function", "name": mode} + + if tool_choice.tools: + return fmt_tools, { + "type": "allowed_tools", + "mode": mode, + "tools": [ + {"type": "function", "name": name} + for name in tool_choice.tools + ], + } + + return fmt_tools, mode diff --git a/src/agentscope/model/_openai_response/_models/gpt-4.1-mini.yaml b/src/agentscope/model/_openai_response/_models/gpt-4.1-mini.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cd34c11fd7077a58c864484fe1d712524276d838 --- /dev/null +++ b/src/agentscope/model/_openai_response/_models/gpt-4.1-mini.yaml @@ -0,0 +1,24 @@ +name: gpt-4.1-mini +label: GPT-4.1 Mini +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + +context_size: 1047576 +output_size: 32768 + +parameter_overrides: + max_tokens: + maximum: 32768 + thinking_enable: + hidden: true + reasoning_effort: + hidden: true diff --git a/src/agentscope/model/_openai_response/_models/gpt-4.1-nano.yaml b/src/agentscope/model/_openai_response/_models/gpt-4.1-nano.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f5ad17022d7c80831d5ef8b2627e105721c5d337 --- /dev/null +++ b/src/agentscope/model/_openai_response/_models/gpt-4.1-nano.yaml @@ -0,0 +1,24 @@ +name: gpt-4.1-nano +label: GPT-4.1 Nano +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + +context_size: 1047576 +output_size: 32768 + +parameter_overrides: + max_tokens: + maximum: 32768 + thinking_enable: + hidden: true + reasoning_effort: + hidden: true diff --git a/src/agentscope/model/_openai_response/_models/gpt-4.1.yaml b/src/agentscope/model/_openai_response/_models/gpt-4.1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5967e8af7030f941da34024550d131f51dd13761 --- /dev/null +++ b/src/agentscope/model/_openai_response/_models/gpt-4.1.yaml @@ -0,0 +1,24 @@ +name: gpt-4.1 +label: GPT-4.1 +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + +context_size: 1047576 +output_size: 32768 + +parameter_overrides: + max_tokens: + maximum: 32768 + thinking_enable: + hidden: true + reasoning_effort: + hidden: true diff --git a/src/agentscope/model/_openai_response/_models/gpt-4o-mini.yaml b/src/agentscope/model/_openai_response/_models/gpt-4o-mini.yaml new file mode 100644 index 0000000000000000000000000000000000000000..75823ebb0dcaf5fec98fd1a96bbe8144aef530fd --- /dev/null +++ b/src/agentscope/model/_openai_response/_models/gpt-4o-mini.yaml @@ -0,0 +1,24 @@ +name: gpt-4o-mini +label: GPT-4o Mini +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + +context_size: 128000 +output_size: 16384 + +parameter_overrides: + max_tokens: + maximum: 16384 + thinking_enable: + hidden: true + reasoning_effort: + hidden: true diff --git a/src/agentscope/model/_openai_response/_models/gpt-4o.yaml b/src/agentscope/model/_openai_response/_models/gpt-4o.yaml new file mode 100644 index 0000000000000000000000000000000000000000..76edf7356249a38b275150ff603b7859c11b3f00 --- /dev/null +++ b/src/agentscope/model/_openai_response/_models/gpt-4o.yaml @@ -0,0 +1,24 @@ +name: gpt-4o +label: GPT-4o +status: active + +input_types: + - text/plain + - image/jpeg + - image/png + - image/gif + - image/webp + +output_types: + - text/plain + +context_size: 128000 +output_size: 16384 + +parameter_overrides: + max_tokens: + maximum: 16384 + thinking_enable: + hidden: true + reasoning_effort: + hidden: true