chatsee-google-adk 0.1.0__tar.gz

chatsee_google_adk-0.1.0/PKG-INFO

@@ -0,0 +1,97 @@
+Metadata-Version: 2.4
+Name: chatsee-google-adk
+Version: 0.1.0
+Summary: ChatSee connector for Google ADK — capture every agent run with one plugin.
+Maintainer-email: ChatSee <contact@chatsee.ai>
+License: MIT
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: chatsee-ai>=0.14.0
+Requires-Dist: google-adk>=2.4
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0; extra == "dev"
+Requires-Dist: chatsee-connector-testing; extra == "dev"
+
+# chatsee-google-adk
+
+ChatSee connector for Google ADK (Agent Development Kit).
+
+```bash
+pip install chatsee-google-adk
+```
+
+```python
+from google.adk.apps import App
+from chatsee_google_adk import instrument_google_adk
+
+chatsee = instrument_google_adk(agent_id="...", tenant_id="...")
+app = App(name="support", root_agent=agent, plugins=[chatsee.plugin])
+```
+
+That is the whole integration. The handle's `.plugin` is an ADK `BasePlugin`,
+and it works the same on a `Runner(plugins=[...])`. `chatsee-ai` is installed
+for you as a dependency.
+
+## What it captures
+
+Per turn: the user's message, the agent's answer, the root agent's
+instruction as the system prompt, every tool call with its arguments, result
+and duration, and **one model call per model call** with the usage ADK
+reported for it.
+
+## No per-agent wiring
+
+ADK fires plugin callbacks globally — every agent, every sub-agent, every
+model call and tool in the run goes past this one plugin. A multi-agent app
+needs no extra registration, and adding a sub-agent later changes nothing.
+
+Every callback returns `None`, which is ADK's "carry on" — the plugin only
+reads what goes past and never alters a response, a tool result or the flow.
+
+## The part worth knowing
+
+**One ADK run is one ChatSee turn**, emitted from `after_run_callback` — one
+`runner.run_async` call, one user message. Not `after_agent_callback`, which
+fires once per sub-agent and would split a single turn into several.
+
+Two consequences worth knowing:
+
+* **A failed run still produces a turn.** ADK skips `after_run` when a run
+  raises, so the connector also emits from `on_run_error_callback`, with the
+  error attached. Otherwise a failing agent would simply go quiet.
+* **Streaming does not inflate the model-call count.** `after_model` fires
+  once per chunk; partial responses are ignored and only the final aggregated
+  response counts, so one call stays one call rather than dozens.
+
+## Sessions and users
+
+By default each turn uses ADK's own `session.id` and `user_id`, so a reused
+ADK session is one ChatSee conversation — usually what you want.
+
+To override, pass `user_id=` / `conversation_id=` to `instrument_google_adk`,
+or call `chatsee.set_session(conversation_id)` to pin every later turn to one
+conversation. `chatsee.stats()` reports turns emitted, capture failures and
+runs in flight.
+
+## Guarantees
+
+* **Fails open.** Every callback is wrapped; a failure inside capture is
+  logged and swallowed. ADK re-raises plugin exceptions into the run, so this
+  matters more here than elsewhere — ChatSee can never break your agent.
+* **Idempotent.** Instrumenting twice with the same identity returns the first
+  handle rather than registering a second plugin.
+* **Concurrency-safe.** Runs are keyed by `invocation_id`, so one Runner
+  serving several users at once does not mix them up, and the send happens off
+  the event loop so one turn's POST does not stall the others.
+
+All three are enforced by the contract suite, not just intended.
+
+## Tests
+
+No network, no API key, no dev tenant — ADK is stubbed, ChatSee is not.
+
+```bash
+pip install -e python/connectors/_shared
+pip install -e python/connectors/frameworks/google-adk
+pytest python/connectors/frameworks/google-adk/tests -q
+```

chatsee_google_adk-0.1.0/README.md

@@ -0,0 +1,83 @@
+# chatsee-google-adk
+
+ChatSee connector for Google ADK (Agent Development Kit).
+
+```bash
+pip install chatsee-google-adk
+```
+
+```python
+from google.adk.apps import App
+from chatsee_google_adk import instrument_google_adk
+
+chatsee = instrument_google_adk(agent_id="...", tenant_id="...")
+app = App(name="support", root_agent=agent, plugins=[chatsee.plugin])
+```
+
+That is the whole integration. The handle's `.plugin` is an ADK `BasePlugin`,
+and it works the same on a `Runner(plugins=[...])`. `chatsee-ai` is installed
+for you as a dependency.
+
+## What it captures
+
+Per turn: the user's message, the agent's answer, the root agent's
+instruction as the system prompt, every tool call with its arguments, result
+and duration, and **one model call per model call** with the usage ADK
+reported for it.
+
+## No per-agent wiring
+
+ADK fires plugin callbacks globally — every agent, every sub-agent, every
+model call and tool in the run goes past this one plugin. A multi-agent app
+needs no extra registration, and adding a sub-agent later changes nothing.
+
+Every callback returns `None`, which is ADK's "carry on" — the plugin only
+reads what goes past and never alters a response, a tool result or the flow.
+
+## The part worth knowing
+
+**One ADK run is one ChatSee turn**, emitted from `after_run_callback` — one
+`runner.run_async` call, one user message. Not `after_agent_callback`, which
+fires once per sub-agent and would split a single turn into several.
+
+Two consequences worth knowing:
+
+* **A failed run still produces a turn.** ADK skips `after_run` when a run
+  raises, so the connector also emits from `on_run_error_callback`, with the
+  error attached. Otherwise a failing agent would simply go quiet.
+* **Streaming does not inflate the model-call count.** `after_model` fires
+  once per chunk; partial responses are ignored and only the final aggregated
+  response counts, so one call stays one call rather than dozens.
+
+## Sessions and users
+
+By default each turn uses ADK's own `session.id` and `user_id`, so a reused
+ADK session is one ChatSee conversation — usually what you want.
+
+To override, pass `user_id=` / `conversation_id=` to `instrument_google_adk`,
+or call `chatsee.set_session(conversation_id)` to pin every later turn to one
+conversation. `chatsee.stats()` reports turns emitted, capture failures and
+runs in flight.
+
+## Guarantees
+
+* **Fails open.** Every callback is wrapped; a failure inside capture is
+  logged and swallowed. ADK re-raises plugin exceptions into the run, so this
+  matters more here than elsewhere — ChatSee can never break your agent.
+* **Idempotent.** Instrumenting twice with the same identity returns the first
+  handle rather than registering a second plugin.
+* **Concurrency-safe.** Runs are keyed by `invocation_id`, so one Runner
+  serving several users at once does not mix them up, and the send happens off
+  the event loop so one turn's POST does not stall the others.
+
+All three are enforced by the contract suite, not just intended.
+
+## Tests
+
+No network, no API key, no dev tenant — ADK is stubbed, ChatSee is not.
+
+```bash
+pip install -e python/connectors/_shared
+pip install -e python/connectors/frameworks/google-adk
+pytest python/connectors/frameworks/google-adk/tests -q
+```

chatsee_google_adk-0.1.0/chatsee_google_adk/__init__.py

@@ -0,0 +1,387 @@
+"""
+ChatSee connector for Google ADK (Agent Development Kit).
+
+One line in the customer's startup, one argument on their App or Runner:
+
+    from chatsee_google_adk import instrument_google_adk
+
+    chatsee = instrument_google_adk(agent_id="...", tenant_id="...")
+    app = App(name="support", root_agent=agent, plugins=[chatsee.plugin])
+
+The handle's ``.plugin`` is an ADK ``BasePlugin``. ADK fires plugin callbacks
+globally -- every agent, sub-agent, model call and tool in the run -- so there
+is no per-agent wiring. The plugin returns None from every callback, so ADK's
+own behaviour is never altered; it only reads what goes past.
+
+Turn boundary: one ADK run (one ``runner.run_async`` call, one user message)
+is one ChatSee turn, emitted from ``after_run_callback``. Not
+``after_agent_callback``, which fires once per sub-agent.
+"""
+from __future__ import annotations
+
+import asyncio
+import logging
+import threading
+import time
+import weakref
+from typing import Any, Dict, Optional
+
+from ._capture import (
+    FRAMEWORK,
+    ModelCall,
+    Run,
+    RunStore,
+    ToolCall,
+    content_text,
+    elapsed_ms,
+    safe_json,
+    split_model,
+    system_instruction,
+)
+
+__version__ = "0.1.0"
+__all__ = ["instrument_google_adk", "ChatseeGoogleADK", "ChatseeADKPlugin", "__version__"]
+
+logger = logging.getLogger(__name__)
+
+try:
+    from google.adk.plugins.base_plugin import BasePlugin
+except ImportError:  # pragma: no cover - google-adk is a declared dependency
+    # Only so the offline contract tests can run without the framework. A
+    # customer always has google-adk; the Runner needs a real BasePlugin.
+    class BasePlugin:  # type: ignore[no-redef]
+        def __init__(self, name: str) -> None:
+            self.name = name
+
+
+PLUGIN_NAME = "chatsee"
+
+
+def _framework_version() -> str:
+    try:
+        from google.adk import version
+
+        return getattr(version, "__version__", "unknown")
+    except Exception:
+        return "unknown"
+
+
+def _invocation_id(ctx: Any) -> Optional[str]:
+    """CallbackContext / ToolContext / InvocationContext all expose it."""
+    return getattr(ctx, "invocation_id", None)
+
+
+def _call_key(ctx: Any) -> str:
+    """Which agent made this model call, so its start and end pair up."""
+    return str(getattr(ctx, "agent_name", None) or "")
+
+
+class ChatseeADKPlugin(BasePlugin):
+    """The ADK plugin. Pass it to ``App(plugins=[...])`` or ``Runner(plugins=[...])``.
+
+    Every callback is wrapped so a failure inside ChatSee is logged and
+    swallowed. ADK re-raises plugin exceptions into the run, so an unguarded
+    bug here would break the customer's agent.
+    """
+
+    def __init__(self, handle: "ChatseeGoogleADK") -> None:
+        super().__init__(name=PLUGIN_NAME)
+        self._handle = handle
+
+    # -- run lifecycle -----------------------------------------------------
+
+    async def on_user_message_callback(self, *, invocation_context: Any,
+                                       user_message: Any) -> None:
+        try:
+            invocation_id = _invocation_id(invocation_context)
+            text = content_text(user_message)
+            if not invocation_id or not text:
+                # Nothing identifies this run, or nothing was asked. Recording
+                # it would attach its calls to no turn.
+                return None
+            session = getattr(invocation_context, "session", None)
+            agent = getattr(invocation_context, "agent", None)
+            self._handle._store.open(invocation_id, Run(
+                user_message=text,
+                session_id=getattr(session, "id", None),
+                user_id=getattr(invocation_context, "user_id", None),
+                bot_id=getattr(agent, "name", None),
+            ))
+        except Exception as exc:  # noqa: BLE001 - ChatSee must never break the agent
+            self._handle._capture_failed("on_user_message", exc)
+        return None
+
+    async def after_run_callback(self, *, invocation_context: Any) -> None:
+        try:
+            run = self._handle._store.take(_invocation_id(invocation_context))
+            if run is not None:
+                await self._handle._emit(run)
+        except Exception as exc:  # noqa: BLE001
+            self._handle._capture_failed("after_run", exc)
+        return None
+
+    async def on_run_error_callback(self, *, invocation_context: Any,
+                                    error: Exception) -> None:
+        # ADK skips after_run when the run raises. Emit the turn here instead,
+        # with the error, so a failed turn is visible rather than missing.
+        try:
+            run = self._handle._store.take(_invocation_id(invocation_context))
+            if run is not None:
+                await self._handle._emit(run, exception=error)
+        except Exception as exc:  # noqa: BLE001
+            self._handle._capture_failed("on_run_error", exc)
+        return None
+
+    # -- model calls -------------------------------------------------------
+
+    async def before_model_callback(self, *, callback_context: Any,
+                                    llm_request: Any) -> None:
+        try:
+            run = self._handle._store.get(_invocation_id(callback_context))
+            if run is not None:
+                if run.system_prompt is None:
+                    # The first model call is the root agent's; its
+                    # instruction is the turn's system prompt.
+                    run.system_prompt = system_instruction(llm_request)
+                key = _call_key(callback_context)
+                run.model_started[key] = time.monotonic()
+                run.model_requested[key] = getattr(llm_request, "model", None)
+        except Exception as exc:  # noqa: BLE001
+            self._handle._capture_failed("before_model", exc)
+        return None
+
+    async def after_model_callback(self, *, callback_context: Any,
+                                   llm_response: Any) -> None:
+        try:
+            run = self._handle._store.get(_invocation_id(callback_context))
+            if run is None:
+                return None
+            if getattr(llm_response, "partial", None):
+                # Streaming fires this once per chunk. Only the final,
+                # aggregated response is a model call; counting chunks would
+                # turn one call into dozens.
+                return None
+            key = _call_key(callback_context)
+            requested = run.model_requested.pop(key, None)
+            _, provider = split_model(requested)
+            model, version_provider = split_model(
+                getattr(llm_response, "model_version", None) or requested)
+            # One log_model_call per model call. A tool turn fires this twice
+            # (the tool request, then the answer); do NOT sum them -- the SDK
+            # sums, and the per-call grain keeps spend attributable.
+            run.model_calls.append(ModelCall(
+                model=model,
+                provider=provider or version_provider,
+                usage=getattr(llm_response, "usage_metadata", None),
+                duration_ms=elapsed_ms(run.model_started.pop(key, None)),
+            ))
+            text = content_text(getattr(llm_response, "content", None))
+            if text:
+                run.bot_message = text
+        except Exception as exc:  # noqa: BLE001
+            self._handle._capture_failed("after_model", exc)
+        return None
+
+    async def on_model_error_callback(self, *, callback_context: Any,
+                                      llm_request: Any, error: Exception) -> None:
+        try:
+            run = self._handle._store.get(_invocation_id(callback_context))
+            if run is not None:
+                key = _call_key(callback_context)
+                run.model_started.pop(key, None)
+                run.model_requested.pop(key, None)
+                run.exception = str(error)
+        except Exception as exc:  # noqa: BLE001
+            self._handle._capture_failed("on_model_error", exc)
+        return None
+
+    # -- tool calls --------------------------------------------------------
+
+    async def before_tool_callback(self, *, tool: Any, tool_args: Dict[str, Any],
+                                   tool_context: Any) -> None:
+        try:
+            run = self._handle._store.get(_invocation_id(tool_context))
+            if run is not None:
+                run.tool_started[_tool_key(tool, tool_context)] = time.monotonic()
+        except Exception as exc:  # noqa: BLE001
+            self._handle._capture_failed("before_tool", exc)
+        return None
+
+    async def after_tool_callback(self, *, tool: Any, tool_args: Dict[str, Any],
+                                  tool_context: Any, result: Any) -> None:
+        self._record_tool(tool, tool_args, tool_context, result=result)
+        return None
+
+    async def on_tool_error_callback(self, *, tool: Any, tool_args: Dict[str, Any],
+                                     tool_context: Any, error: Exception) -> None:
+        self._record_tool(tool, tool_args, tool_context, error=str(error) or type(error).__name__)
+        return None
+
+    def _record_tool(self, tool: Any, tool_args: Any, tool_context: Any,
+                     result: Any = None, error: Optional[str] = None) -> None:
+        try:
+            run = self._handle._store.get(_invocation_id(tool_context))
+            if run is None:
+                return
+            run.tool_calls.append(ToolCall(
+                tool_name=getattr(tool, "name", None) or "unknown",
+                arguments=safe_json(tool_args) or {},
+                result=safe_json(result),
+                error=error,
+                duration_ms=elapsed_ms(run.tool_started.pop(_tool_key(tool, tool_context), None)),
+            ))
+        except Exception as exc:  # noqa: BLE001
+            self._handle._capture_failed("tool", exc)
+
+
+def _tool_key(tool: Any, tool_context: Any) -> str:
+    """Pairs a tool's before and after. Parallel calls differ by function_call_id."""
+    return "%s:%s" % (getattr(tool, "name", None),
+                      getattr(tool_context, "function_call_id", None))
+
+
+class ChatseeGoogleADK:
+    """Handle returned by :func:`instrument_google_adk`.
+
+    ``.plugin`` is the only thing a customer needs. ``.set_session(id)`` pins
+    every turn to one ChatSee conversation; ``.stats()`` reports what the
+    connector emitted.
+    """
+
+    def __init__(self, tracker: Any, user_id: Optional[str] = None,
+                 conversation_id: Optional[str] = None) -> None:
+        self._tracker = tracker
+        self._user_id = user_id
+        self._conversation_id = conversation_id
+        self._store = RunStore()
+        # The tracker holds one turn at a time; a run's start..end replay must
+        # not interleave with another run's.
+        self._emit_lock = threading.Lock()
+        self._emitted = 0
+        self._failures = 0
+        self.plugin = ChatseeADKPlugin(self)
+
+    # -- public ------------------------------------------------------------
+
+    def set_session(self, conversation_id: str) -> None:
+        """Send every later turn under this conversation id.
+
+        Without it, each turn uses ADK's own ``session.id``, so a reused ADK
+        session is one ChatSee conversation -- usually what you want.
+        """
+        self._conversation_id = conversation_id
+
+    def stats(self) -> Dict[str, int]:
+        return {
+            "turns_emitted": self._emitted,
+            "capture_failures": self._failures,
+            "in_flight": self._store.in_flight,
+        }
+
+    # -- internals ---------------------------------------------------------
+
+    def _metadata(self, run: Run) -> Dict[str, Any]:
+        metadata: Dict[str, Any] = {
+            "framework": FRAMEWORK,
+            "framework_version": _framework_version(),
+            "connector_version": __version__,
+        }
+        if run.bot_id:
+            # The agent's name in the customer's code -- never the ChatSee agent_id.
+            metadata["bot_id"] = run.bot_id
+        return metadata
+
+    async def _emit(self, run: Run, exception: Optional[BaseException] = None) -> None:
+        # end_turn POSTs synchronously. Off the event loop, so one run's send
+        # does not stall every other run this Runner is serving.
+        await asyncio.to_thread(self._replay, run, exception)
+
+    def _replay(self, run: Run, exception: Optional[BaseException] = None) -> None:
+        """Hand one finished run to the tracker as start_turn .. end_turn."""
+        try:
+            with self._emit_lock:
+                tracker = self._tracker
+                session_id = self._conversation_id or run.session_id
+                if session_id:
+                    tracker.session_id = session_id
+
+                tracker.start_turn(run.user_message,
+                                   metadata=self._metadata(run),
+                                   system_prompt=run.system_prompt,
+                                   user_id=self._user_id or run.user_id)
+                for call in run.tool_calls:
+                    tracker.log_tool_call(call.tool_name, call.arguments,
+                                          result=call.result, error=call.error,
+                                          duration_ms=call.duration_ms)
+                for call in run.model_calls:
+                    tracker.log_model_call(model=call.model, provider=call.provider,
+                                           usage=call.usage, duration_ms=call.duration_ms)
+                error = exception if exception is not None else run.exception
+                if error is not None:
+                    tracker.log_exception(error)
+                tracker.end_turn(run.bot_message)
+                self._emitted += 1
+        except Exception as exc:  # noqa: BLE001
+            self._capture_failed("emit", exc)
+
+    def _capture_failed(self, where: str, exc: BaseException) -> None:
+        self._failures += 1
+        logger.error("chatsee-google-adk capture failed in %s: %s", where, exc, exc_info=True)
+
+
+# One handle per identity, so instrumenting twice is a no-op rather than a
+# second plugin capturing every turn again.
+_HANDLES: "weakref.WeakValueDictionary[Any, ChatseeGoogleADK]" = weakref.WeakValueDictionary()
+_HANDLES_LOCK = threading.Lock()
+
+
+def instrument_google_adk(
+    agent_id: Optional[str] = None,
+    tenant_id: Optional[str] = None,
+    user_id: Optional[str] = None,
+    conversation_id: Optional[str] = None,
+    api_base_url: Optional[str] = None,
+    _tracker: Any = None,
+) -> ChatseeGoogleADK:
+    """Capture every ADK run as one ChatSee turn.
+
+    ``agent_id`` and ``tenant_id`` say which ChatSee agent the data belongs to,
+    and are all that is required. Register the returned handle's ``.plugin``
+    on your ``App`` (or ``Runner``).
+
+    ``user_id`` overrides ADK's own ``user_id`` for every turn;
+    ``conversation_id`` overrides ADK's ``session.id``. Leave both unset and
+    the connector uses what ADK already knows.
+
+    ``_tracker`` is private and exists so the contract tests can inject a
+    fake. Customers never pass it.
+
+    Calling this twice with the same arguments is a no-op and returns the
+    first handle.
+    """
+    if _tracker is not None:
+        key: Any = ("tracker", id(_tracker))
+    else:
+        if not agent_id or not tenant_id:
+            raise ValueError("agent_id and tenant_id are required.")
+        key = ("identity", agent_id, tenant_id, user_id, conversation_id, api_base_url)
+
+    with _HANDLES_LOCK:
+        handle = _HANDLES.get(key)
+        if handle is not None:
+            return handle
+
+        if _tracker is None:
+            from chatsee import ChatseeTracker
+
+            _tracker = ChatseeTracker(
+                agent_id=agent_id,
+                tenant_id=tenant_id,
+                user_id=user_id,
+                session_id=conversation_id,
+                api_base_url=api_base_url,
+            )
+
+        handle = ChatseeGoogleADK(_tracker, user_id=user_id, conversation_id=conversation_id)
+        _HANDLES[key] = handle
+        return handle