google-adk-extras 0.2.5__tar.gz → 0.2.6__tar.gz

{google_adk_extras-0.2.5/src/google_adk_extras.egg-info → google_adk_extras-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-adk-extras
-Version: 0.2.5
+Version: 0.2.6
 Summary: Production-ready services, credentials, and FastAPI wiring for Google ADK
 Home-page: https://github.com/DeadMeme5441/google-adk-extras
 Author: DeadMeme5441

google_adk_extras-0.2.6/docs/streaming.md

@@ -0,0 +1,119 @@
+# Streaming
+
+The optional streaming layer adds a persistent, bi‑directional channel on top of
+Google ADK’s existing `/run`, `/run_sse`, and `/run_live` primitives. It exposes
+simple endpoints for a chat‑style UI while keeping strict ADK type parity.
+
+- Uplink payload: the exact ADK `RunAgentRequest` (JSON)
+- Downlink events: the exact ADK `Event` JSON
+- Per‑channel binding to `(appName, userId, sessionId)`
+
+## Enable
+
+```python
+from google_adk_extras import AdkBuilder
+
+app = (
+    AdkBuilder()
+      .with_agents_dir("./agents")  # or programmatic loader/instances
+      .build_fastapi_app(enable_streaming=True)
+)
+```
+
+Optional config:
+
+```python
+from google_adk_extras.streaming import StreamingConfig
+
+cfg = StreamingConfig(
+    streaming_path_base="/stream",
+    strict_types=True,
+    create_session_on_open=True,  # create a new ADK session on first subscribe
+    ttl_seconds=900,
+    max_queue_size=128,
+    max_channels_per_user=20,
+    heartbeat_interval=20.0,
+)
+
+app = AdkBuilder().with_agents_dir("./agents").build_fastapi_app(
+    enable_streaming=True,
+    streaming_config=cfg,
+)
+```
+
+## Endpoints (default base: `/stream`)
+
+- `GET  /stream/events/{channelId}?appName=&userId=&sessionId=` — SSE downlink
+  - If `sessionId` is omitted and `create_session_on_open=True`, a session is created.
+  - First, a control message is sent: `event: channel-bound` with `data: {appName,userId,sessionId}`.
+  - Subsequent `data: ...` lines contain raw ADK `Event` JSON.
+
+- `POST /stream/send/{channelId}` — enqueue a single run on the bound channel
+  - Body must be a strict ADK `RunAgentRequest` JSON whose `(appName,userId,sessionId)`
+    match the channel binding.
+
+- `WS   /stream/ws/{channelId}?appName=&userId=&sessionId=` — WebSocket bidi
+  - On connect, a `{"event":"channel-bound",...}` JSON frame is sent with the final `sessionId`.
+  - Client sends strict `RunAgentRequest` JSON frames to enqueue runs.
+  - Server streams raw ADK `Event` JSON frames back.
+
+## Minimal Client Examples
+
+SSE (browser):
+
+```html
+<script>
+const ch = crypto.randomUUID();
+const src = new EventSource(`/stream/events/${ch}?appName=my_app&userId=u1`);
+let sessionId;
+src.addEventListener('channel-bound', (e) => {
+  const info = JSON.parse(e.data);
+  sessionId = info.sessionId;
+  // Now send a RunAgentRequest
+  fetch(`/stream/send/${ch}`, {
+    method: 'POST',
+    headers: {'Content-Type':'application/json'},
+    body: JSON.stringify({
+      appName: "my_app",
+      userId: "u1",
+      sessionId,
+      streaming: true,
+      newMessage: { parts: [{ text: "Hello!" }] }
+    })
+  });
+});
+
+src.onmessage = (e) => {
+  const event = JSON.parse(e.data); // ADK Event JSON
+  console.log('event', event);
+};
+</script>
+```
+
+WebSocket (browser):
+
+```html
+<script>
+const ch = crypto.randomUUID();
+const ws = new WebSocket(`ws://${location.host}/stream/ws/${ch}?appName=my_app&userId=u1`);
+let sessionId;
+ws.onmessage = (msg) => {
+  const data = JSON.parse(msg.data);
+  if (data.event === 'channel-bound') {
+    sessionId = data.sessionId;
+    ws.send(JSON.stringify({
+      appName: 'my_app', userId: 'u1', sessionId, streaming: true,
+      newMessage: { parts: [{ text: 'Hello!' }] }
+    }));
+  } else {
+    console.log('event', data); // ADK Event JSON
+  }
+};
+</script>
+```
+
+### Notes
+- This layer is optional. The core ADK endpoints remain available.
+- By default we preserve ADK wire types. If you need a convenience payload
+  (non‑strict), you can add your own translator at your API boundary.
+

google_adk_extras-0.2.6/examples/streaming_sse_ws.py

@@ -0,0 +1,31 @@
+"""Streaming (SSE + WebSocket) example using AdkBuilder.
+
+Run:
+  uvicorn examples.streaming_sse_ws:app --reload
+
+This example enables the optional streaming layer and exposes routes under
+`/stream`:
+- GET  /stream/events/{channelId}?appName=&userId=&sessionId= (SSE)
+- POST /stream/send/{channelId}  (body = RunAgentRequest JSON)
+- WS   /stream/ws/{channelId}?appName=&userId=&sessionId=
+
+Notes:
+- On open, the server emits a channel-bound control message containing the
+  bound sessionId so the client can populate RunAgentRequest.
+"""
+
+from google_adk_extras import AdkBuilder
+
+
+app = (
+    AdkBuilder()
+    # Use either on-disk agents or programmatic agents. For streaming layer,
+    # the agent loader is optional if you stub the runner in tests; in real apps
+    # provide your agents via one of these methods:
+    # .with_agents_dir("./agents")
+
+    # Enable streaming layer with defaults
+    .with_web_ui(False)
+    .build_fastapi_app(enable_streaming=True)
+)
+

{google_adk_extras-0.2.5 → google_adk_extras-0.2.6}/mkdocs.yml

@@ -14,6 +14,7 @@ nav:
   - Getting Started: getting-started.md
   - Quickstarts: quickstarts.md
   - FastAPI Integration: fastapi.md
+  - Streaming: streaming.md
   - Services:
       - Durable Services: services.md
       - Credentials: credentials.md

{google_adk_extras-0.2.5 → google_adk_extras-0.2.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "google-adk-extras"
-version = "0.2.5"
+version = "0.2.6"
 description = "Production-ready services, credentials, and FastAPI wiring for Google ADK"
 readme = "README.md"
 requires-python = ">=3.10,<3.13"

{google_adk_extras-0.2.5 → google_adk_extras-0.2.6}/src/google_adk_extras/__init__.py

@@ -28,4 +28,4 @@ __all__ = [
     "CustomAgentLoader",
 ]
 
-__version__ = "0.2.5"
+__version__ = "0.2.6"

{google_adk_extras-0.2.5 → google_adk_extras-0.2.6}/src/google_adk_extras/enhanced_fastapi.py

@@ -5,6 +5,7 @@ that properly supports custom credential services.
 """
 
 import json
+import asyncio
 import logging
 import os
 from pathlib import Path
@@ -34,6 +35,7 @@ from google.adk.sessions.database_session_service import DatabaseSessionService
 from google.adk.utils.feature_decorator import working_in_progress
 from google.adk.cli.adk_web_server import AdkWebServer
 from .enhanced_adk_web_server import EnhancedAdkWebServer
+from .streaming import StreamingController, StreamingConfig
 from google.adk.cli.utils import envs
 from google.adk.cli.utils import evals
 from google.adk.cli.utils.agent_change_handler import AgentChangeEventHandler
@@ -64,6 +66,9 @@ def get_enhanced_fast_api_app(
     trace_to_cloud: bool = False,
     reload_agents: bool = False,
     lifespan: Optional[Lifespan[FastAPI]] = None,
+    # Streaming layer (optional)
+    enable_streaming: bool = False,
+    streaming_config: Optional[StreamingConfig] = None,
 ) -> FastAPI:
     """Enhanced version of Google ADK's get_fast_api_app with EnhancedRunner integration.
     
@@ -504,4 +509,96 @@ def get_enhanced_fast_api_app(
                     logger.error("Failed to setup programmatic A2A agent %s: %s", app_name, e)
 
     logger.info("Enhanced FastAPI app created with credential service support")
+
+    # Optional streaming mounts (SSE + WebSocket)
+    if enable_streaming:
+        cfg = streaming_config or StreamingConfig(enable_streaming=True)
+        controller = StreamingController(
+            config=cfg,
+            get_runner_async=adk_web_server.get_runner_async,
+            session_service=session_service,
+        )
+        app.state.streaming_controller = controller
+        @app.on_event("startup")
+        async def _start_streaming():  # pragma: no cover - lifecycle glue
+            controller.start()
+        @app.on_event("shutdown")
+        async def _stop_streaming():  # pragma: no cover - lifecycle glue
+            await controller.stop()
+
+        from fastapi import APIRouter, WebSocket, Query
+        from fastapi.responses import StreamingResponse
+        from google.adk.cli.adk_web_server import RunAgentRequest
+
+        router = APIRouter()
+        base = cfg.streaming_path_base.rstrip("/")
+
+        @router.get(f"{base}/events/{{channel_id}}")
+        async def stream_events(channel_id: str, appName: str = Query(...), userId: str = Query(...), sessionId: Optional[str] = Query(None)):
+            ch = await app.state.streaming_controller.open_or_bind_channel(
+                channel_id=channel_id, app_name=appName, user_id=userId, session_id=sessionId
+            )
+            q = app.state.streaming_controller.subscribe(channel_id, kind="sse")
+
+            async def gen():
+                try:
+                    # Announce channel binding with session id
+                    yield "event: channel-bound\n"
+                    yield f"data: {{\"appName\":\"{appName}\",\"userId\":\"{userId}\",\"sessionId\":\"{ch.session_id}\"}}\n\n"
+                    while True:
+                        payload = await q.get()
+                        yield f"data: {payload}\n\n"
+                except asyncio.CancelledError:
+                    pass
+                finally:
+                    app.state.streaming_controller.unsubscribe(channel_id, q)
+
+            return StreamingResponse(gen(), media_type="text/event-stream")
+
+        @router.post(f"{base}/send/{{channel_id}}")
+        async def send_message(channel_id: str, req: RunAgentRequest):
+            # Validation: channel binding must match
+            await app.state.streaming_controller.enqueue(channel_id, req)
+            return PlainTextResponse("", status_code=204)
+
+        @router.websocket(f"{base}/ws/{{channel_id}}")
+        async def ws_endpoint(websocket: WebSocket, channel_id: str, appName: str, userId: str, sessionId: Optional[str] = None):
+            await websocket.accept()
+            try:
+                await app.state.streaming_controller.open_or_bind_channel(
+                    channel_id=channel_id, app_name=appName, user_id=userId, session_id=sessionId
+                )
+                q = app.state.streaming_controller.subscribe(channel_id, kind="ws")
+                # Send channel binding info including session id
+                await websocket.send_text(json.dumps({"event": "channel-bound", "appName": appName, "userId": userId, "sessionId": app.state.streaming_controller._channels[channel_id].session_id}))
+
+                async def downlink():
+                    try:
+                        while True:
+                            payload = await q.get()
+                            await websocket.send_text(payload)
+                    except asyncio.CancelledError:
+                        pass
+
+                async def uplink():
+                    try:
+                        while True:
+                            text = await websocket.receive_text()
+                            # Strict type parity by default
+                            req = RunAgentRequest.model_validate_json(text)
+                            await app.state.streaming_controller.enqueue(channel_id, req)
+                    except Exception:
+                        return
+
+                down = asyncio.create_task(downlink())
+                up = asyncio.create_task(uplink())
+                await asyncio.wait({down, up}, return_when=asyncio.FIRST_COMPLETED)
+            finally:
+                try:
+                    app.state.streaming_controller.unsubscribe(channel_id, q)
+                except Exception:
+                    pass
+
+        app.include_router(router)
+
     return app

google_adk_extras-0.2.6/src/google_adk_extras/streaming/__init__.py

@@ -0,0 +1,12 @@
+"""Streaming support (SSE/WebSocket) for google-adk-extras.
+
+This package provides an optional, persistent bi-directional streaming layer
+with strict ADK type parity by default. It complements ADK's built-in
+`/run`, `/run_sse`, and `/run_live` endpoints by offering per-channel
+subscription and send semantics for chat-style UIs.
+"""
+
+from .streaming_controller import StreamingConfig, StreamingController
+
+__all__ = ["StreamingConfig", "StreamingController"]
+

google_adk_extras-0.2.6/src/google_adk_extras/streaming/streaming_controller.py

@@ -0,0 +1,262 @@
+import asyncio
+import time
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional, Set, Callable, Awaitable
+
+from fastapi import WebSocket, HTTPException
+from pydantic import BaseModel
+
+from google.adk.events.event import Event
+from google.adk.runners import Runner
+
+
+class StreamingConfig(BaseModel):
+    enable_streaming: bool = False
+    streaming_path_base: str = "/stream"
+    strict_types: bool = True
+    create_session_on_open: bool = True
+    ttl_seconds: int = 900
+    max_queue_size: int = 128
+    max_channels_per_user: int = 20
+    heartbeat_interval: Optional[float] = 20.0
+    reuse_session_policy: str = "per_channel"  # "per_channel" or "external"
+
+
+@dataclass
+class _Subscriber:
+    queue: "asyncio.Queue[str]"
+    kind: str  # "sse" | "ws"
+
+
+@dataclass
+class _Channel:
+    channel_id: str
+    app_name: str
+    user_id: str
+    session_id: str
+    in_q: "asyncio.Queue[Any]" = field(default_factory=asyncio.Queue)
+    subscribers: list[_Subscriber] = field(default_factory=list)
+    worker_task: Optional[asyncio.Task] = None
+    created_at: float = field(default_factory=lambda: time.time())
+    last_activity: float = field(default_factory=lambda: time.time())
+    lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+
+
+class StreamingController:
+    """Manages streaming channels and workers.
+
+    This controller binds a channel to (app_name, user_id, session_id) and
+    runs a background worker per channel to execute streamed runs and push
+    ADK Event JSON to all subscribers.
+    """
+
+    def __init__(
+        self,
+        *,
+        config: StreamingConfig,
+        get_runner_async: Callable[[str], Awaitable[Runner]],
+        session_service,
+    ) -> None:
+        self._config = config
+        self._get_runner_async = get_runner_async
+        self._session_service = session_service
+        self._channels: Dict[str, _Channel] = {}
+        self._gc_task: Optional[asyncio.Task] = None
+
+    def start(self) -> None:
+        if self._gc_task is None:
+            self._gc_task = asyncio.create_task(self._gc_loop())
+
+    async def stop(self) -> None:
+        if self._gc_task:
+            self._gc_task.cancel()
+            with asyncio.CancelledError:
+                pass
+            self._gc_task = None
+        # Cancel workers
+        for ch in list(self._channels.values()):
+            if ch.worker_task and not ch.worker_task.done():
+                ch.worker_task.cancel()
+        self._channels.clear()
+
+    def _ensure_user_limit(self, user_id: str) -> None:
+        if self._config.max_channels_per_user <= 0:
+            return
+        count = sum(1 for c in self._channels.values() if c.user_id == user_id)
+        if count >= self._config.max_channels_per_user:
+            raise HTTPException(status_code=429, detail="Too many channels for this user")
+
+    async def open_or_bind_channel(
+        self,
+        *,
+        channel_id: str,
+        app_name: str,
+        user_id: str,
+        session_id: Optional[str],
+    ) -> _Channel:
+        # Existing channel validation/match
+        if channel_id in self._channels:
+            ch = self._channels[channel_id]
+            if ch.app_name != app_name or ch.user_id != user_id:
+                raise HTTPException(status_code=409, detail="Channel binding conflict")
+            if session_id and session_id != ch.session_id:
+                raise HTTPException(status_code=409, detail="Channel already bound to different session")
+            ch.last_activity = time.time()
+            return ch
+
+        # New channel
+        self._ensure_user_limit(user_id)
+        if not session_id:
+            if not self._config.create_session_on_open:
+                raise HTTPException(status_code=400, detail="sessionId required for this channel")
+            # Create a fresh ADK session
+            create = getattr(self._session_service, "create_session", None)
+            if create is None:
+                # Older ADK interfaces may expose sync variant
+                create = getattr(self._session_service, "create_session_sync", None)
+            if create is None:
+                raise HTTPException(status_code=500, detail="Session service does not support create_session")
+            if asyncio.iscoroutinefunction(create):
+                session = await create(app_name=app_name, user_id=user_id)
+            else:
+                # Call sync and wrap
+                session = create(app_name=app_name, user_id=user_id)
+            session_id = session.id
+        else:
+            # Validate existing session
+            session = await self._session_service.get_session(app_name=app_name, user_id=user_id, session_id=session_id)
+            if not session:
+                raise HTTPException(status_code=404, detail="Session not found")
+
+        ch = _Channel(
+            channel_id=channel_id,
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            in_q=asyncio.Queue(),
+        )
+        self._channels[channel_id] = ch
+        ch.worker_task = asyncio.create_task(self._worker(ch))
+        return ch
+
+    def subscribe(self, channel_id: str, kind: str) -> asyncio.Queue[str]:
+        if channel_id not in self._channels:
+            raise HTTPException(status_code=404, detail="Channel not found")
+        q: asyncio.Queue[str] = asyncio.Queue(maxsize=self._config.max_queue_size)
+        self._channels[channel_id].subscribers.append(_Subscriber(queue=q, kind=kind))
+        self._channels[channel_id].last_activity = time.time()
+        return q
+
+    def unsubscribe(self, channel_id: str, q: asyncio.Queue[str]) -> None:
+        ch = self._channels.get(channel_id)
+        if not ch:
+            return
+        ch.subscribers = [s for s in ch.subscribers if s.queue is not q]
+        ch.last_activity = time.time()
+
+    async def enqueue(self, channel_id: str, req: Any) -> None:
+        ch = self._channels.get(channel_id)
+        if not ch:
+            raise HTTPException(status_code=404, detail="Channel not found")
+        # Validate binding
+        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:
+            raise HTTPException(status_code=409, detail="Request does not match channel binding")
+        await ch.in_q.put(req)
+        ch.last_activity = time.time()
+
+    async def _worker(self, ch: _Channel) -> None:
+        try:
+            while True:
+                req = await ch.in_q.get()
+                ch.last_activity = time.time()
+                try:
+                    runner = await self._get_runner_async(ch.app_name)
+                    # Stream events for this request
+                    async with _aclosing(
+                        runner.run_async(
+                            user_id=ch.user_id,
+                            session_id=ch.session_id,
+                            new_message=req.new_message,
+                            state_delta=getattr(req, "state_delta", None),
+                            run_config=_maybe_run_config_streaming(True),
+                        )
+                    ) as agen:
+                        async for event in agen:
+                            await self._broadcast_event(ch, event)
+                except Exception as e:  # pragma: no cover - safety
+                    await self._broadcast_error(ch, str(e))
+        except asyncio.CancelledError:  # worker shutdown
+            return
+
+    async def _broadcast_event(self, ch: _Channel, event: Event) -> None:
+        payload = event.model_dump_json(exclude_none=True, by_alias=True)
+        for sub in list(ch.subscribers):
+            try:
+                sub.queue.put_nowait(payload)
+            except asyncio.QueueFull:
+                # Drop subscriber on backpressure
+                ch.subscribers = [s for s in ch.subscribers if s is not sub]
+        ch.last_activity = time.time()
+
+    async def _broadcast_heartbeat(self, ch: _Channel) -> None:
+        if self._config.heartbeat_interval is None:
+            return
+        payload = '{"event":"heartbeat"}'
+        for sub in list(ch.subscribers):
+            try:
+                sub.queue.put_nowait(payload)
+            except asyncio.QueueFull:
+                ch.subscribers = [s for s in ch.subscribers if s is not sub]
+
+    async def _broadcast_error(self, ch: _Channel, message: str) -> None:
+        payload = '{"error": %s}' % _json_escape(message)
+        for sub in list(ch.subscribers):
+            try:
+                sub.queue.put_nowait(payload)
+            except asyncio.QueueFull:
+                ch.subscribers = [s for s in ch.subscribers if s is not sub]
+
+    async def _gc_loop(self) -> None:
+        try:
+            while True:
+                await asyncio.sleep(min(10, max(1, int(self._config.ttl_seconds / 3))))
+                now = time.time()
+                for cid, ch in list(self._channels.items()):
+                    idle = now - ch.last_activity
+                    if idle >= self._config.ttl_seconds and not ch.subscribers and ch.in_q.empty():
+                        if ch.worker_task and not ch.worker_task.done():
+                            ch.worker_task.cancel()
+                        self._channels.pop(cid, None)
+        except asyncio.CancelledError:
+            return
+
+
+# Utilities (avoid importing optional internals at module import time)
+def _maybe_run_config_streaming(enabled: bool):
+    # Support multiple ADK versions by resolving RunConfig/StreamingMode from
+    # either google.adk.runners or google.adk.agents.run_config
+    try:
+        from google.adk.runners import RunConfig  # type: ignore
+    except Exception:  # pragma: no cover - version fallback
+        from google.adk.agents.run_config import RunConfig  # type: ignore
+    try:
+        from google.adk.agents.run_config import StreamingMode  # type: ignore
+    except Exception:  # pragma: no cover - defensive
+        StreamingMode = type("StreamingMode", (), {"SSE": "sse", "NONE": None})  # minimal stub
+    return RunConfig(streaming_mode=StreamingMode.SSE if enabled else StreamingMode.NONE)
+
+
+class _aclosing:
+    def __init__(self, agen):
+        self._agen = agen
+    async def __aenter__(self):
+        return self._agen
+    async def __aexit__(self, exc_type, exc, tb):
+        try:
+            await self._agen.aclose()
+        except Exception:
+            pass
+
+
+def _json_escape(s: str) -> str:
+    return '"' + s.replace('\\', '\\\\').replace('"', '\\"') + '"'

{google_adk_extras-0.2.5 → google_adk_extras-0.2.6/src/google_adk_extras.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-adk-extras
-Version: 0.2.5
+Version: 0.2.6
 Summary: Production-ready services, credentials, and FastAPI wiring for Google ADK
 Home-page: https://github.com/DeadMeme5441/google-adk-extras
 Author: DeadMeme5441

{google_adk_extras-0.2.5 → google_adk_extras-0.2.6}/src/google_adk_extras.egg-info/SOURCES.txt

@@ -12,6 +12,7 @@ docs/getting-started.md
 docs/index.md
 docs/quickstarts.md
 docs/services.md
+docs/streaming.md
 docs/troubleshooting.md
 docs/uris.md
 examples/README.md
@@ -20,6 +21,7 @@ examples/custom_loader.py
 examples/fastapi_app.py
 examples/programmatic_a2a_expose.py
 examples/runner_basic.py
+examples/streaming_sse_ws.py
 examples/credentials/google_oauth2.py
 examples/services/artifacts_local.py
 examples/services/memory_yaml.py
@@ -61,4 +63,6 @@ src/google_adk_extras/sessions/mongo_session_service.py
 src/google_adk_extras/sessions/redis_session_service.py
 src/google_adk_extras/sessions/sql_session_service.py
 src/google_adk_extras/sessions/yaml_file_session_service.py
+src/google_adk_extras/streaming/__init__.py
+src/google_adk_extras/streaming/streaming_controller.py
 tests/test_a2a_helpers.py