agentstack-sdk 0.5.2rc2__py3-none-any.whl

agentstack_sdk/a2a/extensions/exceptions.py

@@ -0,0 +1,11 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from agentstack_sdk.a2a.extensions.base import BaseExtensionSpec
+
+
+class ExtensionError(Exception):
+    extension: BaseExtensionSpec
+
+    def __init__(self, spec: BaseExtensionSpec, message: str):
+        super().__init__(f"Exception in extension '{spec.URI}': \n{message}")

agentstack_sdk/a2a/extensions/interactions/__init__.py

@@ -0,0 +1,4 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from .approval import *

agentstack_sdk/a2a/extensions/interactions/approval.py

@@ -0,0 +1,125 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import uuid
+from types import NoneType
+from typing import TYPE_CHECKING, Annotated, Any, Literal
+
+import a2a.types
+from mcp import Implementation, Tool
+from pydantic import BaseModel, Discriminator, Field, TypeAdapter
+
+from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
+from agentstack_sdk.a2a.types import AgentMessage, InputRequired
+
+if TYPE_CHECKING:
+    from agentstack_sdk.server.context import RunContext
+
+
+class ApprovalRejectionError(RuntimeError):
+    pass
+
+
+class GenericApprovalRequest(BaseModel):
+    action: Literal["generic"] = "generic"
+
+    title: str | None = Field(None, description="A human-readable title for the action being approved.")
+    description: str | None = Field(None, description="A human-readable description of the action being approved.")
+
+
+class ToolCallServer(BaseModel):
+    name: str = Field(description="The programmatic name of the server.")
+    title: str | None = Field(description="A human-readable title for the server.")
+    version: str = Field(description="The version of the server.")
+
+
+class ToolCallApprovalRequest(BaseModel):
+    action: Literal["tool-call"] = "tool-call"
+
+    title: str | None = Field(None, description="A human-readable title for the tool call being approved.")
+    description: str | None = Field(None, description="A human-readable description of the tool call being approved.")
+    name: str = Field(description="The programmatic name of the tool.")
+    input: dict[str, Any] | None = Field(description="The input for the tool.")
+    server: ToolCallServer | None = Field(None, description="The server executing the tool.")
+
+    @staticmethod
+    def from_mcp_tool(
+        tool: Tool, input: dict[str, Any] | None, server: Implementation | None = None
+    ) -> ToolCallApprovalRequest:
+        return ToolCallApprovalRequest(
+            name=tool.name,
+            title=tool.annotations.title if tool.annotations else None,
+            description=tool.description,
+            input=input,
+            server=ToolCallServer(name=server.name, title=server.title, version=server.version) if server else None,
+        )
+
+
+ApprovalRequest = Annotated[GenericApprovalRequest | ToolCallApprovalRequest, Discriminator("action")]
+
+
+class ApprovalResponse(BaseModel):
+    decision: Literal["approve", "reject"]
+
+    @property
+    def approved(self) -> bool:
+        return self.decision == "approve"
+
+    def raise_on_rejection(self) -> None:
+        if self.decision == "reject":
+            raise ApprovalRejectionError("Approval request has been rejected")
+
+
+class ApprovalExtensionParams(BaseModel):
+    pass
+
+
+class ApprovalExtensionSpec(BaseExtensionSpec[ApprovalExtensionParams]):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/interactions/approval/v1"
+
+
+class ApprovalExtensionMetadata(BaseModel):
+    pass
+
+
+class ApprovalExtensionServer(BaseExtensionServer[ApprovalExtensionSpec, ApprovalExtensionMetadata]):
+    def create_request_message(self, *, request: ApprovalRequest):
+        return AgentMessage(text="Approval requested", metadata={self.spec.URI: request.model_dump(mode="json")})
+
+    def parse_response(self, *, message: a2a.types.Message):
+        if not message.metadata or not (data := message.metadata.get(self.spec.URI)):
+            raise ValueError("Approval response data is missing")
+        return ApprovalResponse.model_validate(data)
+
+    async def request_approval(
+        self,
+        request: ApprovalRequest,
+        *,
+        context: RunContext,
+    ) -> ApprovalResponse:
+        message = self.create_request_message(request=request)
+        message = await context.yield_async(InputRequired(message=message))
+        if not message:
+            raise RuntimeError("Yield did not return a message")
+        return self.parse_response(message=message)
+
+
+class ApprovalExtensionClient(BaseExtensionClient[ApprovalExtensionSpec, NoneType]):
+    def create_response_message(self, *, response: ApprovalResponse, task_id: str | None):
+        return a2a.types.Message(
+            message_id=str(uuid.uuid4()),
+            role=a2a.types.Role.user,
+            parts=[],
+            task_id=task_id,
+            metadata={self.spec.URI: response.model_dump(mode="json")},
+        )
+
+    def parse_request(self, *, message: a2a.types.Message):
+        if not message.metadata or not (data := message.metadata.get(self.spec.URI)):
+            raise ValueError("Approval request data is missing")
+        return TypeAdapter(ApprovalRequest).validate_python(data)
+
+    def metadata(self) -> dict[str, Any]:
+        return {self.spec.URI: ApprovalExtensionMetadata().model_dump(mode="json")}

agentstack_sdk/a2a/extensions/services/__init__.py

@@ -0,0 +1,8 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from .embedding import *
+from .form import *
+from .llm import *
+from .mcp import *
+from .platform import *

agentstack_sdk/a2a/extensions/services/embedding.py

@@ -0,0 +1,106 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import re
+from types import NoneType
+from typing import TYPE_CHECKING, Any, Self
+
+import pydantic
+from a2a.server.agent_execution.context import RequestContext
+from a2a.types import Message as A2AMessage
+from typing_extensions import override
+
+from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
+
+if TYPE_CHECKING:
+    from agentstack_sdk.server.context import RunContext
+
+
+class EmbeddingFulfillment(pydantic.BaseModel):
+    identifier: str | None = None
+    """
+    Name of the model for identification and optimization purposes. Usually corresponds to LiteLLM identifiers.
+    Should be the name of the provider slash name of the model as it appears in the API.
+    Examples: openai/text-embedding-3-small, vertex_ai/textembedding-gecko, ollama/nomic-embed-text:latest
+    """
+
+    api_base: str
+    """
+    Base URL for an OpenAI-compatible API. It should provide at least /v1/chat/completions
+    """
+
+    api_key: str
+    """
+    API key to attach as a `Authorization: Bearer $api_key` header.
+    """
+
+    api_model: str
+    """
+    Model name to use with the /v1/chat/completions API.
+    """
+
+
+class EmbeddingDemand(pydantic.BaseModel):
+    description: str | None = None
+    """
+    Short description of how the model will be used, if multiple are requested.
+    Intended to be shown in the UI alongside a model picker dropdown.
+    """
+
+    suggested: tuple[str, ...] = ()
+    """
+    Identifiers of models recommended to be used. Usually corresponds to LiteLLM identifiers.
+    Should be the name of the provider slash name of the model as it appears in the API.
+    Examples: openai/text-embedding-3-small, vertex_ai/textembedding-gecko, ollama/nomic-embed-text:latest
+    """
+
+
+class EmbeddingServiceExtensionParams(pydantic.BaseModel):
+    embedding_demands: dict[str, EmbeddingDemand]
+    """Model requests that the agent requires to be provided by the client."""
+
+
+class EmbeddingServiceExtensionSpec(BaseExtensionSpec[EmbeddingServiceExtensionParams]):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/services/embedding/v1"
+
+    @classmethod
+    def single_demand(
+        cls, name: str | None = None, description: str | None = None, suggested: tuple[str, ...] = ()
+    ) -> Self:
+        return cls(
+            params=EmbeddingServiceExtensionParams(
+                embedding_demands={name or "default": EmbeddingDemand(description=description, suggested=suggested)}
+            )
+        )
+
+
+class EmbeddingServiceExtensionMetadata(pydantic.BaseModel):
+    embedding_fulfillments: dict[str, EmbeddingFulfillment] = {}
+    """Provided models corresponding to the model requests."""
+
+
+class EmbeddingServiceExtensionServer(
+    BaseExtensionServer[EmbeddingServiceExtensionSpec, EmbeddingServiceExtensionMetadata]
+):
+    @override
+    def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, request_context: RequestContext):
+        from agentstack_sdk.platform import get_platform_client
+
+        super().handle_incoming_message(message, run_context, request_context)
+        if not self.data:
+            return
+
+        for fullfilment in self.data.embedding_fulfillments.values():
+            platform_url = str(get_platform_client().base_url)
+            fullfilment.api_base = re.sub("{platform_url}", platform_url, fullfilment.api_base)
+
+
+class EmbeddingServiceExtensionClient(BaseExtensionClient[EmbeddingServiceExtensionSpec, NoneType]):
+    def fulfillment_metadata(self, *, embedding_fulfillments: dict[str, EmbeddingFulfillment]) -> dict[str, Any]:
+        return {
+            self.spec.URI: EmbeddingServiceExtensionMetadata(embedding_fulfillments=embedding_fulfillments).model_dump(
+                mode="json"
+            )
+        }

agentstack_sdk/a2a/extensions/services/form.py

@@ -0,0 +1,54 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+
+from __future__ import annotations
+
+from typing import Self, TypeVar, cast
+
+from pydantic import BaseModel, TypeAdapter
+from typing_extensions import TypedDict
+
+from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
+from agentstack_sdk.a2a.extensions.common.form import FormRender, FormResponse
+
+
+class FormDemands(TypedDict):
+    initial_form: FormRender | None
+    # TODO: We can put settings here too
+
+
+class FormServiceExtensionMetadata(BaseModel):
+    form_fulfillments: dict[str, FormResponse] = {}
+
+
+class FormServiceExtensionParams(BaseModel):
+    form_demands: FormDemands
+
+
+class FormServiceExtensionSpec(BaseExtensionSpec[FormServiceExtensionParams]):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/services/form/v1"
+
+    @classmethod
+    def demand(cls, initial_form: FormRender | None) -> Self:
+        return cls(params=FormServiceExtensionParams(form_demands={"initial_form": initial_form}))
+
+
+T = TypeVar("T")
+
+
+class FormServiceExtensionServer(BaseExtensionServer[FormServiceExtensionSpec, FormServiceExtensionMetadata]):
+    def parse_initial_form(self, *, model: type[T] = FormResponse) -> T | None:
+        if self.data is None:
+            return None
+
+        initial_form = self.data.form_fulfillments.get("initial_form")
+
+        if initial_form is None:
+            return None
+        if model is FormResponse:
+            return cast(T, initial_form)
+        return TypeAdapter(model).validate_python(dict(initial_form))
+
+
+class FormServiceExtensionClient(BaseExtensionClient[FormServiceExtensionSpec, FormRender]): ...

agentstack_sdk/a2a/extensions/services/llm.py

@@ -0,0 +1,100 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import re
+from types import NoneType
+from typing import TYPE_CHECKING, Any, Self
+
+import pydantic
+from a2a.server.agent_execution.context import RequestContext
+from a2a.types import Message as A2AMessage
+from typing_extensions import override
+
+from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
+
+if TYPE_CHECKING:
+    from agentstack_sdk.server.context import RunContext
+
+
+class LLMFulfillment(pydantic.BaseModel):
+    identifier: str | None = None
+    """
+    Name of the model for identification and optimization purposes. Usually corresponds to LiteLLM identifiers.
+    Should be the name of the provider slash name of the model as it appears in the API.
+    Examples: openai/gpt-4o, watsonx/ibm/granite-13b-chat-v2, ollama/mistral-small:22b
+    """
+
+    api_base: str
+    """
+    Base URL for an OpenAI-compatible API. It should provide at least /v1/chat/completions
+    """
+
+    api_key: str
+    """
+    API key to attach as a `Authorization: Bearer $api_key` header.
+    """
+
+    api_model: str
+    """
+    Model name to use with the /v1/chat/completions API.
+    """
+
+
+class LLMDemand(pydantic.BaseModel):
+    description: str | None = None
+    """
+    Short description of how the model will be used, if multiple are requested.
+    Intended to be shown in the UI alongside a model picker dropdown.
+    """
+
+    suggested: tuple[str, ...] = ()
+    """
+    Identifiers of models recommended to be used. Usually corresponds to LiteLLM identifiers.
+    Should be the name of the provider slash name of the model as it appears in the API.
+    Examples: openai/gpt-4o, watsonx/ibm/granite-13b-chat-v2, ollama/mistral-small:22b
+    """
+
+
+class LLMServiceExtensionParams(pydantic.BaseModel):
+    llm_demands: dict[str, LLMDemand]
+    """Model requests that the agent requires to be provided by the client."""
+
+
+class LLMServiceExtensionSpec(BaseExtensionSpec[LLMServiceExtensionParams]):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/services/llm/v1"
+
+    @classmethod
+    def single_demand(
+        cls, name: str | None = None, description: str | None = None, suggested: tuple[str, ...] = ()
+    ) -> Self:
+        return cls(
+            params=LLMServiceExtensionParams(
+                llm_demands={name or "default": LLMDemand(description=description, suggested=suggested)}
+            )
+        )
+
+
+class LLMServiceExtensionMetadata(pydantic.BaseModel):
+    llm_fulfillments: dict[str, LLMFulfillment] = {}
+    """Provided models corresponding to the model requests."""
+
+
+class LLMServiceExtensionServer(BaseExtensionServer[LLMServiceExtensionSpec, LLMServiceExtensionMetadata]):
+    @override
+    def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, request_context: RequestContext):
+        from agentstack_sdk.platform import get_platform_client
+
+        super().handle_incoming_message(message, run_context, request_context)
+        if not self.data:
+            return
+
+        for fullfilment in self.data.llm_fulfillments.values():
+            platform_url = str(get_platform_client().base_url)
+            fullfilment.api_base = re.sub("{platform_url}", platform_url, fullfilment.api_base)
+
+
+class LLMServiceExtensionClient(BaseExtensionClient[LLMServiceExtensionSpec, NoneType]):
+    def fulfillment_metadata(self, *, llm_fulfillments: dict[str, LLMFulfillment]) -> dict[str, Any]:
+        return {self.spec.URI: LLMServiceExtensionMetadata(llm_fulfillments=llm_fulfillments).model_dump(mode="json")}

agentstack_sdk/a2a/extensions/services/mcp.py

@@ -0,0 +1,193 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import re
+from contextlib import asynccontextmanager
+from types import NoneType
+from typing import TYPE_CHECKING, Annotated, Any, Literal, Self
+
+import pydantic
+from a2a.server.agent_execution.context import RequestContext
+from a2a.types import Message as A2AMessage
+from mcp.client.stdio import StdioServerParameters, stdio_client
+from mcp.client.streamable_http import streamablehttp_client
+from typing_extensions import override
+
+from agentstack_sdk.a2a.extensions.auth.oauth.oauth import OAuthExtensionServer
+from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
+from agentstack_sdk.a2a.extensions.services.platform import PlatformApiExtensionServer
+from agentstack_sdk.platform.client import get_platform_client
+from agentstack_sdk.util.logging import logger
+
+if TYPE_CHECKING:
+    from agentstack_sdk.server.context import RunContext
+
+_TRANSPORT_TYPES = Literal["streamable_http", "stdio"]
+
+_DEFAULT_DEMAND_NAME = "default"
+_DEFAULT_ALLOWED_TRANSPORTS: list[_TRANSPORT_TYPES] = ["streamable_http"]
+
+
+class StdioTransport(pydantic.BaseModel):
+    type: Literal["stdio"] = "stdio"
+
+    command: str
+    args: list[str]
+    env: dict[str, str] | None = None
+
+
+class StreamableHTTPTransport(pydantic.BaseModel):
+    type: Literal["streamable_http"] = "streamable_http"
+
+    url: str
+    headers: dict[str, str] | None = None
+
+
+MCPTransport = Annotated[StdioTransport | StreamableHTTPTransport, pydantic.Field(discriminator="type")]
+
+
+class MCPFulfillment(pydantic.BaseModel):
+    transport: MCPTransport
+
+
+class MCPDemand(pydantic.BaseModel):
+    description: str | None = None
+    """
+    Short description of how the server will be used, what tools should it contain, etc.
+    """
+
+    suggested: tuple[str, ...] = ()
+    """
+    Identifiers of servers recommended to be used. Usually corresponds to MCP StreamableHTTP URIs.
+    """
+
+    allowed_transports: list[_TRANSPORT_TYPES] = pydantic.Field(default_factory=lambda: _DEFAULT_ALLOWED_TRANSPORTS)
+    """
+    Transports allowed for the server. Specifying other transports will result in rejection.
+    """
+
+
+class MCPServiceExtensionParams(pydantic.BaseModel):
+    mcp_demands: dict[str, MCPDemand]
+    """Server requests that the agent requires to be provided by the client."""
+
+
+class MCPServiceExtensionSpec(BaseExtensionSpec[MCPServiceExtensionParams]):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/services/mcp/v1"
+
+    @classmethod
+    def single_demand(
+        cls,
+        name: str = _DEFAULT_DEMAND_NAME,
+        description: str | None = None,
+        suggested: tuple[str, ...] = (),
+        allowed_transports: list[_TRANSPORT_TYPES] | None = None,
+    ) -> Self:
+        return cls(
+            params=MCPServiceExtensionParams(
+                mcp_demands={
+                    name: MCPDemand(
+                        description=description,
+                        suggested=suggested,
+                        allowed_transports=allowed_transports or _DEFAULT_ALLOWED_TRANSPORTS,
+                    )
+                }
+            )
+        )
+
+
+class MCPServiceExtensionMetadata(pydantic.BaseModel):
+    mcp_fulfillments: dict[str, MCPFulfillment] = {}
+    """Provided servers corresponding to the server requests."""
+
+
+class MCPServiceExtensionServer(BaseExtensionServer[MCPServiceExtensionSpec, MCPServiceExtensionMetadata]):
+    @override
+    def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, request_context: RequestContext):
+        super().handle_incoming_message(message, run_context, request_context)
+        if not self.data:
+            return
+
+        platform_url = str(get_platform_client().base_url)
+        for fullfilment in self.data.mcp_fulfillments.values():
+            if fullfilment.transport.type == "streamable_http":
+                try:
+                    fullfilment.transport.url = re.sub("^{platform_url}", platform_url, str(fullfilment.transport.url))
+                except Exception:
+                    logger.warning("Platform URL substitution failed", exc_info=True)
+
+    @override
+    def parse_client_metadata(self, message: A2AMessage) -> MCPServiceExtensionMetadata | None:
+        metadata = super().parse_client_metadata(message)
+        if metadata:
+            for name, demand in self.spec.params.mcp_demands.items():
+                if not (fulfillment := metadata.mcp_fulfillments.get(name)):
+                    continue
+                if fulfillment.transport.type not in demand.allowed_transports:
+                    raise ValueError(f'Transport "{fulfillment.transport.type}" not allowed for demand "{name}"')
+        return metadata
+
+    def _get_oauth_server(self):
+        for dependency in self._dependencies.values():
+            if isinstance(dependency, OAuthExtensionServer):
+                return dependency
+        return None
+
+    def _get_platform_server(self):
+        for dependency in self._dependencies.values():
+            if isinstance(dependency, PlatformApiExtensionServer):
+                return dependency
+        return None
+
+    @asynccontextmanager
+    async def create_client(self, demand: str = _DEFAULT_DEMAND_NAME):
+        fulfillment = self.data.mcp_fulfillments.get(demand) if self.data else None
+
+        if not fulfillment:
+            yield None
+            return
+
+        transport = fulfillment.transport
+
+        if isinstance(transport, StdioTransport):
+            async with stdio_client(
+                server=StdioServerParameters(command=transport.command, args=transport.args, env=transport.env)
+            ) as (
+                read,
+                write,
+            ):
+                yield (read, write)
+        elif isinstance(transport, StreamableHTTPTransport):
+            async with streamablehttp_client(
+                url=transport.url,
+                headers=transport.headers,
+                auth=await self._create_auth(transport),
+            ) as (
+                read,
+                write,
+                _,
+            ):
+                yield (read, write)
+        else:
+            raise NotImplementedError("Unsupported transport")
+
+    async def _create_auth(self, transport: StreamableHTTPTransport):
+        platform = self._get_platform_server()
+        if (
+            platform
+            and platform.data
+            and platform.data.base_url
+            and transport.url.startswith(str(platform.data.base_url))
+        ):
+            return await platform.create_httpx_auth()
+        oauth = self._get_oauth_server()
+        if oauth:
+            return await oauth.create_httpx_auth(resource_url=pydantic.AnyUrl(transport.url))
+        return None
+
+
+class MCPServiceExtensionClient(BaseExtensionClient[MCPServiceExtensionSpec, NoneType]):
+    def fulfillment_metadata(self, *, mcp_fulfillments: dict[str, MCPFulfillment]) -> dict[str, Any]:
+        return {self.spec.URI: MCPServiceExtensionMetadata(mcp_fulfillments=mcp_fulfillments).model_dump(mode="json")}