agentstack-sdk 0.4.1rc1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (60) hide show
  1. agentstack_sdk-0.4.1rc1/PKG-INFO +69 -0
  2. agentstack_sdk-0.4.1rc1/README.md +45 -0
  3. agentstack_sdk-0.4.1rc1/pyproject.toml +77 -0
  4. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/__init__.py +6 -0
  5. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/__init__.py +2 -0
  6. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/__init__.py +6 -0
  7. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/__init__.py +5 -0
  8. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/oauth/__init__.py +4 -0
  9. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/oauth/oauth.py +144 -0
  10. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/oauth/storage/__init__.py +5 -0
  11. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/oauth/storage/base.py +11 -0
  12. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/oauth/storage/memory.py +38 -0
  13. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/secrets/__init__.py +4 -0
  14. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/secrets/secrets.py +71 -0
  15. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/base.py +196 -0
  16. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/exceptions.py +11 -0
  17. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/services/__init__.py +7 -0
  18. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/services/embedding.py +99 -0
  19. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/services/llm.py +98 -0
  20. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/services/mcp.py +149 -0
  21. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/services/platform.py +117 -0
  22. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/ui/__init__.py +8 -0
  23. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/ui/agent_detail.py +54 -0
  24. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/ui/canvas.py +65 -0
  25. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/ui/citation.py +83 -0
  26. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/ui/form.py +193 -0
  27. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/ui/settings.py +73 -0
  28. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/ui/trajectory.py +66 -0
  29. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/types.py +106 -0
  30. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/__init__.py +10 -0
  31. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/client.py +119 -0
  32. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/common.py +37 -0
  33. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/configuration.py +47 -0
  34. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/context.py +263 -0
  35. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/file.py +213 -0
  36. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/model_provider.py +131 -0
  37. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/provider.py +219 -0
  38. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/provider_build.py +190 -0
  39. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/types.py +45 -0
  40. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/variables.py +44 -0
  41. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/vector_store.py +205 -0
  42. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/py.typed +0 -0
  43. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/__init__.py +4 -0
  44. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/agent.py +555 -0
  45. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/app.py +69 -0
  46. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/constants.py +8 -0
  47. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/context.py +53 -0
  48. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/dependencies.py +118 -0
  49. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/logging.py +63 -0
  50. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/server.py +360 -0
  51. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/store/__init__.py +3 -0
  52. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/store/context_store.py +28 -0
  53. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/store/memory_context_store.py +43 -0
  54. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/store/platform_context_store.py +48 -0
  55. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/telemetry.py +53 -0
  56. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/utils.py +26 -0
  57. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/util/__init__.py +4 -0
  58. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/util/file.py +260 -0
  59. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/util/resource_context.py +44 -0
  60. agentstack_sdk-0.4.1rc1/src/agentstack_sdk/util/utils.py +38 -0
@@ -0,0 +1,69 @@
1
+ Metadata-Version: 2.3
2
+ Name: agentstack-sdk
3
+ Version: 0.4.1rc1
4
+ Summary: Agent Stack SDK
5
+ Author: IBM Corp.
6
+ Requires-Dist: a2a-sdk==0.3.9
7
+ Requires-Dist: objprint>=0.3.0
8
+ Requires-Dist: uvicorn>=0.35.0
9
+ Requires-Dist: asyncclick>=8.1.8
10
+ Requires-Dist: sse-starlette>=2.2.1
11
+ Requires-Dist: starlette>=0.47.2
12
+ Requires-Dist: anyio>=4.9.0
13
+ Requires-Dist: opentelemetry-api>=1.35.0
14
+ Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.35.0
15
+ Requires-Dist: opentelemetry-instrumentation-fastapi>=0.56b0
16
+ Requires-Dist: opentelemetry-sdk>=1.35.0
17
+ Requires-Dist: tenacity>=9.1.2
18
+ Requires-Dist: janus>=2.0.0
19
+ Requires-Dist: httpx
20
+ Requires-Dist: mcp>=1.12.3
21
+ Requires-Dist: fastapi>=0.116.1
22
+ Requires-Python: >=3.11, <3.14
23
+ Description-Content-Type: text/markdown
24
+
25
+ # Agent Stack SDK
26
+
27
+ ## Examples
28
+
29
+ The examples connect to the Agent Stack for LLM inteference.
30
+
31
+ Run using:
32
+
33
+ ```bash
34
+ uv run examples/agent.py
35
+ ```
36
+
37
+ Connect to the agent using the CLI:
38
+
39
+ ```bash
40
+ uv run examples/cli.py
41
+ ```
42
+
43
+ ## Plan
44
+
45
+ - `agentstack_sdk`
46
+ - `a2a`:
47
+ - `extensions`: Shared definitions for A2A extensions
48
+ - `services`: Dependency injection extensions for external services
49
+ - `llm`
50
+ - `embedding`
51
+ - `docling`
52
+ - `file_store`
53
+ - `vector_store`
54
+ - `ui`: User interface extensions for Agent Stack UI
55
+ - `trajectory`
56
+ - `citations`
57
+ - `history`: store and allow requesting the full history of the context
58
+ - `server`
59
+ - `context_storage`: store data associated with context_id
60
+ - `wrapper`: conveniently build A2A agents -- opinionated on how tasks work, `yield`-semantics, autowired
61
+ services
62
+ - `services`: clients for external services
63
+ - `llm`: OpenAI-compatible chat LLM
64
+ - `embedding`: OpenAI-compatible embedding
65
+ - `text_extraction`: Docling-compatible text extraction
66
+ - `file_store`: S3-compatible file storage
67
+ - `vector_store`: some vector store?
68
+ - `client`
69
+ - ?
@@ -0,0 +1,45 @@
1
+ # Agent Stack SDK
2
+
3
+ ## Examples
4
+
5
+ The examples connect to the Agent Stack for LLM inteference.
6
+
7
+ Run using:
8
+
9
+ ```bash
10
+ uv run examples/agent.py
11
+ ```
12
+
13
+ Connect to the agent using the CLI:
14
+
15
+ ```bash
16
+ uv run examples/cli.py
17
+ ```
18
+
19
+ ## Plan
20
+
21
+ - `agentstack_sdk`
22
+ - `a2a`:
23
+ - `extensions`: Shared definitions for A2A extensions
24
+ - `services`: Dependency injection extensions for external services
25
+ - `llm`
26
+ - `embedding`
27
+ - `docling`
28
+ - `file_store`
29
+ - `vector_store`
30
+ - `ui`: User interface extensions for Agent Stack UI
31
+ - `trajectory`
32
+ - `citations`
33
+ - `history`: store and allow requesting the full history of the context
34
+ - `server`
35
+ - `context_storage`: store data associated with context_id
36
+ - `wrapper`: conveniently build A2A agents -- opinionated on how tasks work, `yield`-semantics, autowired
37
+ services
38
+ - `services`: clients for external services
39
+ - `llm`: OpenAI-compatible chat LLM
40
+ - `embedding`: OpenAI-compatible embedding
41
+ - `text_extraction`: Docling-compatible text extraction
42
+ - `file_store`: S3-compatible file storage
43
+ - `vector_store`: some vector store?
44
+ - `client`
45
+ - ?
@@ -0,0 +1,77 @@
1
+ [project]
2
+ name = "agentstack-sdk"
3
+ version = "0.4.1-rc1"
4
+ description = "Agent Stack SDK"
5
+ readme = "README.md"
6
+ authors = [{ name = "IBM Corp." }]
7
+ requires-python = ">=3.11,<3.14"
8
+ dependencies = [
9
+ "a2a-sdk==0.3.9",
10
+ "objprint>=0.3.0",
11
+ "uvicorn>=0.35.0",
12
+ "asyncclick>=8.1.8",
13
+ "sse-starlette>=2.2.1",
14
+ "starlette>=0.47.2",
15
+ "anyio>=4.9.0",
16
+ "opentelemetry-api>=1.35.0",
17
+ "opentelemetry-exporter-otlp-proto-http>=1.35.0",
18
+ "opentelemetry-instrumentation-fastapi>=0.56b0",
19
+ "opentelemetry-sdk>=1.35.0",
20
+ "tenacity>=9.1.2",
21
+ "janus>=2.0.0",
22
+ "httpx", # version determined by a2a-sdk
23
+ "mcp>=1.12.3",
24
+ "fastapi>=0.116.1",
25
+ ]
26
+
27
+ [dependency-groups]
28
+ dev = [
29
+ "beeai-framework[duckduckgo,wikipedia]>=0.1.58",
30
+ "pyright>=1.1.403",
31
+ "pytest>=8.4.1",
32
+ "pytest-asyncio>=1.1.0",
33
+ "pytest-httpx>=0.35.0",
34
+ "ruff>=0.12.3",
35
+ ]
36
+
37
+ [build-system]
38
+ requires = ["uv_build>=0.9.0,<0.10.0"]
39
+ build-backend = "uv_build"
40
+
41
+ [tool.ruff]
42
+ line-length = 120
43
+ target-version = "py311"
44
+ lint.select = [
45
+ "E", # pycodestyle errors
46
+ "W", # pycodestyle warnings
47
+ "F", # pyflakes
48
+ "UP", # pyupgrade
49
+ "I", # isort
50
+ "B", # bugbear
51
+ "N", # pep8-naming
52
+ "C4", # Comprehensions
53
+ "Q", # Quotes
54
+ "SIM", # Simplify
55
+ "RUF", # Ruff
56
+ "TID", # tidy-imports
57
+ "ASYNC", # async
58
+ # TODO: add "DTZ", # DatetimeZ
59
+ # TODO: add "ANN", # annotations
60
+ ]
61
+ lint.ignore = [
62
+ "E501", # line lenght (annyoing)
63
+ "N999", # invalid module name agentstack-server (yeah, we use a dash, deal with it)
64
+ "F403", # redundant with Pyright
65
+ ]
66
+ force-exclude = true
67
+
68
+ [tool.pytest.ini_options]
69
+ markers = ["e2e", "unit"]
70
+ asyncio_mode = "auto"
71
+ asyncio_default_fixture_loop_scope = "function"
72
+ addopts = "-v"
73
+
74
+ [tool.pyright]
75
+ ignore = ["tests/**", "examples/cli.py"]
76
+ venvPath = "."
77
+ venv = ".venv"
@@ -0,0 +1,6 @@
1
+ # Copyright 2025 © BeeAI a Series of LF Projects, LLC
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ from importlib.metadata import version
5
+
6
+ __version__ = version("agentstack-sdk")
@@ -0,0 +1,2 @@
1
+ # Copyright 2025 © BeeAI a Series of LF Projects, LLC
2
+ # SPDX-License-Identifier: Apache-2.0
@@ -0,0 +1,6 @@
1
+ # Copyright 2025 © BeeAI a Series of LF Projects, LLC
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ from .auth import *
5
+ from .services import *
6
+ from .ui import *
@@ -0,0 +1,5 @@
1
+ # Copyright 2025 © BeeAI a Series of LF Projects, LLC
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ from .oauth import *
5
+ from .secrets import *
@@ -0,0 +1,4 @@
1
+ # Copyright 2025 © BeeAI a Series of LF Projects, LLC
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ from .oauth import *
@@ -0,0 +1,144 @@
1
+ # Copyright 2025 © BeeAI a Series of LF Projects, LLC
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ from __future__ import annotations
5
+
6
+ import uuid
7
+ from types import NoneType
8
+ from typing import TYPE_CHECKING, Any, Self
9
+ from urllib.parse import parse_qs
10
+
11
+ import a2a.types
12
+ import pydantic
13
+ from mcp.client.auth import OAuthClientProvider
14
+ from mcp.shared.auth import OAuthClientMetadata
15
+
16
+ from agentstack_sdk.a2a.extensions.auth.oauth.storage import MemoryTokenStorageFactory, TokenStorageFactory
17
+ from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
18
+ from agentstack_sdk.a2a.types import AgentMessage, AuthRequired, RunYieldResume
19
+
20
+ if TYPE_CHECKING:
21
+ from agentstack_sdk.server.context import RunContext
22
+
23
+ _DEFAULT_DEMAND_NAME = "default"
24
+
25
+
26
+ class AuthRequest(pydantic.BaseModel):
27
+ authorization_endpoint_url: pydantic.AnyUrl
28
+
29
+
30
+ class AuthResponse(pydantic.BaseModel):
31
+ redirect_uri: pydantic.AnyUrl
32
+
33
+
34
+ class OAuthFulfillment(pydantic.BaseModel):
35
+ redirect_uri: pydantic.AnyUrl
36
+
37
+
38
+ class OAuthDemand(pydantic.BaseModel):
39
+ redirect_uri: bool = True
40
+
41
+
42
+ class OAuthExtensionParams(pydantic.BaseModel):
43
+ oauth_demands: dict[str, OAuthDemand]
44
+ """Server requests that the agent requires to be provided by the client."""
45
+
46
+
47
+ class OAuthExtensionSpec(BaseExtensionSpec[OAuthExtensionParams]):
48
+ URI: str = "https://a2a-extensions.agentstack.beeai.dev/auth/oauth/v1"
49
+
50
+ @classmethod
51
+ def single_demand(cls, name: str = _DEFAULT_DEMAND_NAME) -> Self:
52
+ return cls(params=OAuthExtensionParams(oauth_demands={name: OAuthDemand()}))
53
+
54
+
55
+ class OAuthExtensionMetadata(pydantic.BaseModel):
56
+ oauth_fulfillments: dict[str, OAuthFulfillment] = {}
57
+ """Provided servers corresponding to the server requests."""
58
+
59
+
60
+ class OAuthExtensionServer(BaseExtensionServer[OAuthExtensionSpec, OAuthExtensionMetadata]):
61
+ def __init__(self, spec: OAuthExtensionSpec, token_storage_factory: TokenStorageFactory | None = None) -> None:
62
+ super().__init__(spec)
63
+ self.token_storage_factory = token_storage_factory or MemoryTokenStorageFactory()
64
+
65
+ def handle_incoming_message(self, message: a2a.types.Message, context: RunContext):
66
+ super().handle_incoming_message(message, context)
67
+ self.context = context
68
+
69
+ def _get_fulfillment_for_resource(self, resource_url: pydantic.AnyUrl):
70
+ if not self.data:
71
+ raise RuntimeError("No fulfillments found")
72
+
73
+ fulfillment = self.data.oauth_fulfillments.get(str(resource_url)) or self.data.oauth_fulfillments.get(
74
+ _DEFAULT_DEMAND_NAME
75
+ )
76
+ if fulfillment:
77
+ return fulfillment
78
+
79
+ raise RuntimeError("Fulfillment not found")
80
+
81
+ async def create_httpx_auth(self, *, resource_url: pydantic.AnyUrl):
82
+ fulfillment = self._get_fulfillment_for_resource(resource_url=resource_url)
83
+
84
+ resume: RunYieldResume = None
85
+
86
+ async def handle_redirect(auth_url: str) -> None:
87
+ nonlocal resume
88
+ if resume:
89
+ raise RuntimeError("Another redirect is already pending")
90
+ message = self.create_auth_request(authorization_endpoint_url=pydantic.AnyUrl(auth_url))
91
+ resume = await self.context.yield_async(AuthRequired(message=message))
92
+
93
+ async def handle_callback() -> tuple[str, str | None]:
94
+ nonlocal resume
95
+ try:
96
+ if not resume:
97
+ raise ValueError("Missing resume data")
98
+ response = self.parse_auth_response(message=resume)
99
+ params = parse_qs(response.redirect_uri.query)
100
+ return params["code"][0], params.get("state", [None])[0]
101
+ finally:
102
+ resume = None
103
+
104
+ # A2A Client is responsible for catching the redirect and forwarding it over the A2A connection
105
+ oauth_auth = OAuthClientProvider(
106
+ server_url=str(resource_url),
107
+ client_metadata=OAuthClientMetadata(
108
+ redirect_uris=[fulfillment.redirect_uri],
109
+ ),
110
+ storage=await self.token_storage_factory.create_storage(),
111
+ redirect_handler=handle_redirect,
112
+ callback_handler=handle_callback,
113
+ )
114
+ return oauth_auth
115
+
116
+ def create_auth_request(self, *, authorization_endpoint_url: pydantic.AnyUrl):
117
+ data = AuthRequest(authorization_endpoint_url=authorization_endpoint_url)
118
+ return AgentMessage(text="Authorization required", metadata={self.spec.URI: data.model_dump(mode="json")})
119
+
120
+ def parse_auth_response(self, *, message: a2a.types.Message):
121
+ if not message or not message.metadata or not (data := message.metadata.get(self.spec.URI)):
122
+ raise RuntimeError("Invalid auth response")
123
+ return AuthResponse.model_validate(data)
124
+
125
+
126
+ class OAuthExtensionClient(BaseExtensionClient[OAuthExtensionSpec, NoneType]):
127
+ def fulfillment_metadata(self, *, oauth_fulfillments: dict[str, Any]) -> dict[str, Any]:
128
+ return {self.spec.URI: OAuthExtensionMetadata(oauth_fulfillments=oauth_fulfillments).model_dump(mode="json")}
129
+
130
+ def parse_auth_request(self, *, message: a2a.types.Message):
131
+ if not message or not message.metadata or not (data := message.metadata.get(self.spec.URI)):
132
+ raise ValueError("Invalid auth request")
133
+ return AuthRequest.model_validate(data)
134
+
135
+ def create_auth_response(self, *, task_id: str, redirect_uri: pydantic.AnyUrl):
136
+ data = AuthResponse(redirect_uri=redirect_uri)
137
+
138
+ return a2a.types.Message(
139
+ message_id=str(uuid.uuid4()),
140
+ role=a2a.types.Role.user,
141
+ parts=[a2a.types.TextPart(text="Authorization completed")], # type: ignore
142
+ task_id=task_id,
143
+ metadata={self.spec.URI: data.model_dump(mode="json")},
144
+ )
@@ -0,0 +1,5 @@
1
+ # Copyright 2025 © BeeAI a Series of LF Projects, LLC
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ from .base import *
5
+ from .memory import *
@@ -0,0 +1,11 @@
1
+ # Copyright 2025 © BeeAI a Series of LF Projects, LLC
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ import abc
5
+
6
+ from mcp.client.auth import TokenStorage
7
+
8
+
9
+ class TokenStorageFactory(abc.ABC):
10
+ @abc.abstractmethod
11
+ async def create_storage(self) -> TokenStorage: ...
@@ -0,0 +1,38 @@
1
+ # Copyright 2025 © BeeAI a Series of LF Projects, LLC
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+
5
+ from mcp.client.auth import TokenStorage
6
+ from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
7
+
8
+ from .base import TokenStorageFactory
9
+
10
+
11
+ class MemoryTokenStorage(TokenStorage):
12
+ def __init__(self):
13
+ self.tokens: OAuthToken | None = None
14
+ self.client_info: OAuthClientInformationFull | None = None
15
+
16
+ async def get_tokens(self) -> OAuthToken | None:
17
+ return self.tokens
18
+
19
+ async def set_tokens(self, tokens: OAuthToken) -> None:
20
+ self.tokens = tokens
21
+
22
+ async def get_client_info(self) -> OAuthClientInformationFull | None:
23
+ return self.client_info
24
+
25
+ async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
26
+ self.client_info = client_info
27
+
28
+
29
+ class MemoryTokenStorageFactory(TokenStorageFactory):
30
+ def __init__(self, *, client_info: OAuthClientInformationFull | None = None):
31
+ super().__init__()
32
+ self._client_info = client_info
33
+
34
+ async def create_storage(self) -> TokenStorage:
35
+ storage = MemoryTokenStorage()
36
+ if self._client_info:
37
+ await storage.set_client_info(self._client_info)
38
+ return storage
@@ -0,0 +1,4 @@
1
+ # Copyright 2025 © BeeAI a Series of LF Projects, LLC
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ from .secrets import *
@@ -0,0 +1,71 @@
1
+ # Copyright 2025 © BeeAI a Series of LF Projects, LLC
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ import typing
5
+ from typing import TYPE_CHECKING
6
+
7
+ import pydantic
8
+ from a2a.types import Message as A2AMessage
9
+
10
+ from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
11
+ from agentstack_sdk.a2a.types import AgentMessage, AuthRequired
12
+
13
+ if TYPE_CHECKING:
14
+ from agentstack_sdk.server.context import RunContext
15
+
16
+
17
+ class SecretDemand(pydantic.BaseModel):
18
+ name: str
19
+ description: str | None = None
20
+
21
+
22
+ class SecretFulfillment(pydantic.BaseModel):
23
+ secret: str
24
+
25
+
26
+ class SecretsServiceExtensionParams(pydantic.BaseModel):
27
+ secret_demands: dict[str, SecretDemand]
28
+
29
+
30
+ class SecretsServiceExtensionMetadata(pydantic.BaseModel):
31
+ secret_fulfillments: dict[str, SecretFulfillment] = {}
32
+
33
+
34
+ class SecretsExtensionSpec(BaseExtensionSpec[SecretsServiceExtensionParams | None]):
35
+ URI: str = "https://a2a-extensions.agentstack.beeai.dev/auth/secrets/v1"
36
+
37
+ @classmethod
38
+ def single_demand(cls, name: str, key: str | None = None, description: str | None = None) -> typing.Self:
39
+ return cls(
40
+ params=SecretsServiceExtensionParams(
41
+ secret_demands={key or "default": SecretDemand(description=description, name=name)}
42
+ )
43
+ )
44
+
45
+
46
+ class SecretsExtensionServer(BaseExtensionServer[SecretsExtensionSpec, SecretsServiceExtensionMetadata]):
47
+ def handle_incoming_message(self, message: A2AMessage, context: "RunContext"):
48
+ super().handle_incoming_message(message, context)
49
+ self.context = context
50
+
51
+ def parse_secret_response(self, message: A2AMessage) -> SecretsServiceExtensionMetadata:
52
+ if not message or not message.metadata or not (data := message.metadata.get(self.spec.URI)):
53
+ raise ValueError("Secrets has not been provided in response.")
54
+
55
+ return SecretsServiceExtensionMetadata.model_validate(data)
56
+
57
+ async def request_secrets(self, params: SecretsServiceExtensionParams) -> SecretsServiceExtensionMetadata:
58
+ resume = await self.context.yield_async(
59
+ AuthRequired(
60
+ message=AgentMessage(
61
+ metadata={self.spec.URI: params.model_dump(mode="json")},
62
+ )
63
+ )
64
+ )
65
+ if isinstance(resume, A2AMessage):
66
+ return self.parse_secret_response(message=resume)
67
+ else:
68
+ raise ValueError("Secrets has not been provided in response.")
69
+
70
+
71
+ class SecretsExtensionClient(BaseExtensionClient[SecretsExtensionSpec, SecretsServiceExtensionParams]): ...