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.
- agentstack_sdk-0.4.1rc1/PKG-INFO +69 -0
- agentstack_sdk-0.4.1rc1/README.md +45 -0
- agentstack_sdk-0.4.1rc1/pyproject.toml +77 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/__init__.py +6 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/__init__.py +2 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/__init__.py +6 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/__init__.py +5 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/oauth/__init__.py +4 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/oauth/oauth.py +144 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/oauth/storage/__init__.py +5 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/oauth/storage/base.py +11 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/oauth/storage/memory.py +38 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/secrets/__init__.py +4 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/auth/secrets/secrets.py +71 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/base.py +196 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/exceptions.py +11 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/services/__init__.py +7 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/services/embedding.py +99 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/services/llm.py +98 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/services/mcp.py +149 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/services/platform.py +117 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/ui/__init__.py +8 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/ui/agent_detail.py +54 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/ui/canvas.py +65 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/ui/citation.py +83 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/ui/form.py +193 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/ui/settings.py +73 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/extensions/ui/trajectory.py +66 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/a2a/types.py +106 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/__init__.py +10 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/client.py +119 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/common.py +37 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/configuration.py +47 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/context.py +263 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/file.py +213 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/model_provider.py +131 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/provider.py +219 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/provider_build.py +190 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/types.py +45 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/variables.py +44 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/platform/vector_store.py +205 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/py.typed +0 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/__init__.py +4 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/agent.py +555 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/app.py +69 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/constants.py +8 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/context.py +53 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/dependencies.py +118 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/logging.py +63 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/server.py +360 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/store/__init__.py +3 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/store/context_store.py +28 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/store/memory_context_store.py +43 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/store/platform_context_store.py +48 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/telemetry.py +53 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/server/utils.py +26 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/util/__init__.py +4 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/util/file.py +260 -0
- agentstack_sdk-0.4.1rc1/src/agentstack_sdk/util/resource_context.py +44 -0
- 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,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,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,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]): ...
|