google-adk-extras 0.2.5__py3-none-any.whl → 0.2.6__py3-none-any.whl

google_adk_extras/__init__.py

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

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/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/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.dist-info → google_adk_extras-0.2.6.dist-info}/METADATA

@@ -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.dist-info → google_adk_extras-0.2.6.dist-info}/RECORD

@@ -1,8 +1,8 @@
-google_adk_extras/__init__.py,sha256=D0P8SjXVhABGaVYLQmkP9LlXRMLlsRfyhWZ4XsiWaOs,830
+google_adk_extras/__init__.py,sha256=cIaqZH7E__SmojqMjCqOA4FV6S6-k1pH9_mGCfkCfU4,830
 google_adk_extras/adk_builder.py,sha256=anTpd-UYtPugRSqdYNmWa_uesVjHGEKO2mntkcU-J6g,41179
 google_adk_extras/custom_agent_loader.py,sha256=6rQyBTAxvuwFPs3QGt8hWYZKsQVpbI4WIy6OS73tsac,6001
 google_adk_extras/enhanced_adk_web_server.py,sha256=rML_m4Um9QCcnnlHvCOyLGjLHDcyladEunX-x4JX43Q,5414
-google_adk_extras/enhanced_fastapi.py,sha256=qMAawDdduVp3DMZtFJq1f1l6tD8o5_lyHf3Vf4k5JyM,23790
+google_adk_extras/enhanced_fastapi.py,sha256=TzIr159B0tnEqj2dXKmV_tJvkA66Fq2I-kL5cz9VGSk,28304
 google_adk_extras/enhanced_runner.py,sha256=b7O1a9-4S49LduILOEDs6IxjCI4w_E39sc-Hs4y3Rys,1410
 google_adk_extras/artifacts/__init__.py,sha256=_IsKDgf6wanWR0HXvSpK9SiLa3n5URKLtazkKyH1P-o,931
 google_adk_extras/artifacts/base_custom_artifact_service.py,sha256=O9rkc250B3yDRYbyDI0EvTrCKvnih5_DQas5OF-hRMY,9721
@@ -30,8 +30,10 @@ google_adk_extras/sessions/mongo_session_service.py,sha256=r3jZ3PmDpbZ0veNzXzuAj
 google_adk_extras/sessions/redis_session_service.py,sha256=yyfXZozeFWJ2S_kz7zXqz--f_ymE6HMMpT3MhcpFXIE,10434
 google_adk_extras/sessions/sql_session_service.py,sha256=TaOeEVWnwQ_8nvDZBW7e3qhzR_ecuGsjvZ_kh6Guq8g,14558
 google_adk_extras/sessions/yaml_file_session_service.py,sha256=SpTh8YHIALcoxzmturhcZ4ReHKQrJI1CxoYiJQ-baRc,11819
-google_adk_extras-0.2.5.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-google_adk_extras-0.2.5.dist-info/METADATA,sha256=QmkrJkhlsJQb09By4L-NlRFxfU9P_WfjJJIoo1YwCq8,10707
-google_adk_extras-0.2.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-google_adk_extras-0.2.5.dist-info/top_level.txt,sha256=DDWgVkz8G8ihPzznxAWyKa2jgJW3F6Fwy__qMddoKTs,18
-google_adk_extras-0.2.5.dist-info/RECORD,,
+google_adk_extras/streaming/__init__.py,sha256=rcjmlCJHTlvUiCrx6qNGw5ObCnEtfENkGTvzfEiGL0M,461
+google_adk_extras/streaming/streaming_controller.py,sha256=Qxg4yqakd3mBPJ9sISe4gO_pvBpYttvSq5SIK4W5ZPQ,10505
+google_adk_extras-0.2.6.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+google_adk_extras-0.2.6.dist-info/METADATA,sha256=NXbByeiIsdHl2UcZ5F-i22thmm-rFcTr1poTFdVY5cI,10707
+google_adk_extras-0.2.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+google_adk_extras-0.2.6.dist-info/top_level.txt,sha256=DDWgVkz8G8ihPzznxAWyKa2jgJW3F6Fwy__qMddoKTs,18
+google_adk_extras-0.2.6.dist-info/RECORD,,