google-adk-extras 0.2.3__tar.gz → 0.2.6__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 (70) hide show
  1. {google_adk_extras-0.2.3/src/google_adk_extras.egg-info → google_adk_extras-0.2.6}/PKG-INFO +1 -1
  2. google_adk_extras-0.2.6/docs/streaming.md +119 -0
  3. google_adk_extras-0.2.6/examples/streaming_sse_ws.py +31 -0
  4. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/mkdocs.yml +1 -0
  5. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/pyproject.toml +1 -1
  6. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/__init__.py +1 -1
  7. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/custom_agent_loader.py +15 -1
  8. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/enhanced_fastapi.py +138 -4
  9. google_adk_extras-0.2.6/src/google_adk_extras/streaming/__init__.py +12 -0
  10. google_adk_extras-0.2.6/src/google_adk_extras/streaming/streaming_controller.py +262 -0
  11. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6/src/google_adk_extras.egg-info}/PKG-INFO +1 -1
  12. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras.egg-info/SOURCES.txt +4 -0
  13. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/LICENSE +0 -0
  14. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/MANIFEST.in +0 -0
  15. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/README.md +0 -0
  16. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/docs/agent-loading.md +0 -0
  17. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/docs/credentials.md +0 -0
  18. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/docs/examples.md +0 -0
  19. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/docs/fastapi.md +0 -0
  20. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/docs/getting-started.md +0 -0
  21. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/docs/index.md +0 -0
  22. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/docs/quickstarts.md +0 -0
  23. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/docs/services.md +0 -0
  24. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/docs/troubleshooting.md +0 -0
  25. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/docs/uris.md +0 -0
  26. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/examples/README.md +0 -0
  27. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/examples/consume_remote_a2a.py +0 -0
  28. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/examples/credentials/google_oauth2.py +0 -0
  29. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/examples/custom_loader.py +0 -0
  30. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/examples/fastapi_app.py +0 -0
  31. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/examples/programmatic_a2a_expose.py +0 -0
  32. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/examples/runner_basic.py +0 -0
  33. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/examples/services/artifacts_local.py +0 -0
  34. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/examples/services/memory_yaml.py +0 -0
  35. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/examples/services/sessions_sql.py +0 -0
  36. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/setup.cfg +0 -0
  37. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/setup.py +0 -0
  38. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/adk_builder.py +0 -0
  39. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/artifacts/__init__.py +0 -0
  40. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/artifacts/base_custom_artifact_service.py +0 -0
  41. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/artifacts/local_folder_artifact_service.py +0 -0
  42. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/artifacts/mongo_artifact_service.py +0 -0
  43. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/artifacts/s3_artifact_service.py +0 -0
  44. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/artifacts/sql_artifact_service.py +0 -0
  45. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/credentials/__init__.py +0 -0
  46. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/credentials/base_custom_credential_service.py +0 -0
  47. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/credentials/github_oauth2_credential_service.py +0 -0
  48. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/credentials/google_oauth2_credential_service.py +0 -0
  49. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/credentials/http_basic_auth_credential_service.py +0 -0
  50. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/credentials/jwt_credential_service.py +0 -0
  51. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/credentials/microsoft_oauth2_credential_service.py +0 -0
  52. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/credentials/x_oauth2_credential_service.py +0 -0
  53. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/enhanced_adk_web_server.py +0 -0
  54. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/enhanced_runner.py +0 -0
  55. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/memory/__init__.py +0 -0
  56. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/memory/base_custom_memory_service.py +0 -0
  57. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/memory/mongo_memory_service.py +0 -0
  58. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/memory/redis_memory_service.py +0 -0
  59. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/memory/sql_memory_service.py +0 -0
  60. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/memory/yaml_file_memory_service.py +0 -0
  61. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/sessions/__init__.py +0 -0
  62. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/sessions/base_custom_session_service.py +0 -0
  63. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/sessions/mongo_session_service.py +0 -0
  64. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/sessions/redis_session_service.py +0 -0
  65. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/sessions/sql_session_service.py +0 -0
  66. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras/sessions/yaml_file_session_service.py +0 -0
  67. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras.egg-info/dependency_links.txt +0 -0
  68. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras.egg-info/requires.txt +0 -0
  69. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/src/google_adk_extras.egg-info/top_level.txt +0 -0
  70. {google_adk_extras-0.2.3 → google_adk_extras-0.2.6}/tests/test_a2a_helpers.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: google-adk-extras
3
- Version: 0.2.3
3
+ Version: 0.2.6
4
4
  Summary: Production-ready services, credentials, and FastAPI wiring for Google ADK
5
5
  Home-page: https://github.com/DeadMeme5441/google-adk-extras
6
6
  Author: DeadMeme5441
@@ -0,0 +1,119 @@
1
+ # Streaming
2
+
3
+ The optional streaming layer adds a persistent, bi‑directional channel on top of
4
+ Google ADK’s existing `/run`, `/run_sse`, and `/run_live` primitives. It exposes
5
+ simple endpoints for a chat‑style UI while keeping strict ADK type parity.
6
+
7
+ - Uplink payload: the exact ADK `RunAgentRequest` (JSON)
8
+ - Downlink events: the exact ADK `Event` JSON
9
+ - Per‑channel binding to `(appName, userId, sessionId)`
10
+
11
+ ## Enable
12
+
13
+ ```python
14
+ from google_adk_extras import AdkBuilder
15
+
16
+ app = (
17
+ AdkBuilder()
18
+ .with_agents_dir("./agents") # or programmatic loader/instances
19
+ .build_fastapi_app(enable_streaming=True)
20
+ )
21
+ ```
22
+
23
+ Optional config:
24
+
25
+ ```python
26
+ from google_adk_extras.streaming import StreamingConfig
27
+
28
+ cfg = StreamingConfig(
29
+ streaming_path_base="/stream",
30
+ strict_types=True,
31
+ create_session_on_open=True, # create a new ADK session on first subscribe
32
+ ttl_seconds=900,
33
+ max_queue_size=128,
34
+ max_channels_per_user=20,
35
+ heartbeat_interval=20.0,
36
+ )
37
+
38
+ app = AdkBuilder().with_agents_dir("./agents").build_fastapi_app(
39
+ enable_streaming=True,
40
+ streaming_config=cfg,
41
+ )
42
+ ```
43
+
44
+ ## Endpoints (default base: `/stream`)
45
+
46
+ - `GET /stream/events/{channelId}?appName=&userId=&sessionId=` — SSE downlink
47
+ - If `sessionId` is omitted and `create_session_on_open=True`, a session is created.
48
+ - First, a control message is sent: `event: channel-bound` with `data: {appName,userId,sessionId}`.
49
+ - Subsequent `data: ...` lines contain raw ADK `Event` JSON.
50
+
51
+ - `POST /stream/send/{channelId}` — enqueue a single run on the bound channel
52
+ - Body must be a strict ADK `RunAgentRequest` JSON whose `(appName,userId,sessionId)`
53
+ match the channel binding.
54
+
55
+ - `WS /stream/ws/{channelId}?appName=&userId=&sessionId=` — WebSocket bidi
56
+ - On connect, a `{"event":"channel-bound",...}` JSON frame is sent with the final `sessionId`.
57
+ - Client sends strict `RunAgentRequest` JSON frames to enqueue runs.
58
+ - Server streams raw ADK `Event` JSON frames back.
59
+
60
+ ## Minimal Client Examples
61
+
62
+ SSE (browser):
63
+
64
+ ```html
65
+ <script>
66
+ const ch = crypto.randomUUID();
67
+ const src = new EventSource(`/stream/events/${ch}?appName=my_app&userId=u1`);
68
+ let sessionId;
69
+ src.addEventListener('channel-bound', (e) => {
70
+ const info = JSON.parse(e.data);
71
+ sessionId = info.sessionId;
72
+ // Now send a RunAgentRequest
73
+ fetch(`/stream/send/${ch}`, {
74
+ method: 'POST',
75
+ headers: {'Content-Type':'application/json'},
76
+ body: JSON.stringify({
77
+ appName: "my_app",
78
+ userId: "u1",
79
+ sessionId,
80
+ streaming: true,
81
+ newMessage: { parts: [{ text: "Hello!" }] }
82
+ })
83
+ });
84
+ });
85
+
86
+ src.onmessage = (e) => {
87
+ const event = JSON.parse(e.data); // ADK Event JSON
88
+ console.log('event', event);
89
+ };
90
+ </script>
91
+ ```
92
+
93
+ WebSocket (browser):
94
+
95
+ ```html
96
+ <script>
97
+ const ch = crypto.randomUUID();
98
+ const ws = new WebSocket(`ws://${location.host}/stream/ws/${ch}?appName=my_app&userId=u1`);
99
+ let sessionId;
100
+ ws.onmessage = (msg) => {
101
+ const data = JSON.parse(msg.data);
102
+ if (data.event === 'channel-bound') {
103
+ sessionId = data.sessionId;
104
+ ws.send(JSON.stringify({
105
+ appName: 'my_app', userId: 'u1', sessionId, streaming: true,
106
+ newMessage: { parts: [{ text: 'Hello!' }] }
107
+ }));
108
+ } else {
109
+ console.log('event', data); // ADK Event JSON
110
+ }
111
+ };
112
+ </script>
113
+ ```
114
+
115
+ ### Notes
116
+ - This layer is optional. The core ADK endpoints remain available.
117
+ - By default we preserve ADK wire types. If you need a convenience payload
118
+ (non‑strict), you can add your own translator at your API boundary.
119
+
@@ -0,0 +1,31 @@
1
+ """Streaming (SSE + WebSocket) example using AdkBuilder.
2
+
3
+ Run:
4
+ uvicorn examples.streaming_sse_ws:app --reload
5
+
6
+ This example enables the optional streaming layer and exposes routes under
7
+ `/stream`:
8
+ - GET /stream/events/{channelId}?appName=&userId=&sessionId= (SSE)
9
+ - POST /stream/send/{channelId} (body = RunAgentRequest JSON)
10
+ - WS /stream/ws/{channelId}?appName=&userId=&sessionId=
11
+
12
+ Notes:
13
+ - On open, the server emits a channel-bound control message containing the
14
+ bound sessionId so the client can populate RunAgentRequest.
15
+ """
16
+
17
+ from google_adk_extras import AdkBuilder
18
+
19
+
20
+ app = (
21
+ AdkBuilder()
22
+ # Use either on-disk agents or programmatic agents. For streaming layer,
23
+ # the agent loader is optional if you stub the runner in tests; in real apps
24
+ # provide your agents via one of these methods:
25
+ # .with_agents_dir("./agents")
26
+
27
+ # Enable streaming layer with defaults
28
+ .with_web_ui(False)
29
+ .build_fastapi_app(enable_streaming=True)
30
+ )
31
+
@@ -14,6 +14,7 @@ nav:
14
14
  - Getting Started: getting-started.md
15
15
  - Quickstarts: quickstarts.md
16
16
  - FastAPI Integration: fastapi.md
17
+ - Streaming: streaming.md
17
18
  - Services:
18
19
  - Durable Services: services.md
19
20
  - Credentials: credentials.md
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "google-adk-extras"
3
- version = "0.2.3"
3
+ version = "0.2.6"
4
4
  description = "Production-ready services, credentials, and FastAPI wiring for Google ADK"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.10,<3.13"
@@ -28,4 +28,4 @@ __all__ = [
28
28
  "CustomAgentLoader",
29
29
  ]
30
30
 
31
- __version__ = "0.2.3"
31
+ __version__ = "0.2.6"
@@ -146,6 +146,20 @@ class CustomAgentLoader(BaseAgentLoader):
146
146
  sorted_agents = sorted(agent_names)
147
147
  logger.debug("Total registered agents: %d", len(sorted_agents))
148
148
  return sorted_agents
149
+
150
+ # Compatibility with ADK's AgentLoader API used by AgentChangeEventHandler
151
+ def remove_agent_from_cache(self, name: str) -> None:
152
+ """No-op cache invalidation for compatibility with ADK hot reload.
153
+
154
+ ADK's file-watcher calls `agent_loader.remove_agent_from_cache(current_app)`
155
+ when files change. Our loader does not cache filesystem-loaded agents,
156
+ but we provide this method to satisfy the expected interface.
157
+
158
+ Args:
159
+ name: Agent name to invalidate (ignored here).
160
+ """
161
+ # Nothing to do; present for interface compatibility.
162
+ logger.debug("CustomAgentLoader.remove_agent_from_cache(%s) - no-op", name)
149
163
 
150
164
 
151
165
  def __repr__(self) -> str:
@@ -153,4 +167,4 @@ class CustomAgentLoader(BaseAgentLoader):
153
167
  with self._lock:
154
168
  registered_count = len(self._registered_agents)
155
169
 
156
- return f"CustomAgentLoader(registered={registered_count})"
170
+ return f"CustomAgentLoader(registered={registered_count})"
@@ -5,6 +5,7 @@ that properly supports custom credential services.
5
5
  """
6
6
 
7
7
  import json
8
+ import asyncio
8
9
  import logging
9
10
  import os
10
11
  from pathlib import Path
@@ -34,6 +35,7 @@ from google.adk.sessions.database_session_service import DatabaseSessionService
34
35
  from google.adk.utils.feature_decorator import working_in_progress
35
36
  from google.adk.cli.adk_web_server import AdkWebServer
36
37
  from .enhanced_adk_web_server import EnhancedAdkWebServer
38
+ from .streaming import StreamingController, StreamingConfig
37
39
  from google.adk.cli.utils import envs
38
40
  from google.adk.cli.utils import evals
39
41
  from google.adk.cli.utils.agent_change_handler import AgentChangeEventHandler
@@ -64,6 +66,9 @@ def get_enhanced_fast_api_app(
64
66
  trace_to_cloud: bool = False,
65
67
  reload_agents: bool = False,
66
68
  lifespan: Optional[Lifespan[FastAPI]] = None,
69
+ # Streaming layer (optional)
70
+ enable_streaming: bool = False,
71
+ streaming_config: Optional[StreamingConfig] = None,
67
72
  ) -> FastAPI:
68
73
  """Enhanced version of Google ADK's get_fast_api_app with EnhancedRunner integration.
69
74
 
@@ -151,7 +156,7 @@ def get_enhanced_fast_api_app(
151
156
  agent_engine_id = agent_engine_id_or_resource_name
152
157
  return project, location, agent_engine_id
153
158
 
154
- # Build the Memory service (same as ADK)
159
+ # Build the Memory service (enhanced to recognize extras URIs)
155
160
  if memory_service_uri:
156
161
  if memory_service_uri.startswith("rag://"):
157
162
  from google.adk.memory.vertex_ai_rag_memory_service import VertexAiRagMemoryService
@@ -172,6 +177,19 @@ def get_enhanced_fast_api_app(
172
177
  location=location,
173
178
  agent_engine_id=agent_engine_id,
174
179
  )
180
+ elif memory_service_uri.startswith("yaml://"):
181
+ from .memory.yaml_file_memory_service import YamlFileMemoryService
182
+ base_directory = memory_service_uri.split("://")[1]
183
+ memory_service = YamlFileMemoryService(base_directory=base_directory)
184
+ elif memory_service_uri.startswith("redis://"):
185
+ from .memory.redis_memory_service import RedisMemoryService
186
+ memory_service = RedisMemoryService(connection_string=memory_service_uri) # type: ignore[arg-type]
187
+ elif memory_service_uri.startswith(("sqlite://", "postgresql://", "mysql://")):
188
+ from .memory.sql_memory_service import SQLMemoryService
189
+ memory_service = SQLMemoryService(database_url=memory_service_uri)
190
+ elif memory_service_uri.startswith("mongodb://"):
191
+ from .memory.mongo_memory_service import MongoMemoryService
192
+ memory_service = MongoMemoryService(connection_string=memory_service_uri)
175
193
  else:
176
194
  raise click.ClickException(
177
195
  "Unsupported memory service URI: %s" % memory_service_uri
@@ -179,7 +197,7 @@ def get_enhanced_fast_api_app(
179
197
  else:
180
198
  memory_service = InMemoryMemoryService()
181
199
 
182
- # Build the Session service (same as ADK)
200
+ # Build the Session service (enhanced to recognize extras URIs)
183
201
  if session_service_uri:
184
202
  if session_service_uri.startswith("agentengine://"):
185
203
  agent_engine_id_or_resource_name = session_service_uri.split("://")[1]
@@ -191,8 +209,18 @@ def get_enhanced_fast_api_app(
191
209
  location=location,
192
210
  agent_engine_id=agent_engine_id,
193
211
  )
212
+ elif session_service_uri.startswith("yaml://"):
213
+ from .sessions.yaml_file_session_service import YamlFileSessionService
214
+ base_directory = session_service_uri.split("://")[1]
215
+ session_service = YamlFileSessionService(base_directory=base_directory)
216
+ elif session_service_uri.startswith("redis://"):
217
+ from .sessions.redis_session_service import RedisSessionService
218
+ session_service = RedisSessionService(connection_string=session_service_uri) # type: ignore[arg-type]
219
+ elif session_service_uri.startswith("mongodb://"):
220
+ from .sessions.mongo_session_service import MongoSessionService
221
+ session_service = MongoSessionService(connection_string=session_service_uri)
194
222
  else:
195
- # Database session additional settings
223
+ # Treat remaining schemes as database URLs (sqlite/postgres/mysql)
196
224
  if session_db_kwargs is None:
197
225
  session_db_kwargs = {}
198
226
  session_service = DatabaseSessionService(
@@ -201,11 +229,25 @@ def get_enhanced_fast_api_app(
201
229
  else:
202
230
  session_service = InMemorySessionService()
203
231
 
204
- # Build the Artifact service (same as ADK)
232
+ # Build the Artifact service (enhanced to recognize extras URIs)
205
233
  if artifact_service_uri:
206
234
  if artifact_service_uri.startswith("gs://"):
207
235
  gcs_bucket = artifact_service_uri.split("://")[1]
208
236
  artifact_service = GcsArtifactService(bucket_name=gcs_bucket)
237
+ elif artifact_service_uri.startswith("local://"):
238
+ from .artifacts.local_folder_artifact_service import LocalFolderArtifactService
239
+ base_directory = artifact_service_uri.split("://")[1]
240
+ artifact_service = LocalFolderArtifactService(base_directory=base_directory)
241
+ elif artifact_service_uri.startswith("s3://"):
242
+ from .artifacts.s3_artifact_service import S3ArtifactService
243
+ bucket_name = artifact_service_uri.split("://")[1]
244
+ artifact_service = S3ArtifactService(bucket_name=bucket_name)
245
+ elif artifact_service_uri.startswith(("sqlite://", "postgresql://", "mysql://")):
246
+ from .artifacts.sql_artifact_service import SQLArtifactService
247
+ artifact_service = SQLArtifactService(database_url=artifact_service_uri)
248
+ elif artifact_service_uri.startswith("mongodb://"):
249
+ from .artifacts.mongo_artifact_service import MongoArtifactService
250
+ artifact_service = MongoArtifactService(connection_string=artifact_service_uri)
209
251
  else:
210
252
  raise click.ClickException(
211
253
  "Unsupported artifact service URI: %s" % artifact_service_uri
@@ -467,4 +509,96 @@ def get_enhanced_fast_api_app(
467
509
  logger.error("Failed to setup programmatic A2A agent %s: %s", app_name, e)
468
510
 
469
511
  logger.info("Enhanced FastAPI app created with credential service support")
512
+
513
+ # Optional streaming mounts (SSE + WebSocket)
514
+ if enable_streaming:
515
+ cfg = streaming_config or StreamingConfig(enable_streaming=True)
516
+ controller = StreamingController(
517
+ config=cfg,
518
+ get_runner_async=adk_web_server.get_runner_async,
519
+ session_service=session_service,
520
+ )
521
+ app.state.streaming_controller = controller
522
+ @app.on_event("startup")
523
+ async def _start_streaming(): # pragma: no cover - lifecycle glue
524
+ controller.start()
525
+ @app.on_event("shutdown")
526
+ async def _stop_streaming(): # pragma: no cover - lifecycle glue
527
+ await controller.stop()
528
+
529
+ from fastapi import APIRouter, WebSocket, Query
530
+ from fastapi.responses import StreamingResponse
531
+ from google.adk.cli.adk_web_server import RunAgentRequest
532
+
533
+ router = APIRouter()
534
+ base = cfg.streaming_path_base.rstrip("/")
535
+
536
+ @router.get(f"{base}/events/{{channel_id}}")
537
+ async def stream_events(channel_id: str, appName: str = Query(...), userId: str = Query(...), sessionId: Optional[str] = Query(None)):
538
+ ch = await app.state.streaming_controller.open_or_bind_channel(
539
+ channel_id=channel_id, app_name=appName, user_id=userId, session_id=sessionId
540
+ )
541
+ q = app.state.streaming_controller.subscribe(channel_id, kind="sse")
542
+
543
+ async def gen():
544
+ try:
545
+ # Announce channel binding with session id
546
+ yield "event: channel-bound\n"
547
+ yield f"data: {{\"appName\":\"{appName}\",\"userId\":\"{userId}\",\"sessionId\":\"{ch.session_id}\"}}\n\n"
548
+ while True:
549
+ payload = await q.get()
550
+ yield f"data: {payload}\n\n"
551
+ except asyncio.CancelledError:
552
+ pass
553
+ finally:
554
+ app.state.streaming_controller.unsubscribe(channel_id, q)
555
+
556
+ return StreamingResponse(gen(), media_type="text/event-stream")
557
+
558
+ @router.post(f"{base}/send/{{channel_id}}")
559
+ async def send_message(channel_id: str, req: RunAgentRequest):
560
+ # Validation: channel binding must match
561
+ await app.state.streaming_controller.enqueue(channel_id, req)
562
+ return PlainTextResponse("", status_code=204)
563
+
564
+ @router.websocket(f"{base}/ws/{{channel_id}}")
565
+ async def ws_endpoint(websocket: WebSocket, channel_id: str, appName: str, userId: str, sessionId: Optional[str] = None):
566
+ await websocket.accept()
567
+ try:
568
+ await app.state.streaming_controller.open_or_bind_channel(
569
+ channel_id=channel_id, app_name=appName, user_id=userId, session_id=sessionId
570
+ )
571
+ q = app.state.streaming_controller.subscribe(channel_id, kind="ws")
572
+ # Send channel binding info including session id
573
+ await websocket.send_text(json.dumps({"event": "channel-bound", "appName": appName, "userId": userId, "sessionId": app.state.streaming_controller._channels[channel_id].session_id}))
574
+
575
+ async def downlink():
576
+ try:
577
+ while True:
578
+ payload = await q.get()
579
+ await websocket.send_text(payload)
580
+ except asyncio.CancelledError:
581
+ pass
582
+
583
+ async def uplink():
584
+ try:
585
+ while True:
586
+ text = await websocket.receive_text()
587
+ # Strict type parity by default
588
+ req = RunAgentRequest.model_validate_json(text)
589
+ await app.state.streaming_controller.enqueue(channel_id, req)
590
+ except Exception:
591
+ return
592
+
593
+ down = asyncio.create_task(downlink())
594
+ up = asyncio.create_task(uplink())
595
+ await asyncio.wait({down, up}, return_when=asyncio.FIRST_COMPLETED)
596
+ finally:
597
+ try:
598
+ app.state.streaming_controller.unsubscribe(channel_id, q)
599
+ except Exception:
600
+ pass
601
+
602
+ app.include_router(router)
603
+
470
604
  return app
@@ -0,0 +1,12 @@
1
+ """Streaming support (SSE/WebSocket) for google-adk-extras.
2
+
3
+ This package provides an optional, persistent bi-directional streaming layer
4
+ with strict ADK type parity by default. It complements ADK's built-in
5
+ `/run`, `/run_sse`, and `/run_live` endpoints by offering per-channel
6
+ subscription and send semantics for chat-style UIs.
7
+ """
8
+
9
+ from .streaming_controller import StreamingConfig, StreamingController
10
+
11
+ __all__ = ["StreamingConfig", "StreamingController"]
12
+
@@ -0,0 +1,262 @@
1
+ import asyncio
2
+ import time
3
+ from dataclasses import dataclass, field
4
+ from typing import Any, Dict, Optional, Set, Callable, Awaitable
5
+
6
+ from fastapi import WebSocket, HTTPException
7
+ from pydantic import BaseModel
8
+
9
+ from google.adk.events.event import Event
10
+ from google.adk.runners import Runner
11
+
12
+
13
+ class StreamingConfig(BaseModel):
14
+ enable_streaming: bool = False
15
+ streaming_path_base: str = "/stream"
16
+ strict_types: bool = True
17
+ create_session_on_open: bool = True
18
+ ttl_seconds: int = 900
19
+ max_queue_size: int = 128
20
+ max_channels_per_user: int = 20
21
+ heartbeat_interval: Optional[float] = 20.0
22
+ reuse_session_policy: str = "per_channel" # "per_channel" or "external"
23
+
24
+
25
+ @dataclass
26
+ class _Subscriber:
27
+ queue: "asyncio.Queue[str]"
28
+ kind: str # "sse" | "ws"
29
+
30
+
31
+ @dataclass
32
+ class _Channel:
33
+ channel_id: str
34
+ app_name: str
35
+ user_id: str
36
+ session_id: str
37
+ in_q: "asyncio.Queue[Any]" = field(default_factory=asyncio.Queue)
38
+ subscribers: list[_Subscriber] = field(default_factory=list)
39
+ worker_task: Optional[asyncio.Task] = None
40
+ created_at: float = field(default_factory=lambda: time.time())
41
+ last_activity: float = field(default_factory=lambda: time.time())
42
+ lock: asyncio.Lock = field(default_factory=asyncio.Lock)
43
+
44
+
45
+ class StreamingController:
46
+ """Manages streaming channels and workers.
47
+
48
+ This controller binds a channel to (app_name, user_id, session_id) and
49
+ runs a background worker per channel to execute streamed runs and push
50
+ ADK Event JSON to all subscribers.
51
+ """
52
+
53
+ def __init__(
54
+ self,
55
+ *,
56
+ config: StreamingConfig,
57
+ get_runner_async: Callable[[str], Awaitable[Runner]],
58
+ session_service,
59
+ ) -> None:
60
+ self._config = config
61
+ self._get_runner_async = get_runner_async
62
+ self._session_service = session_service
63
+ self._channels: Dict[str, _Channel] = {}
64
+ self._gc_task: Optional[asyncio.Task] = None
65
+
66
+ def start(self) -> None:
67
+ if self._gc_task is None:
68
+ self._gc_task = asyncio.create_task(self._gc_loop())
69
+
70
+ async def stop(self) -> None:
71
+ if self._gc_task:
72
+ self._gc_task.cancel()
73
+ with asyncio.CancelledError:
74
+ pass
75
+ self._gc_task = None
76
+ # Cancel workers
77
+ for ch in list(self._channels.values()):
78
+ if ch.worker_task and not ch.worker_task.done():
79
+ ch.worker_task.cancel()
80
+ self._channels.clear()
81
+
82
+ def _ensure_user_limit(self, user_id: str) -> None:
83
+ if self._config.max_channels_per_user <= 0:
84
+ return
85
+ count = sum(1 for c in self._channels.values() if c.user_id == user_id)
86
+ if count >= self._config.max_channels_per_user:
87
+ raise HTTPException(status_code=429, detail="Too many channels for this user")
88
+
89
+ async def open_or_bind_channel(
90
+ self,
91
+ *,
92
+ channel_id: str,
93
+ app_name: str,
94
+ user_id: str,
95
+ session_id: Optional[str],
96
+ ) -> _Channel:
97
+ # Existing channel validation/match
98
+ if channel_id in self._channels:
99
+ ch = self._channels[channel_id]
100
+ if ch.app_name != app_name or ch.user_id != user_id:
101
+ raise HTTPException(status_code=409, detail="Channel binding conflict")
102
+ if session_id and session_id != ch.session_id:
103
+ raise HTTPException(status_code=409, detail="Channel already bound to different session")
104
+ ch.last_activity = time.time()
105
+ return ch
106
+
107
+ # New channel
108
+ self._ensure_user_limit(user_id)
109
+ if not session_id:
110
+ if not self._config.create_session_on_open:
111
+ raise HTTPException(status_code=400, detail="sessionId required for this channel")
112
+ # Create a fresh ADK session
113
+ create = getattr(self._session_service, "create_session", None)
114
+ if create is None:
115
+ # Older ADK interfaces may expose sync variant
116
+ create = getattr(self._session_service, "create_session_sync", None)
117
+ if create is None:
118
+ raise HTTPException(status_code=500, detail="Session service does not support create_session")
119
+ if asyncio.iscoroutinefunction(create):
120
+ session = await create(app_name=app_name, user_id=user_id)
121
+ else:
122
+ # Call sync and wrap
123
+ session = create(app_name=app_name, user_id=user_id)
124
+ session_id = session.id
125
+ else:
126
+ # Validate existing session
127
+ session = await self._session_service.get_session(app_name=app_name, user_id=user_id, session_id=session_id)
128
+ if not session:
129
+ raise HTTPException(status_code=404, detail="Session not found")
130
+
131
+ ch = _Channel(
132
+ channel_id=channel_id,
133
+ app_name=app_name,
134
+ user_id=user_id,
135
+ session_id=session_id,
136
+ in_q=asyncio.Queue(),
137
+ )
138
+ self._channels[channel_id] = ch
139
+ ch.worker_task = asyncio.create_task(self._worker(ch))
140
+ return ch
141
+
142
+ def subscribe(self, channel_id: str, kind: str) -> asyncio.Queue[str]:
143
+ if channel_id not in self._channels:
144
+ raise HTTPException(status_code=404, detail="Channel not found")
145
+ q: asyncio.Queue[str] = asyncio.Queue(maxsize=self._config.max_queue_size)
146
+ self._channels[channel_id].subscribers.append(_Subscriber(queue=q, kind=kind))
147
+ self._channels[channel_id].last_activity = time.time()
148
+ return q
149
+
150
+ def unsubscribe(self, channel_id: str, q: asyncio.Queue[str]) -> None:
151
+ ch = self._channels.get(channel_id)
152
+ if not ch:
153
+ return
154
+ ch.subscribers = [s for s in ch.subscribers if s.queue is not q]
155
+ ch.last_activity = time.time()
156
+
157
+ async def enqueue(self, channel_id: str, req: Any) -> None:
158
+ ch = self._channels.get(channel_id)
159
+ if not ch:
160
+ raise HTTPException(status_code=404, detail="Channel not found")
161
+ # Validate binding
162
+ if getattr(req, "app_name", None) != ch.app_name or getattr(req, "user_id", None) != ch.user_id or getattr(req, "session_id", None) != ch.session_id:
163
+ raise HTTPException(status_code=409, detail="Request does not match channel binding")
164
+ await ch.in_q.put(req)
165
+ ch.last_activity = time.time()
166
+
167
+ async def _worker(self, ch: _Channel) -> None:
168
+ try:
169
+ while True:
170
+ req = await ch.in_q.get()
171
+ ch.last_activity = time.time()
172
+ try:
173
+ runner = await self._get_runner_async(ch.app_name)
174
+ # Stream events for this request
175
+ async with _aclosing(
176
+ runner.run_async(
177
+ user_id=ch.user_id,
178
+ session_id=ch.session_id,
179
+ new_message=req.new_message,
180
+ state_delta=getattr(req, "state_delta", None),
181
+ run_config=_maybe_run_config_streaming(True),
182
+ )
183
+ ) as agen:
184
+ async for event in agen:
185
+ await self._broadcast_event(ch, event)
186
+ except Exception as e: # pragma: no cover - safety
187
+ await self._broadcast_error(ch, str(e))
188
+ except asyncio.CancelledError: # worker shutdown
189
+ return
190
+
191
+ async def _broadcast_event(self, ch: _Channel, event: Event) -> None:
192
+ payload = event.model_dump_json(exclude_none=True, by_alias=True)
193
+ for sub in list(ch.subscribers):
194
+ try:
195
+ sub.queue.put_nowait(payload)
196
+ except asyncio.QueueFull:
197
+ # Drop subscriber on backpressure
198
+ ch.subscribers = [s for s in ch.subscribers if s is not sub]
199
+ ch.last_activity = time.time()
200
+
201
+ async def _broadcast_heartbeat(self, ch: _Channel) -> None:
202
+ if self._config.heartbeat_interval is None:
203
+ return
204
+ payload = '{"event":"heartbeat"}'
205
+ for sub in list(ch.subscribers):
206
+ try:
207
+ sub.queue.put_nowait(payload)
208
+ except asyncio.QueueFull:
209
+ ch.subscribers = [s for s in ch.subscribers if s is not sub]
210
+
211
+ async def _broadcast_error(self, ch: _Channel, message: str) -> None:
212
+ payload = '{"error": %s}' % _json_escape(message)
213
+ for sub in list(ch.subscribers):
214
+ try:
215
+ sub.queue.put_nowait(payload)
216
+ except asyncio.QueueFull:
217
+ ch.subscribers = [s for s in ch.subscribers if s is not sub]
218
+
219
+ async def _gc_loop(self) -> None:
220
+ try:
221
+ while True:
222
+ await asyncio.sleep(min(10, max(1, int(self._config.ttl_seconds / 3))))
223
+ now = time.time()
224
+ for cid, ch in list(self._channels.items()):
225
+ idle = now - ch.last_activity
226
+ if idle >= self._config.ttl_seconds and not ch.subscribers and ch.in_q.empty():
227
+ if ch.worker_task and not ch.worker_task.done():
228
+ ch.worker_task.cancel()
229
+ self._channels.pop(cid, None)
230
+ except asyncio.CancelledError:
231
+ return
232
+
233
+
234
+ # Utilities (avoid importing optional internals at module import time)
235
+ def _maybe_run_config_streaming(enabled: bool):
236
+ # Support multiple ADK versions by resolving RunConfig/StreamingMode from
237
+ # either google.adk.runners or google.adk.agents.run_config
238
+ try:
239
+ from google.adk.runners import RunConfig # type: ignore
240
+ except Exception: # pragma: no cover - version fallback
241
+ from google.adk.agents.run_config import RunConfig # type: ignore
242
+ try:
243
+ from google.adk.agents.run_config import StreamingMode # type: ignore
244
+ except Exception: # pragma: no cover - defensive
245
+ StreamingMode = type("StreamingMode", (), {"SSE": "sse", "NONE": None}) # minimal stub
246
+ return RunConfig(streaming_mode=StreamingMode.SSE if enabled else StreamingMode.NONE)
247
+
248
+
249
+ class _aclosing:
250
+ def __init__(self, agen):
251
+ self._agen = agen
252
+ async def __aenter__(self):
253
+ return self._agen
254
+ async def __aexit__(self, exc_type, exc, tb):
255
+ try:
256
+ await self._agen.aclose()
257
+ except Exception:
258
+ pass
259
+
260
+
261
+ def _json_escape(s: str) -> str:
262
+ return '"' + s.replace('\\', '\\\\').replace('"', '\\"') + '"'
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: google-adk-extras
3
- Version: 0.2.3
3
+ Version: 0.2.6
4
4
  Summary: Production-ready services, credentials, and FastAPI wiring for Google ADK
5
5
  Home-page: https://github.com/DeadMeme5441/google-adk-extras
6
6
  Author: DeadMeme5441
@@ -12,6 +12,7 @@ docs/getting-started.md
12
12
  docs/index.md
13
13
  docs/quickstarts.md
14
14
  docs/services.md
15
+ docs/streaming.md
15
16
  docs/troubleshooting.md
16
17
  docs/uris.md
17
18
  examples/README.md
@@ -20,6 +21,7 @@ examples/custom_loader.py
20
21
  examples/fastapi_app.py
21
22
  examples/programmatic_a2a_expose.py
22
23
  examples/runner_basic.py
24
+ examples/streaming_sse_ws.py
23
25
  examples/credentials/google_oauth2.py
24
26
  examples/services/artifacts_local.py
25
27
  examples/services/memory_yaml.py
@@ -61,4 +63,6 @@ src/google_adk_extras/sessions/mongo_session_service.py
61
63
  src/google_adk_extras/sessions/redis_session_service.py
62
64
  src/google_adk_extras/sessions/sql_session_service.py
63
65
  src/google_adk_extras/sessions/yaml_file_session_service.py
66
+ src/google_adk_extras/streaming/__init__.py
67
+ src/google_adk_extras/streaming/streaming_controller.py
64
68
  tests/test_a2a_helpers.py