agentstack-sdk 0.5.2rc2__py3-none-any.whl

agentstack_sdk/a2a/extensions/services/platform.py

@@ -0,0 +1,141 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from types import NoneType
+from typing import TYPE_CHECKING
+
+import pydantic
+from a2a.server.agent_execution.context import RequestContext
+from a2a.types import Message as A2AMessage
+from fastapi.security.utils import get_authorization_scheme_param
+from pydantic.networks import HttpUrl
+from typing_extensions import override
+
+from agentstack_sdk.a2a.extensions.base import (
+    BaseExtensionClient,
+    BaseExtensionServer,
+    BaseExtensionSpec,
+)
+from agentstack_sdk.a2a.extensions.exceptions import ExtensionError
+from agentstack_sdk.platform import use_platform_client
+from agentstack_sdk.platform.client import PlatformClient
+from agentstack_sdk.server.middleware.platform_auth_backend import PlatformAuthenticatedUser
+from agentstack_sdk.util.httpx import BearerAuth
+
+if TYPE_CHECKING:
+    from agentstack_sdk.server.context import RunContext
+
+
+class PlatformApiExtensionMetadata(pydantic.BaseModel):
+    base_url: HttpUrl | None = None
+    auth_token: pydantic.Secret[str] | None = None
+    expires_at: pydantic.AwareDatetime | None = None
+
+
+class PlatformApiExtension(pydantic.BaseModel):
+    """
+    Request authentication token and url to be able to access the agentstack API
+    """
+
+
+class PlatformApiExtensionParams(pydantic.BaseModel):
+    auto_use: bool = True
+
+
+class PlatformApiExtensionSpec(BaseExtensionSpec[PlatformApiExtensionParams]):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/services/platform_api/v1"
+
+    def __init__(self, params: PlatformApiExtensionParams | None = None) -> None:
+        super().__init__(params or PlatformApiExtensionParams())
+
+
+class PlatformApiExtensionServer(BaseExtensionServer[PlatformApiExtensionSpec, PlatformApiExtensionMetadata]):
+    context_id: str | None = None
+
+    @asynccontextmanager
+    @override
+    async def lifespan(self) -> AsyncIterator[None]:
+        """Called when entering the agent context after the first message was parsed (__call__ was already called)"""
+        if self.data and self.spec.params.auto_use:
+            async with self.use_client():
+                yield
+        else:
+            yield
+
+    def _get_header_token(self, request_context: RequestContext) -> pydantic.Secret[str] | None:
+        header_token = None
+        call_context = request_context.call_context
+        assert call_context
+        if isinstance(call_context.user, PlatformAuthenticatedUser):
+            header_token = call_context.user.auth_token.get_secret_value()
+        elif auth_header := call_context.state.get("headers", {}).get("authorization", None):
+            _scheme, header_token = get_authorization_scheme_param(auth_header)
+        return pydantic.Secret(header_token) if header_token else None
+
+    @override
+    def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, request_context: RequestContext):
+        super().handle_incoming_message(message, run_context, request_context)
+        # we assume that request context id is the same ID as the platform context id
+        # if different IDs are passed, api requests to platform using this token will fail
+        self.context_id = request_context.context_id
+
+        self._metadata_from_client = self._metadata_from_client or PlatformApiExtensionMetadata()
+        data = self._metadata_from_client
+        data.base_url = data.base_url or HttpUrl(os.getenv("PLATFORM_URL", "http://127.0.0.1:8333"))
+        data.auth_token = data.auth_token or self._get_header_token(request_context)
+
+        if not data.auth_token:
+            raise ExtensionError(self.spec, "Platform extension metadata was not provided")
+
+    @asynccontextmanager
+    async def use_client(self) -> AsyncIterator[PlatformClient]:
+        if not self.data or not self.data.auth_token:
+            raise ExtensionError(self.spec, "Platform extension metadata was not provided")
+        async with use_platform_client(
+            context_id=self.context_id,
+            base_url=str(self.data.base_url),
+            auth_token=self.data.auth_token.get_secret_value(),
+        ) as client:
+            yield client
+
+    async def create_httpx_auth(self) -> BearerAuth:
+        if not self.data or not self.data.auth_token:
+            raise ExtensionError(self.spec, "Platform extension metadata was not provided")
+        return BearerAuth(token=self.data.auth_token.get_secret_value())
+
+
+class PlatformApiExtensionClient(BaseExtensionClient[PlatformApiExtensionSpec, NoneType]):
+    def api_auth_metadata(
+        self,
+        *,
+        auth_token: pydantic.Secret[str] | str,
+        expires_at: pydantic.AwareDatetime | None = None,
+        base_url: HttpUrl | None = None,
+    ) -> dict[str, dict[str, str]]:
+        return {
+            self.spec.URI: {
+                **PlatformApiExtensionMetadata(
+                    base_url=base_url,
+                    auth_token=pydantic.Secret("replaced below"),
+                    expires_at=expires_at,
+                ).model_dump(mode="json"),
+                "auth_token": auth_token if isinstance(auth_token, str) else auth_token.get_secret_value(),
+            }
+        }
+
+
+class _PlatformSelfRegistrationExtension(pydantic.BaseModel):
+    """Internal extension"""
+
+
+class _PlatformSelfRegistrationExtensionParams(pydantic.BaseModel):
+    self_registration_id: str
+
+
+class _PlatformSelfRegistrationExtensionSpec(BaseExtensionSpec[_PlatformSelfRegistrationExtensionParams]):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/services/platform-self-registration/v1"

agentstack_sdk/a2a/extensions/tools/__init__.py

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

agentstack_sdk/a2a/extensions/tools/call.py

@@ -0,0 +1,114 @@
+# 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, Any, Literal
+
+import a2a.types
+from mcp import Tool
+from mcp.types import Implementation
+from pydantic import BaseModel, Field
+
+from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
+from agentstack_sdk.a2a.extensions.tools.exceptions import ToolCallRejectionError
+from agentstack_sdk.a2a.types import AgentMessage, InputRequired
+
+if TYPE_CHECKING:
+    from agentstack_sdk.server.context import RunContext
+
+
+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 ToolCallRequest(BaseModel):
+    name: str = Field(description="The programmatic name of the tool.")
+    title: str | None = Field(None, description="A human-readable title for the tool.")
+    description: str | None = Field(None, description="A human-readable description 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
+    ) -> ToolCallRequest:
+        return ToolCallRequest(
+            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,
+        )
+
+
+class ToolCallResponse(BaseModel):
+    action: Literal["accept", "reject"]
+
+
+class ToolCallExtensionParams(BaseModel):
+    pass
+
+
+class ToolCallExtensionSpec(BaseExtensionSpec[ToolCallExtensionParams]):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/tools/call/v1"
+
+
+class ToolCallExtensionMetadata(BaseModel):
+    pass
+
+
+class ToolCallExtensionServer(BaseExtensionServer[ToolCallExtensionSpec, ToolCallExtensionMetadata]):
+    def create_request_message(self, *, request: ToolCallRequest):
+        return AgentMessage(
+            text="Tool call approval requested", metadata={self.spec.URI: request.model_dump(mode="json")}
+        )
+
+    def parse_response(self, *, message: a2a.types.Message):
+        if not message or not message.metadata or not (data := message.metadata.get(self.spec.URI)):
+            raise RuntimeError("Invalid mcp response")
+        return ToolCallResponse.model_validate(data)
+
+    async def request_tool_call_approval(
+        self,
+        request: ToolCallRequest,
+        *,
+        context: RunContext,
+    ) -> ToolCallResponse:
+        message = self.create_request_message(request=request)
+        message = await context.yield_async(InputRequired(message=message))
+        if message:
+            result = self.parse_response(message=message)
+            match result.action:
+                case "accept":
+                    return result
+                case "reject":
+                    raise ToolCallRejectionError("User has rejected the tool call")
+
+        else:
+            raise RuntimeError("Yield did not return a message")
+
+
+class ToolCallExtensionClient(BaseExtensionClient[ToolCallExtensionSpec, NoneType]):
+    def create_response_message(self, *, response: ToolCallResponse, 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 or not message.metadata or not (data := message.metadata.get(self.spec.URI)):
+            raise ValueError("Invalid tool call request")
+        return ToolCallRequest.model_validate(data)
+
+    def metadata(self) -> dict[str, Any]:
+        return {self.spec.URI: ToolCallExtensionMetadata().model_dump(mode="json")}

agentstack_sdk/a2a/extensions/tools/exceptions.py

@@ -0,0 +1,6 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+
+class ToolCallRejectionError(RuntimeError):
+    pass

agentstack_sdk/a2a/extensions/ui/__init__.py

@@ -0,0 +1,10 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from .agent_detail import *
+from .canvas import *
+from .citation import *
+from .error import *
+from .form_request import *
+from .settings import *
+from .trajectory import *

agentstack_sdk/a2a/extensions/ui/agent_detail.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 types import NoneType
+
+import pydantic
+
+from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
+
+
+class AgentDetailTool(pydantic.BaseModel):
+    name: str
+    description: str
+
+
+class AgentDetailContributor(pydantic.BaseModel):
+    name: str
+    email: str | None = None
+    url: str | None = None
+
+
+class EnvVar(pydantic.BaseModel):
+    name: str
+    description: str | None = None
+    required: bool = False
+
+
+class AgentDetail(pydantic.BaseModel, extra="allow"):
+    interaction_mode: str | None = pydantic.Field("multi-turn", examples=["multi-turn", "single-turn"])
+    user_greeting: str | None = None
+    input_placeholder: str | None = None
+    tools: list[AgentDetailTool] | None = None
+    framework: str | None = None
+    license: str | None = None
+    programming_language: str | None = "Python"
+    homepage_url: str | None = None
+    source_code_url: str | None = None
+    container_image_url: str | None = None
+    author: AgentDetailContributor | None = None
+    contributors: list[AgentDetailContributor] | None = None
+    variables: list[EnvVar] | None = None
+
+
+class AgentDetailExtensionSpec(BaseExtensionSpec[AgentDetail]):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/ui/agent-detail/v1"
+
+
+class AgentDetailExtensionServer(BaseExtensionServer[AgentDetailExtensionSpec, NoneType]): ...
+
+
+class AgentDetailExtensionClient(BaseExtensionClient[AgentDetailExtensionSpec, AgentDetail]): ...

agentstack_sdk/a2a/extensions/ui/canvas.py

@@ -0,0 +1,71 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pydantic
+from a2a.server.agent_execution.context import RequestContext
+from a2a.types import Artifact, TextPart
+from a2a.types import Message as A2AMessage
+from typing_extensions import override
+
+if TYPE_CHECKING:
+    from agentstack_sdk.server.context import RunContext
+
+from agentstack_sdk.a2a.extensions.base import (
+    BaseExtensionServer,
+    NoParamsBaseExtensionSpec,
+)
+
+
+class CanvasEditRequestMetadata(pydantic.BaseModel):
+    start_index: int
+    end_index: int
+    description: str
+    artifact_id: str
+
+
+class CanvasEditRequest(pydantic.BaseModel):
+    start_index: int
+    end_index: int
+    description: str
+    artifact: Artifact
+
+
+class CanvasExtensionSpec(NoParamsBaseExtensionSpec):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/ui/canvas/v1"
+
+
+class CanvasExtensionServer(BaseExtensionServer[CanvasExtensionSpec, CanvasEditRequestMetadata]):
+    @override
+    def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, request_context: RequestContext):
+        if message.metadata and self.spec.URI in message.metadata and message.parts:
+            message.parts = [part for part in message.parts if not isinstance(part.root, TextPart)]
+
+        super().handle_incoming_message(message, run_context, request_context)
+        self.context = run_context
+
+    async def parse_canvas_edit_request(self, *, message: A2AMessage) -> CanvasEditRequest | None:
+        if not message or not message.metadata or not (data := message.metadata.get(self.spec.URI)):
+            return None
+
+        metadata = CanvasEditRequestMetadata.model_validate(data)
+
+        try:
+            artifact = await anext(
+                artifact
+                async for artifact in self.context.load_history()
+                if isinstance(artifact, Artifact) and artifact.parts
+                if artifact.artifact_id == metadata.artifact_id
+            )
+        except StopAsyncIteration as e:
+            raise ValueError(f"Artifact {metadata.artifact_id} not found in history") from e
+
+        return CanvasEditRequest(
+            start_index=metadata.start_index,
+            end_index=metadata.end_index,
+            description=metadata.description,
+            artifact=artifact,
+        )

agentstack_sdk/a2a/extensions/ui/citation.py

@@ -0,0 +1,78 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from types import NoneType
+
+import pydantic
+from a2a.types import DataPart, FilePart, Part, TextPart
+
+from agentstack_sdk.a2a.extensions.base import (
+    BaseExtensionClient,
+    BaseExtensionServer,
+    NoParamsBaseExtensionSpec,
+)
+from agentstack_sdk.a2a.types import AgentMessage, Metadata
+
+
+class Citation(pydantic.BaseModel):
+    """
+    Represents an inline citation, providing info about information source. This
+    is supposed to be rendered as an inline icon, optionally marking a text
+    range it belongs to.
+
+    If Citation is included together with content in the message part,
+    the citation belongs to that content and renders at the Part position.
+    This way may be used for non-text content, like images and files.
+
+    Alternatively, `start_index` and `end_index` may define a text range,
+    counting characters in the current Message across all Parts containing plain
+    text, where the citation will be rendered. If one of `start_index` and
+    `end_index` is missing or their values are equal, the citation renders only
+    as an inline icon at that position.
+
+    If both `start_index` and `end_index` are not present and Part has empty
+    content, the citation renders as inline icon only at the Part position.
+
+    Properties:
+    - url: URL of the source document.
+    - title: Title of the source document.
+    - description: Accompanying text, which may be a general description of the
+                   source document, or a specific snippet.
+    """
+
+    start_index: int | None = None
+    end_index: int | None = None
+    url: str | None = None
+    title: str | None = None
+    description: str | None = None
+
+
+class CitationMetadata(pydantic.BaseModel):
+    citations: list[Citation] = pydantic.Field(default_factory=list)
+
+
+class CitationExtensionSpec(NoParamsBaseExtensionSpec):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/ui/citation/v1"
+
+
+class CitationExtensionServer(BaseExtensionServer[CitationExtensionSpec, NoneType]):
+    def citation_metadata(self, *, citations: list[Citation]) -> Metadata:
+        return Metadata({self.spec.URI: CitationMetadata(citations=citations).model_dump(mode="json")})
+
+    def message(
+        self,
+        text: str | None = None,
+        parts: list[Part | TextPart | FilePart | DataPart] | None = None,
+        *,
+        citations: list[Citation],
+    ) -> AgentMessage:
+        return AgentMessage(
+            text=text,
+            parts=parts or [],
+            metadata=self.citation_metadata(citations=citations),
+        )
+
+
+class CitationExtensionClient(BaseExtensionClient[CitationExtensionSpec, CitationMetadata]): ...