platform extension wip

Signed-off-by: Radek Ježek <radek.jezek@ibm.com>

Author: jezekra1

agents/community/gpt-researcher/gpt_researcher_agent/agent.py

@@ -12,7 +12,7 @@
 from beeai_sdk.a2a.extensions import TrajectoryExtensionServer, TrajectoryExtensionSpec, AgentDetail
 from beeai_sdk.a2a.types import RunYield
 from beeai_sdk.server import Server
-from beeai_sdk.server.context import Context
+from beeai_sdk.server.context import RunContext
 
 server = Server()
 
@@ -59,7 +59,7 @@
     ],
 )
 async def gpt_researcher(
-    message: Message, context: Context, trajectory: Annotated[TrajectoryExtensionServer, TrajectoryExtensionSpec()]
+    message: Message, context: RunContext, trajectory: Annotated[TrajectoryExtensionServer, TrajectoryExtensionSpec()]
 ) -> AsyncGenerator[RunYield, None]:
     """
     The agent conducts in-depth local and web research using a language model to generate comprehensive reports with

agents/official/beeai-framework/chat/src/chat/agent.py

@@ -41,7 +41,7 @@
 )
 from beeai_sdk.a2a.types import AgentMessage
 from beeai_sdk.server import Server
-from beeai_sdk.server.context import Context
+from beeai_sdk.server.context import RunContext
 from chat.helpers.citations import extract_citations
 from chat.helpers.trajectory import TrajectoryContent
 from openinference.instrumentation.beeai import BeeAIInstrumentor
@@ -136,7 +136,7 @@
 )
 async def chat(
     message: Message,
-    context: Context,
+    context: RunContext,
     trajectory: Annotated[TrajectoryExtensionServer, TrajectoryExtensionSpec()],
     citation: Annotated[CitationExtensionServer, CitationExtensionSpec()],
 ):

agents/official/beeai-framework/chat/uv.lock

@@ -10,22 +10,30 @@
 from beeai_sdk.a2a.extensions import LLMServiceExtensionServer, LLMServiceExtensionSpec
 from beeai_sdk.a2a.extensions.ui.trajectory import TrajectoryExtensionServer, TrajectoryExtensionSpec
 from beeai_sdk.a2a.types import RunYield
+from beeai_sdk.platform import File
 from beeai_sdk.server import Server
-from beeai_sdk.server.context import Context
+from beeai_sdk.server.context import RunContext
+from src.beeai_sdk.a2a.extensions.services.platform import (
+    PlatformApiExtensionServer,
+    PlatformApiExtensionSpec,
+)
 
 server = Server()
 
 
 @server.agent()
 async def dependent_agent(
     message: Message,
-    context: Context,
+    context: RunContext,
     trajectory: Annotated[TrajectoryExtensionServer, TrajectoryExtensionSpec()],
     # does not typecheck, does not ruff check
     llm: Annotated[LLMServiceExtensionServer, LLMServiceExtensionSpec.single_demand()],
+    _: Annotated[PlatformApiExtensionServer, PlatformApiExtensionSpec()],
 ) -> AsyncGenerator[RunYield, Message]:
     """Awaits a user message"""
 
+    await File.create(filename="my_file.txt", content=b"hello world", content_type="text/plain")
+
     yield trajectory.trajectory_metadata(title="context_param", content=str(context))
     yield trajectory.trajectory_metadata(title="message_param", content=str(message.model_dump()))
     yield trajectory.message(trajectory_title="llm_param", trajectory_content=str(llm.data))

apps/beeai-sdk/examples/dependencies.py

@@ -10,15 +10,15 @@
 from beeai_sdk.a2a.extensions.services.mcp import MCPServiceExtensionServer, MCPServiceExtensionSpec
 from beeai_sdk.a2a.types import RunYield
 from beeai_sdk.server import Server
-from beeai_sdk.server.context import Context
+from beeai_sdk.server.context import RunContext
 
 server = Server()
 
 
 @server.agent()
 async def mcp_agent(
     message: Message,
-    context: Context,
+    context: RunContext,
     mcp: Annotated[
         MCPServiceExtensionServer,
         MCPServiceExtensionSpec.single_demand(),

apps/beeai-sdk/examples/mcp_agent.py

@@ -20,7 +20,7 @@ dependencies = [
     "tenacity>=9.1.2",
     "janus>=2.0.0",
     "uvloop>=0.21.0",
-    "httpx",                                          # version determined by a2a-sdk
+    "httpx", # version determined by a2a-sdk
     "mcp>=1.12.3",
 ]
 
@@ -42,18 +42,18 @@ build-backend = "uv_build"
 line-length = 120
 target-version = "py311"
 lint.select = [
-    "E",     # pycodestyle errors
-    "W",     # pycodestyle warnings
-    "F",     # pyflakes
-    "UP",    # pyupgrade
-    "I",     # isort
-    "B",     # bugbear
-    "N",     # pep8-naming
-    "C4",    # Comprehensions
-    "Q",     # Quotes
-    "SIM",   # Simplify
-    "RUF",   # Ruff
-    "TID",   # tidy-imports
+    "E", # pycodestyle errors
+    "W", # pycodestyle warnings
+    "F", # pyflakes
+    "UP", # pyupgrade
+    "I", # isort
+    "B", # bugbear
+    "N", # pep8-naming
+    "C4", # Comprehensions
+    "Q", # Quotes
+    "SIM", # Simplify
+    "RUF", # Ruff
+    "TID", # tidy-imports
     "ASYNC", # async
     # TODO: add "DTZ",   # DatetimeZ
     # TODO: add "ANN",   # annotations
@@ -71,3 +71,5 @@ addopts = "-v"
 
 [tool.pyright]
 ignore = ["tests/**"]
+venvPath = "."
+venv = ".venv"

apps/beeai-sdk/pyproject.toml

@@ -16,7 +16,7 @@
 
 
 if typing.TYPE_CHECKING:
-    from beeai_sdk.server.context import Context
+    from beeai_sdk.server.context import RunContext
 
 
 def _get_generic_args(cls: type, base_class: type) -> tuple[typing.Any, ...]:
@@ -141,14 +141,18 @@ def parse_client_metadata(self, message: a2a.types.Message) -> MetadataFromClien
             else pydantic.TypeAdapter(self.MetadataFromClient).validate_python(message.metadata[self.spec.URI])
         )
 
-    def handle_incoming_message(self, message: a2a.types.Message, context: Context):
+    def handle_incoming_message(self, message: a2a.types.Message, context: RunContext):
         if self._metadata_from_client is None:
             self._metadata_from_client = self.parse_client_metadata(message)
 
-    def __call__(self, message: a2a.types.Message, context: Context) -> typing.Self:
+    def __call__(self, message: a2a.types.Message, context: RunContext) -> typing.Self:
         self.handle_incoming_message(message, context)
         return self
 
+    async def initialize(self):
+        """Called when entering the agent context after the first message was parsed (__call__ was already called)"""
+        pass
+
 
 class BaseExtensionClient(abc.ABC, typing.Generic[ExtensionSpecT, MetadataFromServerT]):
     MetadataFromServer: type[MetadataFromServerT]

apps/beeai-sdk/src/beeai_sdk/a2a/extensions/base.py

@@ -0,0 +1,11 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from beeai_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}")

apps/beeai-sdk/src/beeai_sdk/a2a/extensions/exceptions.py

@@ -0,0 +1,117 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import asyncio
+import os
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from types import NoneType
+
+import a2a.types
+import pydantic
+from pydantic.networks import HttpUrl
+
+from beeai_sdk.a2a.extensions.base import (
+    BaseExtensionClient,
+    BaseExtensionServer,
+    BaseExtensionSpec,
+)
+from beeai_sdk.a2a.extensions.exceptions import ExtensionError
+from beeai_sdk.platform import use_platform_client
+from beeai_sdk.platform.client import PlatformClient
+
+
+class PlatformApiExtensionMetadata(pydantic.BaseModel):
+    base_url: HttpUrl | None = None
+    auth_token: pydantic.Secret[str]
+    expires_at: pydantic.AwareDatetime | None = None
+
+
+class PlatformApiExtension(pydantic.BaseModel):
+    """
+    Request authentication token and url to be able to access the beeai platform API
+    """
+
+
+class PlatformApiExtensionParams(pydantic.BaseModel):
+    auto_use: bool = True
+
+
+class PlatformApiExtensionSpec(BaseExtensionSpec[PlatformApiExtensionParams]):
+    URI: str = "https://a2a-extensions.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
+    _should_exit: asyncio.Event = asyncio.Event()
+    _wait_for_startup: asyncio.Event = asyncio.Event()
+    _lifecycle_task: asyncio.Task | None = None
+    _lifecycle_task_exception: Exception | None = None
+    _initialized: bool = False
+
+    def parse_client_metadata(self, message: a2a.types.Message) -> PlatformApiExtensionMetadata | None:
+        self.context_id = message.context_id
+        # we assume that the 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
+        return super().parse_client_metadata(message)
+
+    async def managed_platform_client_lifespan(self):
+        try:
+            async with self.use_client():
+                self._wait_for_startup.set()
+                await self._should_exit.wait()
+        except Exception as e:
+            self._lifecycle_task_exception = e
+            raise
+        finally:
+            self._wait_for_startup.set()
+
+    async def initialize(self):
+        """Called when entering the agent context after the first message was parsed (__call__ was already called)"""
+        if self._initialized:
+            raise ExtensionError(self.spec, "Platform extension was already initialized")
+
+        if self.spec.params.auto_use:
+            self._lifecycle_task = asyncio.create_task(self.managed_platform_client_lifespan())
+            self._lifecycle_task.add_done_callback(lambda *_args: None)
+            await self._wait_for_startup.wait()
+            if self._lifecycle_task_exception:
+                raise self._lifecycle_task_exception
+        self._initialized = True
+
+    def __del__(self):
+        self._should_exit.set()
+
+    @asynccontextmanager
+    async def use_client(self) -> AsyncIterator[PlatformClient]:
+        if not self.data:
+            raise ExtensionError(self.spec, "Platform extension metadata was not provided")
+        base_url = str(self.data.base_url or os.getenv("PLATFORM_URL", "http://127.0.0.1:8333"))
+        auth_token = self.data.auth_token.get_secret_value()
+        async with use_platform_client(context_id=self.context_id, base_url=base_url, auth_token=auth_token) as client:
+            yield client
+
+
+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(),
+            }
+        }