From 4716801d913986a29944df817a9821c80e4d3f2b Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Wed, 23 Sep 2026 13:52:50 -0800 Subject: [PATCH 1/2] feat(AIC-3412): add the Google ADK agents handler ADK is a wildcard agent runtime, so the handler serves any provider. Gemini is the default transport and Vertex is an explicit opt-in. Co-authored-by: Cursor --- .env.example | 8 + .release-please-manifest.json | 3 +- .sdk_metadata.json | 13 + AGENTS.md | 5 +- examples/google_adk_agents_example.py | 41 ++ main.py | 2 + packages/google-adk-agents/README.md | 17 + packages/google-adk-agents/agents.md | 20 + packages/google-adk-agents/pyproject.toml | 36 + .../__init__.py | 26 + .../graph.py | 40 + .../handler.py | 694 ++++++++++++++++++ .../native_graph.py | 289 ++++++++ .../spans.py | 184 +++++ .../google-adk-agents/tests/test_graph.py | 60 ++ .../google-adk-agents/tests/test_handler.py | 635 ++++++++++++++++ .../tests/test_native_graph.py | 268 +++++++ .../google-adk-agents/tests/test_spans.py | 98 +++ pyproject.toml | 9 + release-please-config.json | 11 + tests/test_cross_handler_parity.py | 8 +- uv.lock | 32 +- 22 files changed, 2486 insertions(+), 13 deletions(-) create mode 100644 examples/google_adk_agents_example.py create mode 100644 packages/google-adk-agents/README.md create mode 100644 packages/google-adk-agents/agents.md create mode 100644 packages/google-adk-agents/pyproject.toml create mode 100644 packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/__init__.py create mode 100644 packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/graph.py create mode 100644 packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/handler.py create mode 100644 packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/native_graph.py create mode 100644 packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/spans.py create mode 100644 packages/google-adk-agents/tests/test_graph.py create mode 100644 packages/google-adk-agents/tests/test_handler.py create mode 100644 packages/google-adk-agents/tests/test_native_graph.py create mode 100644 packages/google-adk-agents/tests/test_spans.py diff --git a/.env.example b/.env.example index 5bb04800..fb7be1bb 100644 --- a/.env.example +++ b/.env.example @@ -5,6 +5,14 @@ LD_SDK_KEY=your-sdk-key-here OPENAI_API_KEY=your-openai-key-here ANTHROPIC_API_KEY=your-anthropic-key-here +# Google ADK. Gemini Developer API is the default; set one of these. +GOOGLE_API_KEY=your-gemini-key-here +# GEMINI_API_KEY= +# GOOGLE_GENAI_API_KEY= +# Vertex is opt-in in the handler (use_vertexai=True). Do not set these for the Gemini path. +# GOOGLE_CLOUD_PROJECT= +# GOOGLE_CLOUD_LOCATION= + # LaunchDarkly endpoints (leave blank to use production defaults) LD_BASE_URI= LD_STREAM_URI= diff --git a/.release-please-manifest.json b/.release-please-manifest.json index 15fdf658..77bb919d 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -6,5 +6,6 @@ "packages/openai-agents": "0.2.3", "packages/openai-messages": "0.2.3", "packages/langchain-agents": "0.2.3", - "packages/langchain-messages": "0.2.3" + "packages/langchain-messages": "0.2.3", + "packages/google-adk-agents": "0.1.0" } diff --git a/.sdk_metadata.json b/.sdk_metadata.json index 5d7708f3..f194de0f 100644 --- a/.sdk_metadata.json +++ b/.sdk_metadata.json @@ -102,6 +102,19 @@ "aiAgentGraph": { "introduced": "0.1" }, "aiTrackMetrics": { "introduced": "0.1" } } + }, + "python-ai-google-adk-agents": { + "name": "Google ADK Agents AI SDK for Python", + "type": "ai-provider", + "path": "packages/google-adk-agents", + "languages": ["Python"], + "aiSdkNames": ["launchdarkly-ai-google-adk-agents"], + "releases": { "tag-prefix": "launchdarkly-ai-google-adk-agents-" }, + "features": { + "aiAgentConfig": { "introduced": "0.1" }, + "aiAgentGraph": { "introduced": "0.1" }, + "aiTrackMetrics": { "introduced": "0.1" } + } } } } diff --git a/AGENTS.md b/AGENTS.md index 305613a9..f4dfe024 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -87,6 +87,7 @@ graph TD claude["launchdarkly-ai-claude-agents"] openai["launchdarkly-ai-openai-agents"] langchain["launchdarkly-ai-langchain-agents"] + googleAdk["launchdarkly-ai-google-adk-agents"] newHandler["launchdarkly-ai-new-provider\n(future)"] end subgraph tier0 ["Tier 0 — Core"] @@ -97,11 +98,13 @@ graph TD app --> claude app --> openai app --> langchain + app --> googleAdk app --> newHandler app --> ai claude --> client openai --> client langchain --> client + googleAdk --> client newHandler --> client ai --> client ``` @@ -110,7 +113,7 @@ graph TD - **Tier 0 — Core** (`launchdarkly-ai-server`): The foundation. Owns all LaunchDarkly integration, telemetry orchestration, shared data types, and the primary entry points (`config()`, `graph()`, `resolve_graph()`). Has no dependency on any other `launchdarkly-ai-*` package. - **Tier 0 — Convenience barrel** (`launchdarkly-ai-python`): A pure re-export package that makes all of `launchdarkly-ai-server` available under a shorter install name. No new logic — intended as the default install for most Python applications. -- **Tier 1 — Handler packages** (`launchdarkly-ai-claude-agents`, `launchdarkly-ai-claude-messages`, `launchdarkly-ai-openai-agents`, `launchdarkly-ai-openai-messages`, `launchdarkly-ai-langchain-agents`, `launchdarkly-ai-langchain-messages`, …): Each wraps a specific AI provider SDK. Depends on `launchdarkly-ai-server` for shared types and utilities. Must not depend on other Tier 1 packages. +- **Tier 1 — Handler packages** (`launchdarkly-ai-claude-agents`, `launchdarkly-ai-claude-messages`, `launchdarkly-ai-openai-agents`, `launchdarkly-ai-openai-messages`, `launchdarkly-ai-langchain-agents`, `launchdarkly-ai-langchain-messages`, `launchdarkly-ai-google-adk-agents`, …): Each wraps a specific AI provider SDK. Depends on `launchdarkly-ai-server` for shared types and utilities. Must not depend on other Tier 1 packages. - **Tier 2 — Consumer applications** (e.g. `main.py`, downstream projects): Imports from one or more handler packages and either `launchdarkly-ai-python` or `launchdarkly-ai-server`. Owns tool implementations and orchestration logic. No `launchdarkly-ai-*` package should ever depend on Tier 2 code. ### Rules diff --git a/examples/google_adk_agents_example.py b/examples/google_adk_agents_example.py new file mode 100644 index 00000000..f2d6097e --- /dev/null +++ b/examples/google_adk_agents_example.py @@ -0,0 +1,41 @@ +""" +Example: google_adk_agents() — Google ADK agents handler (wildcard provider). + +Gemini Developer API is the default (`GOOGLE_API_KEY`, `GEMINI_API_KEY`, or +`GOOGLE_GENAI_API_KEY`). Vertex is opt-in: pass `use_vertexai=True` plus project +and location, or set `GOOGLE_CLOUD_PROJECT` and `GOOGLE_CLOUD_LOCATION`. + +Usage (via main.py): + python main.py google-adk-agents "" +""" + +from __future__ import annotations + +import json + +from examples.tools import ( + fetch_launchdarkly_documentation, + get_preferences, + search_ld_documentation, + web_search, +) +from examples.utils import new_context, write_output +from launchdarkly_ai_google_adk_agents import google_adk_agents + + +async def run(key: str, user_input: str) -> None: + response = await google_adk_agents( + key, + user_input, + new_context(), + tool_handlers={ + "get-user-preferences": get_preferences, + "search-ld-documentation": search_ld_documentation, + "fetch-ld-documentation": fetch_launchdarkly_documentation, + "fetch-launchdarkly-documentation": fetch_launchdarkly_documentation, + "web-search": web_search, + }, + ) + + print(json.dumps(response, indent=2, default=str)) + write_output(response) diff --git a/main.py b/main.py index fc5c1e5b..d01fe711 100644 --- a/main.py +++ b/main.py @@ -15,6 +15,7 @@ python main.py claude-agents launch-darkly-documentation-summarizer "What is the LaunchDarkly AI SDK?" python main.py openai-agents launch-darkly-documentation-summarizer-open-ai-only "What is the LaunchDarkly AI SDK?" python main.py langchain-agents launch-darkly-documentation-summarizer "What is the LaunchDarkly AI SDK?" + python main.py google-adk-agents launch-darkly-documentation-summarizer "What is the LaunchDarkly AI SDK?" python main.py langchain-thinking launch-darkly-documentation-summarizer-messages-claude "Reason it out yourself without any tools: what is 17 times 23?" python main.py native-graph travel-agent-flow "Book me a flight to Paris" python main.py native-graph-langchain travel-agent-flow "Book me a flight to Paris" @@ -56,6 +57,7 @@ "openai-only": "examples.openai_only", "langchain": "examples.langchain_example", "langchain-agents": "examples.langchain_agents_example", + "google-adk-agents": "examples.google_adk_agents_example", "langchain-messages": "examples.langchain_messages_example", "langchain-thinking": "examples.langchain_thinking", "native-graph": "examples.native_graph", diff --git a/packages/google-adk-agents/README.md b/packages/google-adk-agents/README.md new file mode 100644 index 00000000..c8304ec3 --- /dev/null +++ b/packages/google-adk-agents/README.md @@ -0,0 +1,17 @@ +# Google ADK agents + +`launchdarkly-ai-google-adk-agents` is a wildcard agent handler for [Google ADK](https://google.github.io/adk-docs/). It serves any `agent` variation that does not have a more specific provider handler registered. + +```python +from launchdarkly_ai_google_adk_agents import google_adk_agents + +response = await google_adk_agents(key, user_input, context, tool_handlers=tools) +``` + +Gemini Developer API is the default. Set `GOOGLE_API_KEY`, `GEMINI_API_KEY`, or `GOOGLE_GENAI_API_KEY`. Vertex is explicit: `create_google_adk_agents_handler(use_vertexai=True, project=..., location=...)`. A Gemini credential is never sent on the Vertex path, and a failed call is not retried on the other transport. + +Non-Gemini models use ADK's `LiteLlm` adapter (`pip install 'google-adk[extensions]'`). The model id is `{provider}/{model}` unless the configured name already contains a slash. Pass `model=` to supply your own ADK model and skip both constructors. + +`google-adk` is not declared on this package. Release 2.9.2 caps `opentelemetry-sdk<=1.42.1`, which conflicts with this workspace's exporter. Install ADK in the application environment. + +`to_adk_agents(graph)` compiles a LaunchDarkly agent graph into an ADK `Workflow`. `google_adk_graph(key)` pre-wires this handler for `graph()`. diff --git a/packages/google-adk-agents/agents.md b/packages/google-adk-agents/agents.md new file mode 100644 index 00000000..0cf9dbe4 --- /dev/null +++ b/packages/google-adk-agents/agents.md @@ -0,0 +1,20 @@ +# Agent Guide — `launchdarkly-ai-google-adk-agents` + +Wildcard agent handler (`provides_for = ("*", "agent")`) for Google ADK. An exact provider handler still wins when both are registered. + +Gemini Developer API is the default. Vertex is opt-in (`use_vertexai=True`) and requires project and location at factory time. Do not send an API key in Vertex mode, and do not fall back from one transport to the other. + +Python non-Gemini models go through ADK's `google.adk.models.lite_llm.LiteLlm` as `provider/model`. That is ADK's adapter, not `launchdarkly-ai-litellm`. TypeScript `@google/adk` 2.1 has no equivalent; inject `model` or the handler throws. + +`google-adk` is imported lazily and is not a workspace dependency, because 2.9.2 caps `opentelemetry-sdk`. Install it in the app environment (`google-adk[extensions]` when you need LiteLLM). + +## File map + +| File | Responsibility | +|---|---| +| `handler.py` | Factory, session seeding, `LaunchDarklyTelemetryPlugin`, run and stream | +| `spans.py` | `invoke_agent` / `chat {model}` / `execute_tool {name}` | +| `graph.py` | `google_adk_graph()` | +| `native_graph.py` | `to_adk_agents()` | + +Span calls go through the `spans` module attribute so tests can patch them. Tool spans are abandoned from a `finally` (`abandon_open_spans` does not call `fail_span`). `gen_ai.system` is `google_adk`. `gen_ai.provider.name` is the serving provider; Google, Gemini, and Vertex normalize to `gcp.gemini`. diff --git a/packages/google-adk-agents/pyproject.toml b/packages/google-adk-agents/pyproject.toml new file mode 100644 index 00000000..8a48a74a --- /dev/null +++ b/packages/google-adk-agents/pyproject.toml @@ -0,0 +1,36 @@ +[project] +name = "launchdarkly-ai-google-adk-agents" +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [ + "launchdarkly-ai-server", + "opentelemetry-api>=1.25", +] +# google-adk[extensions]>=2.9.2 is the runtime dependency. It is intentionally +# not declared yet: 2.9.2 caps opentelemetry-sdk<=1.42.1, which the workspace +# cannot solve next to opentelemetry-exporter-otlp-proto-http>=1.43. +description = "Google ADK agent handler for LaunchDarkly AI SDK" +readme = "README.md" +license = "Apache-2.0" +authors = [{name = "LaunchDarkly", email = "team@launchdarkly.com"}] +keywords = ["launchdarkly", "ai", "google-adk", "agents"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.12", + "Topic :: Software Development :: Libraries", +] + +[project.urls] +Homepage = "https://github.com/launchdarkly/python-ai-sdk" +Repository = "https://github.com/launchdarkly/python-ai-sdk" +"Bug Tracker" = "https://github.com/launchdarkly/python-ai-sdk/issues" + +[build-system] +requires = ["hatchling>=1.32,<2"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/launchdarkly_ai_google_adk_agents"] diff --git a/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/__init__.py b/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/__init__.py new file mode 100644 index 00000000..184c4088 --- /dev/null +++ b/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/__init__.py @@ -0,0 +1,26 @@ +"""LaunchDarkly AI SDK integration for Google ADK agents. + +See https://launchdarkly.com/docs for usage. +""" + +__version__ = "0.1.0" # x-release-please-version + +from launchdarkly_ai_server import register_ai_sdk_package + +from .graph import google_adk_graph +from .handler import ( + create_google_adk_agents_handler, + google_adk_agents, + history_contents, +) +from .native_graph import to_adk_agents + +__all__ = [ + "create_google_adk_agents_handler", + "google_adk_agents", + "google_adk_graph", + "history_contents", + "to_adk_agents", +] + +register_ai_sdk_package("launchdarkly-ai-google-adk-agents", __version__) diff --git a/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/graph.py b/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/graph.py new file mode 100644 index 00000000..8ce5af98 --- /dev/null +++ b/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/graph.py @@ -0,0 +1,40 @@ +"""Graph convenience wrapper for the Google ADK agents handler.""" + +from __future__ import annotations + +from typing import Any + +from launchdarkly_ai_server import graph + +from .handler import create_google_adk_agents_handler + + +def google_adk_graph( + key: str, + *, + use_vertexai: bool = False, + project: str | None = None, + location: str | None = None, + model: Any = None, + capture_content: bool = False, + **options: Any, +) -> Any: + """Runs an agent graph with one wildcard ADK handler pre-bound. + + ``handlers`` passed by the caller is ignored. Auth and model options go to the + handler factory, not to ``graph()``. + """ + options.pop("handlers", None) + return graph( + key, + handlers=[ + create_google_adk_agents_handler( + use_vertexai=use_vertexai, + project=project, + location=location, + model=model, + capture_content=capture_content, + ) + ], + **options, + ) diff --git a/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/handler.py b/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/handler.py new file mode 100644 index 00000000..42cbf5e7 --- /dev/null +++ b/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/handler.py @@ -0,0 +1,694 @@ +"""Google ADK agent handler. + +Gemini Developer API is the default transport. Vertex is an explicit opt-in and never +shares an API key with that path. A failed Gemini call is not retried on Vertex. + +Non-Gemini models go through ADK's own ``LiteLlm`` adapter (``provider/model``). That +is ``google.adk.models.lite_llm``, not the LaunchDarkly LiteLLM package. +""" + +from __future__ import annotations + +import base64 +import copy +import inspect +import os +from types import SimpleNamespace +from typing import Any + +from launchdarkly_ai_server import ( + NATIVE_TOOL_KEY, + AiConfigRep, + NativeTool, + config, + create_handler, + end_unfinished_spans, + parse_template, + set_input_content_attributes, + set_output_content_attributes, + set_tool_call_content_attributes, + text_message, +) + +from . import spans as spanlib + +try: + from google.adk.agents.llm_agent import LlmAgent + from google.adk.models.google_llm import Gemini + from google.adk.plugins.base_plugin import BasePlugin + from google.adk.runners import InMemoryRunner + from google.adk.tools.function_tool import FunctionTool +except ImportError: # pragma: no cover - unit tests patch these names + LlmAgent = None + Gemini = None + InMemoryRunner = None + FunctionTool = None + + class BasePlugin: # type: ignore[no-redef] + def __init__(self, name: str) -> None: + self.name = name + + +try: + from google.adk.models.lite_llm import LiteLlm +except Exception: # pragma: no cover - extensions extra, or the package is absent + LiteLlm = None + +try: + from google.adk.events.event import Event +except ImportError: # pragma: no cover + Event = None + +try: + from google.genai import types as genai_types +except ImportError: # pragma: no cover + genai_types = None + +_GOOGLE_PROVIDERS = frozenset( + {"google", "gemini", "vertex", "google-genai", "gcp.gemini"} +) +_APP_NAME = "launchdarkly" + + +class LaunchDarklyTelemetryPlugin(BasePlugin): # type: ignore[misc] + """ADK plugin that opens ``chat`` and ``execute_tool`` spans around the runner.""" + + def __init__( + self, config: AiConfigRep, parent: Any, capture_content: bool = False + ) -> None: + super().__init__("launchdarkly") + self.config = config + self.parent = parent + self.capture_content = capture_content + self._tool_spans: dict[str, Any] = {} + self._native_stubs: dict[str, Any] = {} + self._instruction = "" + self._user_text = "" + + def bind_turn(self, instruction: str, user_text: str, handlers: Any) -> None: + self._instruction = instruction + self._user_text = user_text + self._native_stubs = native_stubs(handlers) + + async def after_model_callback( + self, *, callback_context: Any, llm_response: Any + ) -> None: + span = spanlib.start_model_span(self.config, self.parent) + messages = [text_message("user", self._user_text)] if self._user_text else [] + set_input_content_attributes( + span, + self.capture_content, + system_instructions=self._instruction or None, + messages=messages, + ) + output = event_text(llm_response) + if output: + set_output_content_attributes( + span, self.capture_content, [text_message("assistant", output)] + ) + prompt, candidates, _total = spanlib.usage_counts( + getattr(llm_response, "usage_metadata", None) + ) + spanlib.finish_model_span( + span, self.config, spanlib.span_usage_from_counts(prompt, candidates) + ) + + async def before_tool_callback( + self, *, tool: Any, tool_args: Any, tool_context: Any + ) -> None: + name = str(getattr(tool, "name", "") or "tool") + call_id = _call_id(tool_context, name) + self._tool_spans[call_id] = spanlib.start_tool_span(name, call_id, self.parent) + stub = self._native_stubs.get(name) + if callable(stub): + stub() + + async def after_tool_callback( + self, *, tool: Any, tool_args: Any, tool_context: Any, result: Any + ) -> None: + name = str(getattr(tool, "name", "") or "tool") + call_id = _call_id(tool_context, name) + span = self._tool_spans.pop(call_id, None) + set_tool_call_content_attributes( + span, self.capture_content, arguments=tool_args, result=result + ) + spanlib.succeed_span(span) + + async def on_tool_error_callback( + self, *, tool: Any, tool_args: Any, tool_context: Any, error: BaseException + ) -> dict[str, str]: + name = str(getattr(tool, "name", "") or "tool") + call_id = _call_id(tool_context, name) + span = self._tool_spans.pop(call_id, None) or spanlib.start_tool_span( + name, call_id, self.parent + ) + spanlib.fail_span(span, error) + return {"error": str(error)} + + def abandon_open_spans(self, ended: set[int] | None = None) -> None: + tracker = ended if isinstance(ended, set) else set() + for span in list(self._tool_spans.values()): + spanlib.abandon_open_spans([span], tracker) + self._tool_spans.clear() + + def close_open_spans(self, err: Exception, ended: set[int] | None = None) -> None: + tracker = ended if isinstance(ended, set) else set() + for span in list(self._tool_spans.values()): + spanlib.fail_span(span, err, tracker) + self._tool_spans.clear() + + +def history_contents(history: list[dict[str, Any]] | None) -> list[Any]: + """Prior turns as ADK contents. Image bytes are raw, not a data URL.""" + contents: list[Any] = [] + for message in history or []: + role = "model" if message.get("role") == "assistant" else "user" + contents.append(_content(role, _parts(message.get("content")))) + return contents + + +def create_google_adk_agents_handler( + *, + api_key: str | None = None, + use_vertexai: bool = False, + project: str | None = None, + location: str | None = None, + model: Any = None, + capture_content: bool = False, +) -> Any: + """Builds a wildcard agent handler. Vertex project and location are required up front.""" + if use_vertexai: + project = project or os.environ.get("GOOGLE_CLOUD_PROJECT") + location = location or os.environ.get("GOOGLE_CLOUD_LOCATION") + if not project or not location: + raise ValueError( + "Vertex mode requires a Google Cloud project and location " + "(arguments or GOOGLE_CLOUD_PROJECT / GOOGLE_CLOUD_LOCATION)" + ) + + async def _run( + cfg: AiConfigRep, + user_input: str | None = None, + tool_handlers: dict[str, Any] | None = None, + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, + ) -> dict[str, Any]: + variables = variables or {} + ended: set[int] = set() + root = spanlib.start_root_span(cfg, variables) + plugin = LaunchDarklyTelemetryPlugin( + cfg, spanlib.parent_context_of(root), capture_content + ) + try: + output, usage = await _execute( + cfg, + user_input, + tool_handlers, + variables, + history, + plugin, + api_key=api_key, + use_vertexai=use_vertexai, + project=project, + location=location, + model=model, + output_schema=cfg.get("outputFormat"), + ) + spanlib.finish_root_span( + root, + cfg, + spanlib.span_usage_from_counts(usage["input"], usage["output"]), + ) + spanlib.succeed_span(root) + return {"output": output, "usage": usage} + except Exception as err: + spanlib.fail_span(root, err, ended) + plugin.close_open_spans(err, ended) + raise + finally: + plugin.abandon_open_spans(ended) + end_unfinished_spans(root) + + async def _stream( + cfg: AiConfigRep, + user_input: str | None = None, + tool_handlers: dict[str, Any] | None = None, + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, + ) -> Any: + variables = variables or {} + ended: set[int] = set() + root = spanlib.start_root_span(cfg, variables) + plugin = LaunchDarklyTelemetryPlugin( + cfg, spanlib.parent_context_of(root), capture_content + ) + try: + runner, session, user_id = await _open_run( + cfg, + user_input, + tool_handlers, + variables, + history, + plugin, + api_key=api_key, + use_vertexai=use_vertexai, + project=project, + location=location, + model=model, + output_schema=None, + ) + output = "" + usage = _zero_usage() + async for event in runner.run_async( + user_id=user_id, + session_id=session.id, + new_message=_user_message(user_input), + ): + text = event_text(event) + if getattr(event, "partial", False): + if text: + yield {"type": "chunk", "text": text} + continue + if _is_final(event): + output = text + usage = _add_usage(usage, event) + yield {"type": "done", "output": output, "usage": usage} + spanlib.finish_root_span( + root, + cfg, + spanlib.span_usage_from_counts(usage["input"], usage["output"]), + ) + spanlib.succeed_span(root) + except Exception as err: + spanlib.fail_span(root, err, ended) + plugin.close_open_spans(err, ended) + raise + finally: + plugin.abandon_open_spans(ended) + end_unfinished_spans(root) + + return create_handler(("*", "agent"), _run, _stream, capture_content) + + +def google_adk_agents( + config_key: str, + user_input: str, + context: Any, + **kwargs: Any, +) -> Any: + """Convenience wrapper: one wildcard ADK handler and ``config().invoke()``.""" + variables = kwargs.pop("variables", None) + capture_content = kwargs.pop("capture_content", False) + api_key = kwargs.pop("api_key", None) + use_vertexai = kwargs.pop("use_vertexai", False) + project = kwargs.pop("project", None) + location = kwargs.pop("location", None) + model = kwargs.pop("model", None) + return config( + key=config_key, + handler=create_google_adk_agents_handler( + api_key=api_key, + use_vertexai=use_vertexai, + project=project, + location=location, + model=model, + capture_content=capture_content, + ), + **kwargs, + ).invoke(user_input, context, variables=variables) + + +async def _execute( + cfg: AiConfigRep, + user_input: str | None, + tool_handlers: dict[str, Any] | None, + variables: dict[str, Any], + history: list[dict[str, Any]] | None, + plugin: LaunchDarklyTelemetryPlugin, + *, + api_key: str | None, + use_vertexai: bool, + project: str | None, + location: str | None, + model: Any, + output_schema: Any, +) -> tuple[str, dict[str, int]]: + runner, session, user_id = await _open_run( + cfg, + user_input, + tool_handlers, + variables, + history, + plugin, + api_key=api_key, + use_vertexai=use_vertexai, + project=project, + location=location, + model=model, + output_schema=output_schema, + ) + output = "" + usage = _zero_usage() + async for event in runner.run_async( + user_id=user_id, + session_id=session.id, + new_message=_user_message(user_input), + ): + if _is_final(event): + output = event_text(event) + usage = _add_usage(usage, event) + return output, usage + + +async def _open_run( + cfg: AiConfigRep, + user_input: str | None, + tool_handlers: dict[str, Any] | None, + variables: dict[str, Any], + history: list[dict[str, Any]] | None, + plugin: LaunchDarklyTelemetryPlugin, + *, + api_key: str | None, + use_vertexai: bool, + project: str | None, + location: str | None, + model: Any, + output_schema: Any, +) -> tuple[Any, Any, str]: + resolved = _resolve_model( + cfg, + model, + api_key=api_key, + use_vertexai=use_vertexai, + project=project, + location=location, + ) + agent_kwargs: dict[str, Any] = { + "name": "agent", + "model": resolved, + "instruction": _instruction(cfg, variables), + "tools": _tools(cfg, tool_handlers), + } + if output_schema is not None: + agent_kwargs["output_schema"] = normalize_output_schema(output_schema) + plugin.bind_turn(_instruction(cfg, variables), user_input or "", tool_handlers) + agent = LlmAgent(**agent_kwargs) + runner = InMemoryRunner(agent, app_name=_APP_NAME, plugins=[plugin]) + user_id = _user_id(variables) + session = await runner.session_service.create_session( + app_name=_APP_NAME, user_id=user_id + ) + for content in history_contents(history): + await runner.session_service.append_event(session, _session_event(content)) + return runner, session, user_id + + +def _resolve_model( + cfg: AiConfigRep, + injected: Any, + *, + api_key: str | None, + use_vertexai: bool, + project: str | None, + location: str | None, +) -> Any: + if injected is not None: + if callable(injected) and not isinstance(injected, type): + return injected(cfg) + return injected + provider = str((cfg.get("provider") or {}).get("name") or "") + name = str((cfg.get("model") or {}).get("name") or "") + if "/" in name or provider.lower() not in _GOOGLE_PROVIDERS: + model_id = name if "/" in name else f"{provider.lower()}/{name}" + if LiteLlm is None: + raise RuntimeError( + "Non-Gemini models require LiteLLM (install google-adk[extensions])" + ) + return LiteLlm(model=model_id) + client_kwargs: dict[str, Any] = {} + if use_vertexai: + client_kwargs = {"vertexai": True, "project": project, "location": location} + elif api_key: + client_kwargs = {"api_key": api_key} + if client_kwargs: + return Gemini(model=name, client_kwargs=client_kwargs) + return Gemini(model=name) + + +def _instruction(cfg: AiConfigRep, variables: dict[str, Any]) -> str: + raw = cfg.get("instructions") + if not raw: + parts = [ + message.get("content", "") + for message in cfg.get("messages") or [] + if message.get("role") == "system" + and isinstance(message.get("content"), str) + ] + raw = "\n\n".join(parts) + return parse_template(str(raw or ""), variables) + + +def normalize_output_schema(schema: Any) -> Any: + """Copy a JSON schema and mark every property required.""" + if not isinstance(schema, dict): + return schema + copied = copy.deepcopy(schema) + properties = copied.get("properties") + if isinstance(properties, dict): + copied.setdefault("required", list(properties)) + copied.setdefault("additionalProperties", False) + return copied + + +def native_of(func: Any) -> NativeTool | None: + if isinstance(func, NativeTool): + return func + native = getattr(func, NATIVE_TOOL_KEY, None) + return native if isinstance(native, NativeTool) else None + + +def native_stubs(handlers: dict[str, Any] | None) -> dict[str, Any]: + stubs: dict[str, Any] = {} + for key, func in (handlers or {}).items(): + native = native_of(func) + if native is None or not callable(func): + continue + stubs[native.tool_name] = func + stubs[str(key)] = func + return stubs + + +def _tools(cfg: AiConfigRep, handlers: dict[str, Any] | None) -> list[Any]: + catalog = cfg.get("tools") or {} + built: list[Any] = [] + for key, spec in catalog.items(): + info = spec if isinstance(spec, dict) else {} + name = str(info.get("name") or key) + func = (handlers or {}).get(name) or (handlers or {}).get(key) + if func is None: + continue + native = native_of(func) + if native is not None: + builtin = builtin_tool(native.tool_name) + if builtin is not None: + built.append(builtin) + continue + built.append( + make_function_tool( + FunctionTool, + func, + name=name, + description=str(info.get("description") or ""), + parameters=info.get("parameters"), + ) + ) + return built + + +def adapt_handler(func: Any, name: str) -> Any: + """ADK calls tools with keywords. Mapping-style handlers still get one dict.""" + params = [ + param + for param in inspect.signature(func).parameters.values() + if param.name not in {"self", "tool_context", "cls"} + ] + single_mapping = len(params) == 1 and params[0].name in {"args", "input"} + + def _call(*args: Any, **kwargs: Any) -> Any: + if single_mapping and kwargs and not args: + return func(kwargs) + return func(*args, **kwargs) + + if inspect.iscoroutinefunction(func): + adapted: Any = _async_adapter(_call) + else: + adapted = _sync_adapter(_call) + adapted.__name__ = name + adapted.__doc__ = inspect.getdoc(func) or "" + return adapted + + +def _sync_adapter(call: Any) -> Any: + def wrapped(*args: Any, **kwargs: Any) -> Any: + return call(*args, **kwargs) + + return wrapped + + +def _async_adapter(call: Any) -> Any: + async def wrapped(*args: Any, **kwargs: Any) -> Any: + return await call(*args, **kwargs) + + return wrapped + + +def make_function_tool( + function_tool: Any, + func: Any, + *, + name: str, + description: str = "", + parameters: Any = None, +) -> Any: + adapted = adapt_handler(func, name) + try: + return function_tool( + adapted, name=name, description=description, parameters=parameters + ) + except TypeError: + tool = function_tool(adapted) + for attr, value in ( + ("name", name), + ("description", description), + ("parameters", parameters), + ): + try: + setattr(tool, attr, value) + except Exception: + continue + return tool + + +def builtin_tool(name: str) -> Any: + try: + import google.adk.tools as tools + except ImportError: + return None + exported = getattr(tools, name, None) + if exported is None or exported is FunctionTool: + return None + return exported + + +def _call_id(tool_context: Any, fallback: str) -> str: + for attr in ("function_call_id", "functionCallId"): + value = ( + tool_context.get(attr) + if isinstance(tool_context, dict) + else getattr(tool_context, attr, None) + ) + if value: + return str(value) + return fallback + + +def _user_id(variables: dict[str, Any]) -> str: + context = variables.get("ldContext") + if isinstance(context, dict) and context.get("key"): + return str(context["key"]) + return "user" + + +def _user_message(user_input: str | None) -> Any: + return _content("user", [_text_part(user_input or "")]) + + +def _session_event(content: Any) -> Any: + author = getattr(content, "role", None) or "user" + if Event is not None: + try: + return Event(author=author, content=content) + except Exception: + pass + return SimpleNamespace(author=author, content=content) + + +def event_text(event: Any) -> str: + content = getattr(event, "content", None) + if content is None and isinstance(event, dict): + content = event.get("content") + parts = getattr(content, "parts", None) if content is not None else None + if parts is None and isinstance(content, dict): + parts = content.get("parts") + texts: list[str] = [] + for part in parts or []: + text = ( + part.get("text") if isinstance(part, dict) else getattr(part, "text", None) + ) + if text: + texts.append(str(text)) + return "".join(texts) + + +def _is_final(event: Any) -> bool: + if getattr(event, "partial", False): + return False + checker = getattr(event, "is_final_response", None) + if callable(checker): + return bool(checker()) + return True + + +def _zero_usage() -> dict[str, int]: + return {"input": 0, "output": 0, "total": 0} + + +def _add_usage(usage: dict[str, int], event: Any) -> dict[str, int]: + prompt, candidates, total = spanlib.usage_counts( + getattr(event, "usage_metadata", None) + ) + return { + "input": usage["input"] + prompt, + "output": usage["output"] + candidates, + "total": usage["total"] + total, + } + + +def _parts(content: Any) -> list[Any]: + if isinstance(content, str): + return [_text_part(content)] + if not isinstance(content, list): + return [_text_part("")] + parts: list[Any] = [] + for block in content: + if not isinstance(block, dict): + continue + if block.get("type") == "text": + parts.append(_text_part(str(block.get("text") or ""))) + elif block.get("type") == "image": + source = block.get("source") or {} + raw = base64.b64decode(str(source.get("data") or "")) + parts.append(_image_part(str(source.get("media_type") or "image/png"), raw)) + return parts + + +def _text_part(text: str) -> Any: + if genai_types is not None: + return genai_types.Part(text=text) + return SimpleNamespace(text=text, inline_data=None) + + +def _image_part(mime_type: str, data: bytes) -> Any: + if genai_types is not None: + return genai_types.Part( + inline_data=genai_types.Blob(mime_type=mime_type, data=data) + ) + return SimpleNamespace( + text=None, inline_data=SimpleNamespace(mime_type=mime_type, data=data) + ) + + +def _content(role: str, parts: list[Any]) -> Any: + if genai_types is not None: + return genai_types.Content(role=role, parts=parts) + return SimpleNamespace(role=role, parts=parts) diff --git a/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/native_graph.py b/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/native_graph.py new file mode 100644 index 00000000..02a36631 --- /dev/null +++ b/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/native_graph.py @@ -0,0 +1,289 @@ +"""Native ADK workflow adapter for a LaunchDarkly agent graph.""" + +from __future__ import annotations + +import time +import uuid +from typing import Any + +from launchdarkly_ai_server import ( + get_client, + make_track_data, + parse_template, + to_ld_context, +) + +from .handler import ( + _session_event, + event_text, + history_contents, + make_function_tool, + native_of, +) +from .spans import usage_counts + +try: + from google.adk.agents.llm_agent import LlmAgent + from google.adk.runners import InMemoryRunner + from google.adk.tools.function_tool import FunctionTool + from google.adk.workflow import START, Workflow +except ImportError: # pragma: no cover - unit tests patch these names + LlmAgent = None + InMemoryRunner = None + FunctionTool = None + Workflow = None + START = "START" + +try: + from opentelemetry import trace + from opentelemetry.trace import StatusCode as SpanStatusCode +except ImportError: # pragma: no cover + trace = None # type: ignore[assignment] + SpanStatusCode = None # type: ignore[assignment,misc] + +_TRACER = "@launchdarkly/ai-google-adk-agents" +_APP = "launchdarkly" + + +class _AdkAgents: + def __init__(self, graph_def: Any, options: dict[str, Any]) -> None: + self._graph = graph_def + self._options = options + + async def invoke( + self, + user_input: str | None = None, + context: Any = None, + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, + ) -> dict[str, Any]: + graph_def = self._graph + if not getattr(graph_def, "enabled", True): + raise RuntimeError(f"Graph {graph_def.key} is disabled") + if getattr(graph_def, "root", None) is None: + raise RuntimeError("Graph has no root") + + ld_context = None + if context is not None: + ld_context = to_ld_context(get_client(), context) + span = _graph_span(graph_def.key) + start = time.monotonic() + run_id = str(uuid.uuid4()) + variables = variables or {} + tool_handlers = self._options.get("tool_handlers") or {} + handoff: dict[str, str | None] = {"target": None} + agents = _build_agents(graph_def, variables, tool_handlers, handoff) + root_agent = agents[graph_def.root.key] + Workflow(name="graph", edges=[(START, root_agent)]) + + path: list[str] = [] + usage = {"input": 0, "output": 0, "total": 0} + output = "" + try: + current = graph_def.root + seen: set[str] = set() + seed = history + while current is not None and current.key not in seen: + seen.add(current.key) + path.append(current.key) + handoff["target"] = None + output, step = await _run_agent( + agents[current.key], + user_input or "", + seed, + ) + seed = None + for key in usage: + usage[key] += step[key] + target = handoff["target"] + if target and ld_context is not None: + get_client().track( + "$ld:ai:graph:handoff_success", + ld_context, + make_track_data(current, graph_def.key, run_id), + 1, + ) + current = graph_def.get_node(target) if target else None + if span is not None: + span.set_attribute("ld.ai.graph.path", "->".join(path)) + if SpanStatusCode is not None: + span.set_status(SpanStatusCode.OK) + if ld_context is not None: + _track_success( + graph_def.root, + graph_def.key, + run_id, + ld_context, + path, + usage, + start, + ) + return {"response": output, "usage": usage} + except Exception as exc: + if span is not None: + span.record_exception(exc) + if SpanStatusCode is not None: + span.set_status(SpanStatusCode.ERROR, str(exc)) + if ld_context is not None: + get_client().track( + "$ld:ai:graph:invocation_failure", + ld_context, + make_track_data(graph_def.root, graph_def.key, run_id), + 1, + ) + raise + finally: + if span is not None: + span.end() + + +def to_adk_agents(graph_def: Any, **options: Any) -> _AdkAgents: + """Compiles a LaunchDarkly agent graph into an ADK workflow.""" + return _AdkAgents(graph_def, options) + + +def _graph_span(key: str) -> Any: + if trace is None: + return None + span = trace.get_tracer(_TRACER).start_span("ld.ai.graph") + span.set_attribute("ld.ai.graph.key", key) + return span + + +def _track_success( + root: Any, + graph_key: str, + run_id: str, + ld_context: Any, + path: list[str], + usage: dict[str, int], + start: float, +) -> None: + client = get_client() + data = make_track_data(root, graph_key, run_id) + duration = int((time.monotonic() - start) * 1000) + client.track("$ld:ai:graph:duration:total", ld_context, data, duration) + client.track("$ld:ai:graph:total_tokens", ld_context, data, usage["total"]) + client.track("$ld:ai:graph:path", ld_context, data, len(path)) + client.track("$ld:ai:graph:invocation_success", ld_context, data, 1) + + +def _build_agents( + graph_def: Any, + variables: dict[str, Any], + tool_handlers: dict[str, Any], + handoff: dict[str, str | None], +) -> dict[str, Any]: + agents: dict[str, Any] = {} + _visit(graph_def, graph_def.root, agents, variables, tool_handlers, handoff) + return agents + + +def _visit( + graph_def: Any, + node: Any, + agents: dict[str, Any], + variables: dict[str, Any], + tool_handlers: dict[str, Any], + handoff: dict[str, str | None], +) -> None: + if node is None or node.key in agents: + return + targets: list[str] = [] + tools: list[Any] = [] + for edge in node.edges or []: + target_key = getattr(edge, "target_key", None) or getattr( + edge, "targetKey", None + ) + if not target_key: + continue + targets.append(str(target_key)) + tools.append(_transfer_tool(str(target_key), handoff)) + config = node.config or {} + tools = _node_tools(config, tool_handlers) + tools + agents[node.key] = LlmAgent( + name=_agent_name(node.key), + model=str((config.get("model") or {}).get("name") or ""), + instruction=parse_template(str(config.get("instructions") or ""), variables), + tools=tools, + ) + for target_key in targets: + _visit( + graph_def, + graph_def.get_node(target_key), + agents, + variables, + tool_handlers, + handoff, + ) + + +def _node_tools(config: dict[str, Any], handlers: dict[str, Any]) -> list[Any]: + built: list[Any] = [] + for key, spec in (config.get("tools") or {}).items(): + info = spec if isinstance(spec, dict) else {} + name = str(info.get("name") or key) + func = handlers.get(name) or handlers.get(key) + if func is None or native_of(func) is not None: + continue + built.append( + make_function_tool( + FunctionTool, + func, + name=name, + description=str(info.get("description") or ""), + parameters=info.get("parameters"), + ) + ) + return built + + +def _transfer_tool(target_key: str, handoff: dict[str, str | None]) -> Any: + def _select() -> str: + handoff["target"] = target_key + return target_key + + return make_function_tool(FunctionTool, _select, name=f"transfer_to_{target_key}") + + +def _agent_name(key: str) -> str: + cleaned = "".join(ch if ch.isalnum() or ch == "_" else "_" for ch in key) + return cleaned or "agent" + + +async def _run_agent( + agent: Any, + user_input: str, + history: list[dict[str, Any]] | None, +) -> tuple[str, dict[str, int]]: + runner = InMemoryRunner(agent, app_name=_APP) + session = await runner.session_service.create_session(app_name=_APP, user_id="user") + if history: + for content in history_contents(history): + await runner.session_service.append_event(session, _session_event(content)) + output = "" + usage = {"input": 0, "output": 0, "total": 0} + async for event in runner.run_async( + user_id="user", + session_id=session.id, + new_message=_message(user_input), + ): + if getattr(event, "partial", False): + continue + checker = getattr(event, "is_final_response", None) + if callable(checker) and not checker(): + continue + output = event_text(event) + prompt, candidates, total = usage_counts(getattr(event, "usage_metadata", None)) + usage = { + "input": usage["input"] + prompt, + "output": usage["output"] + candidates, + "total": usage["total"] + total, + } + return output, usage + + +def _message(text: str) -> Any: + return type( + "Content", (), {"role": "user", "parts": [type("Part", (), {"text": text})()]} + )() diff --git a/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/spans.py b/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/spans.py new file mode 100644 index 00000000..81291abd --- /dev/null +++ b/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/spans.py @@ -0,0 +1,184 @@ +"""Span construction for the Google ADK agents handler. + +The shape is ``invoke_agent`` root, one ``chat {model}`` child per model turn, one +``execute_tool {name}`` child per tool call. Tool spans are siblings of the ``chat`` span: +both take the root's context. ``gen_ai.system`` is the framework (``google_adk``). +``gen_ai.provider.name`` is who served the model. Google-family providers normalize to +``gcp.gemini``. +""" + +from __future__ import annotations + +from typing import Any + +from launchdarkly_ai_server import ( + AiConfigRep, + SpanUsage, + end_span_once, + number_or_zero, + set_ld_span_attributes, + set_model_identity_attributes, + set_usage_span_attributes, +) + +try: + from opentelemetry import trace + from opentelemetry.trace import StatusCode as SpanStatusCode + + _HAS_OTEL = True +except ImportError: # pragma: no cover + trace = None # type: ignore[assignment] + SpanStatusCode = None # type: ignore[assignment,misc] + _HAS_OTEL = False + +TRACER_NAME = "@launchdarkly/ai-google-adk-agents" + +_GOOGLE_PROVIDERS = frozenset( + {"google", "gemini", "vertex", "google-genai", "gcp.gemini"} +) + + +def model_name(config: AiConfigRep) -> str: + return str((config.get("model") or {}).get("name") or "") + + +def serving_provider(config: AiConfigRep) -> str: + """Who served the model, for ``gen_ai.provider.name``. + + Google, Gemini, and Vertex are one provider in the semconv enum. Every other + configured name is passed through lower-cased. An empty name stays on Gemini, + which is this handler's default transport. + """ + name = str((config.get("provider") or {}).get("name") or "").lower() + if not name or name in _GOOGLE_PROVIDERS: + return "gcp.gemini" + return name + + +def start_root_span(config: AiConfigRep, variables: dict[str, Any]) -> Any: + """Opens the ``invoke_agent`` root. It is the only span with LaunchDarkly identity.""" + if not _HAS_OTEL: + return None + span = trace.get_tracer(TRACER_NAME).start_span("invoke_agent") + span.set_attribute("gen_ai.operation.name", "invoke_agent") + set_model_identity_attributes( + span, + serving_provider(config), + model_name(config), + legacy_system="google_adk", + ) + set_ld_span_attributes(span, variables) + return span + + +def parent_context_of(span: Any) -> Any: + if not _HAS_OTEL or span is None: + return None + return trace.set_span_in_context(span) + + +def start_model_span(config: AiConfigRep, parent: Any) -> Any: + """Opens one ``chat {model}`` span. Children do not repeat LaunchDarkly identity.""" + if not _HAS_OTEL: + return None + name = model_name(config) + span = trace.get_tracer(TRACER_NAME).start_span(f"chat {name}", context=parent) + span.set_attribute("gen_ai.operation.name", "chat") + set_model_identity_attributes( + span, serving_provider(config), name, legacy_system="google_adk" + ) + return span + + +def start_tool_span(tool_name: str, tool_call_id: str, parent: Any) -> Any: + """Opens one ``execute_tool {name}`` span.""" + if not _HAS_OTEL: + return None + span = trace.get_tracer(TRACER_NAME).start_span( + f"execute_tool {tool_name}", context=parent + ) + span.set_attribute("gen_ai.operation.name", "execute_tool") + span.set_attribute("gen_ai.tool.name", tool_name) + span.set_attribute("gen_ai.tool.call.id", tool_call_id) + return span + + +def span_usage_from_counts(input_tokens: int, output_tokens: int) -> SpanUsage: + return SpanUsage(input=input_tokens, output=output_tokens) + + +def finish_root_span(span: Any, config: AiConfigRep, usage: SpanUsage) -> None: + if span is None: + return + span.set_attribute("gen_ai.response.model", model_name(config)) + set_usage_span_attributes(span, usage) + + +def finish_model_span( + span: Any, + config: AiConfigRep, + usage: SpanUsage | None = None, +) -> None: + if span is None: + return + span.set_attribute("gen_ai.response.model", model_name(config)) + set_usage_span_attributes(span, usage or SpanUsage()) + if SpanStatusCode is not None: + span.set_status(SpanStatusCode.OK) + span.end() + + +def succeed_span(span: Any) -> None: + if span is None: + return + if SpanStatusCode is not None: + span.set_status(SpanStatusCode.OK) + span.end() + + +def mark_ok(span: Any) -> None: + if span is None or SpanStatusCode is None: + return + span.set_status(SpanStatusCode.OK) + + +def fail_span(span: Any, error: BaseException, tracker: set[int] | None = None) -> None: + if span is None: + return + span.record_exception(error) + if SpanStatusCode is not None: + span.set_status(SpanStatusCode.ERROR, str(error)) + if tracker is not None: + end_span_once(span, tracker) + else: + span.end() + + +def abandon_open_spans(spans: list[Any], ended: set[int] | None = None) -> None: + """Ends spans still open when a consumer walks away. Does not record a failure.""" + tracker = ended if ended is not None else set() + for span in spans: + end_span_once(span, tracker, abandoned=True) + + +def usage_counts(metadata: Any) -> tuple[int, int, int]: + """ADK ``usage_metadata`` as ``(input, output, total)``. Missing metadata is zeros.""" + if metadata is None: + return 0, 0, 0 + prompt = _field(metadata, "prompt_token_count", "promptTokenCount") + candidates = _field(metadata, "candidates_token_count", "candidatesTokenCount") + total = _field(metadata, "total_token_count", "totalTokenCount") + if total == 0 and (prompt or candidates): + total = prompt + candidates + return prompt, candidates, total + + +def _field(metadata: Any, *keys: str) -> int: + for key in keys: + if isinstance(metadata, dict): + if key in metadata: + return number_or_zero(metadata[key]) + continue + if hasattr(metadata, key): + return number_or_zero(getattr(metadata, key)) + return 0 diff --git a/packages/google-adk-agents/tests/test_graph.py b/packages/google-adk-agents/tests/test_graph.py new file mode 100644 index 00000000..059524f0 --- /dev/null +++ b/packages/google-adk-agents/tests/test_graph.py @@ -0,0 +1,60 @@ +"""Graph convenience wrapper. Reference: TESTING.md §2.1 and §2.x.4.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +from launchdarkly_ai_google_adk_agents.graph import google_adk_graph + + +class TestGoogleAdkGraph: + def test_forwards_the_flag_key(self) -> None: + with patch("launchdarkly_ai_google_adk_agents.graph.graph") as mock_graph: + mock_graph.return_value = MagicMock() + google_adk_graph("my-flag") + assert mock_graph.call_args.args[0] == "my-flag" + + def test_prewires_one_wildcard_agent_handler(self) -> None: + with patch("launchdarkly_ai_google_adk_agents.graph.graph") as mock_graph: + mock_graph.return_value = MagicMock() + google_adk_graph("my-flag") + handlers = mock_graph.call_args.kwargs["handlers"] + assert len(handlers) == 1 + assert handlers[0].provides_for == ("*", "agent") + + def test_caller_cannot_replace_handlers(self) -> None: + with patch("launchdarkly_ai_google_adk_agents.graph.graph") as mock_graph: + mock_graph.return_value = MagicMock() + google_adk_graph( + "my-flag", handlers=["nope"], tool_handlers={"lookup": lambda: None} + ) + kwargs = mock_graph.call_args.kwargs + assert len(kwargs["handlers"]) == 1 + assert kwargs["handlers"][0].provides_for == ("*", "agent") + assert "lookup" in kwargs["tool_handlers"] + + def test_vertex_and_model_options_reach_the_handler(self) -> None: + model = object() + with ( + patch("launchdarkly_ai_google_adk_agents.graph.graph") as mock_graph, + patch( + "launchdarkly_ai_google_adk_agents.graph.create_google_adk_agents_handler" + ) as factory, + ): + factory.return_value = MagicMock(provides_for=("*", "agent")) + mock_graph.return_value = MagicMock() + google_adk_graph( + "my-flag", + use_vertexai=True, + project="p", + location="us-central1", + model=model, + capture_content=True, + ) + factory.assert_called_once_with( + use_vertexai=True, + project="p", + location="us-central1", + model=model, + capture_content=True, + ) diff --git a/packages/google-adk-agents/tests/test_handler.py b/packages/google-adk-agents/tests/test_handler.py new file mode 100644 index 00000000..97ee7e8b --- /dev/null +++ b/packages/google-adk-agents/tests/test_handler.py @@ -0,0 +1,635 @@ +"""Google ADK agents handler tests. + +Reference: TESTING.md §1 and §2.x Google ADK, Appendix A.15. +The ADK runtime is mocked. These tests must not open a network connection. +""" + +from __future__ import annotations + +import inspect +from types import SimpleNamespace +from typing import Any, ClassVar +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import launchdarkly_ai_google_adk_agents.handler as handler_mod +import launchdarkly_ai_google_adk_agents.spans as spans_mod +from launchdarkly_ai_google_adk_agents.handler import ( + create_google_adk_agents_handler, + history_contents, +) + +BASE_CONFIG: dict[str, Any] = { + "model": {"name": "gemini-2.5-flash"}, + "provider": {"name": "Google"}, + "instructions": "Be helpful.", +} + + +def _usage(prompt: int = 10, candidates: int = 5, total: int = 15) -> SimpleNamespace: + return SimpleNamespace( + prompt_token_count=prompt, + candidates_token_count=candidates, + total_token_count=total, + ) + + +def _event( + text: str = "answer", + *, + partial: bool = False, + usage: Any = None, + final: bool = True, +) -> SimpleNamespace: + event = SimpleNamespace( + content=SimpleNamespace(parts=[SimpleNamespace(text=text, inline_data=None)]), + partial=partial, + usage_metadata=usage if usage is not None else _usage(), + ) + event.is_final_response = lambda: final and not partial + return event + + +def _message_text(message: Any) -> str: + parts = message.parts if hasattr(message, "parts") else message["parts"] + part = parts[0] + if hasattr(part, "text"): + return part.text + return part["text"] + + +class _FakeFunctionTool: + created: ClassVar[list[_FakeFunctionTool]] = [] + + def __init__(self, func: Any, **kwargs: Any) -> None: + self.func = func + self.name = kwargs.get("name", getattr(func, "__name__", "")) + self.description = kwargs.get("description") + self.parameters = kwargs.get("parameters") + _FakeFunctionTool.created.append(self) + + +class _FakeRunner: + last: ClassVar[_FakeRunner | None] = None + events: ClassVar[list[Any]] = [] + explode: ClassVar[BaseException | None] = None + + def __init__( + self, + agent: Any = None, + *, + app_name: str | None = None, + plugins: list[Any] | None = None, + **kwargs: Any, + ) -> None: + self.agent = agent + self.app_name = app_name + self.plugins = list(plugins or []) + self.kwargs = kwargs + session = SimpleNamespace(id="sess-1") + self.session_service = SimpleNamespace( + create_session=AsyncMock(return_value=session), + append_event=AsyncMock(), + ) + self.run_calls: list[dict[str, Any]] = [] + _FakeRunner.last = self + + async def run_async(self, **kwargs: Any) -> Any: + self.run_calls.append(kwargs) + if _FakeRunner.explode is not None: + raise _FakeRunner.explode + plugin = self.plugins[0] if self.plugins else None + if plugin is not None and hasattr(plugin, "after_model_callback"): + response = SimpleNamespace( + content=_FakeRunner.events[-1].content if _FakeRunner.events else None, + usage_metadata=_usage(), + partial=False, + ) + result = plugin.after_model_callback( + callback_context=SimpleNamespace(), + llm_response=response, + ) + if inspect.isawaitable(result): + await result + for event in _FakeRunner.events: + yield event + + +@pytest.fixture(autouse=True) +def _adk_mocks(monkeypatch: pytest.MonkeyPatch) -> None: + _FakeFunctionTool.created = [] + _FakeRunner.last = None + _FakeRunner.events = [_event()] + _FakeRunner.explode = None + monkeypatch.setattr(handler_mod, "InMemoryRunner", _FakeRunner) + monkeypatch.setattr( + handler_mod, + "LlmAgent", + MagicMock(side_effect=lambda **kw: SimpleNamespace(**kw)), + ) + monkeypatch.setattr(handler_mod, "FunctionTool", _FakeFunctionTool) + monkeypatch.setattr( + handler_mod, "Gemini", MagicMock(side_effect=lambda **kw: SimpleNamespace(**kw)) + ) + monkeypatch.setattr( + handler_mod, + "LiteLlm", + MagicMock(side_effect=lambda **kw: SimpleNamespace(**kw)), + ) + monkeypatch.setattr( + spans_mod, "start_root_span", MagicMock(return_value=MagicMock(name="root")) + ) + monkeypatch.setattr( + spans_mod, "start_model_span", MagicMock(return_value=MagicMock(name="model")) + ) + monkeypatch.setattr( + spans_mod, "start_tool_span", MagicMock(return_value=MagicMock(name="tool")) + ) + monkeypatch.setattr(spans_mod, "finish_model_span", MagicMock()) + monkeypatch.setattr(spans_mod, "finish_root_span", MagicMock()) + monkeypatch.setattr(spans_mod, "fail_span", MagicMock()) + + +def _gemini_client_kwargs() -> dict[str, Any]: + kwargs = handler_mod.Gemini.call_args.kwargs + return dict(kwargs.get("client_kwargs") or {}) + + +class TestFactory: + def test_provides_for_is_wildcard_agent(self) -> None: + handler = create_google_adk_agents_handler() + assert handler.provides_for == ("*", "agent") + + def test_capture_content_defaults_false(self) -> None: + handler = create_google_adk_agents_handler() + assert handler.capture_content is False + + def test_capture_content_can_be_enabled(self) -> None: + handler = create_google_adk_agents_handler(capture_content=True) + assert handler.capture_content is True + + def test_factory_returns_independent_handlers(self) -> None: + first = create_google_adk_agents_handler() + second = create_google_adk_agents_handler() + assert first is not second + + +class TestAuth: + async def test_default_transport_is_gemini_without_vertex(self) -> None: + await create_google_adk_agents_handler()(BASE_CONFIG, "hello") + assert handler_mod.Gemini.call_args.kwargs["model"] == "gemini-2.5-flash" + assert _gemini_client_kwargs().get("vertexai") is not True + handler_mod.LiteLlm.assert_not_called() + + async def test_explicit_api_key_is_forwarded_only_in_gemini_mode(self) -> None: + await create_google_adk_agents_handler(api_key="gemini-key")( + BASE_CONFIG, "hello" + ) + client_kwargs = _gemini_client_kwargs() + assert client_kwargs["api_key"] == "gemini-key" + assert client_kwargs.get("vertexai") is not True + + async def test_vertex_opt_in_sets_project_and_location_without_api_key( + self, + ) -> None: + await create_google_adk_agents_handler( + use_vertexai=True, + project="ld-proj", + location="us-central1", + api_key="should-not-be-used", + )(BASE_CONFIG, "hello") + client_kwargs = _gemini_client_kwargs() + assert client_kwargs["vertexai"] is True + assert client_kwargs["project"] == "ld-proj" + assert client_kwargs["location"] == "us-central1" + assert "api_key" not in client_kwargs + + async def test_vertex_reads_project_and_location_from_the_environment( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "env-proj") + monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "global") + await create_google_adk_agents_handler(use_vertexai=True)(BASE_CONFIG, "hello") + client_kwargs = _gemini_client_kwargs() + assert client_kwargs["project"] == "env-proj" + assert client_kwargs["location"] == "global" + + def test_vertex_without_project_or_location_raises( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("GOOGLE_CLOUD_PROJECT", raising=False) + monkeypatch.delenv("GOOGLE_CLOUD_LOCATION", raising=False) + with pytest.raises(ValueError, match="project"): + create_google_adk_agents_handler(use_vertexai=True) + + async def test_injected_model_skips_gemini_and_litellm(self) -> None: + model = object() + await create_google_adk_agents_handler(model=model)(BASE_CONFIG, "hello") + assert handler_mod.LlmAgent.call_args.kwargs["model"] is model + handler_mod.Gemini.assert_not_called() + handler_mod.LiteLlm.assert_not_called() + + async def test_model_factory_receives_the_config(self) -> None: + seen: list[dict[str, Any]] = [] + + def factory(config: dict[str, Any]) -> object: + seen.append(config) + return "built-model" + + await create_google_adk_agents_handler(model=factory)(BASE_CONFIG, "hello") + assert seen == [BASE_CONFIG] + assert handler_mod.LlmAgent.call_args.kwargs["model"] == "built-model" + + async def test_non_gemini_provider_uses_litellm_model_id(self) -> None: + config = { + "model": {"name": "gpt-4o"}, + "provider": {"name": "OpenAI"}, + "instructions": "Be helpful.", + } + await create_google_adk_agents_handler()(config, "hello") + assert handler_mod.LiteLlm.call_args.kwargs["model"] == "openai/gpt-4o" + handler_mod.Gemini.assert_not_called() + + async def test_slashed_model_name_is_passed_through_to_litellm(self) -> None: + config = { + "model": {"name": "anthropic/claude-sonnet-4"}, + "provider": {"name": "Anthropic"}, + "instructions": "Be helpful.", + } + await create_google_adk_agents_handler()(config, "hello") + assert ( + handler_mod.LiteLlm.call_args.kwargs["model"] == "anthropic/claude-sonnet-4" + ) + + +class TestPromptAndRun: + async def test_instructions_are_templated_onto_the_agent(self) -> None: + config = {**BASE_CONFIG, "instructions": "Hello {{name}} {{missing}}"} + await create_google_adk_agents_handler()( + config, "question", None, {"name": "Ada"} + ) + assert ( + handler_mod.LlmAgent.call_args.kwargs["instruction"] + == "Hello Ada {{missing}}" + ) + + async def test_instructions_win_over_messages(self) -> None: + config = { + **BASE_CONFIG, + "instructions": "from-instructions", + "messages": [{"role": "system", "content": "from-messages"}], + } + await create_google_adk_agents_handler()(config, "question") + assert ( + handler_mod.LlmAgent.call_args.kwargs["instruction"] == "from-instructions" + ) + + async def test_user_input_is_the_new_message_text(self) -> None: + await create_google_adk_agents_handler()(BASE_CONFIG, "hello") + runner = _FakeRunner.last + assert runner is not None + assert _message_text(runner.run_calls[0]["new_message"]) == "hello" + + async def test_missing_user_input_sends_an_empty_text_part(self) -> None: + await create_google_adk_agents_handler()(BASE_CONFIG, None) + runner = _FakeRunner.last + assert runner is not None + assert _message_text(runner.run_calls[0]["new_message"]) == "" + + async def test_runner_receives_exactly_one_plugin(self) -> None: + await create_google_adk_agents_handler()(BASE_CONFIG, "hello") + runner = _FakeRunner.last + assert runner is not None + assert len(runner.plugins) == 1 + + async def test_final_event_usage_metadata_is_returned(self) -> None: + _FakeRunner.events = [_event("the answer", usage=_usage(11, 7, 18))] + result = await create_google_adk_agents_handler()(BASE_CONFIG, "hello") + assert result["output"] == "the answer" + assert result["usage"] == {"input": 11, "output": 7, "total": 18} + + async def test_missing_usage_is_zeros(self) -> None: + _FakeRunner.events = [_event("ok", usage=None)] + _FakeRunner.events[0].usage_metadata = None + result = await create_google_adk_agents_handler()(BASE_CONFIG, "hello") + assert result["usage"] == {"input": 0, "output": 0, "total": 0} + + async def test_partial_events_are_not_double_counted(self) -> None: + _FakeRunner.events = [ + _event("hel", partial=True, usage=_usage(3, 1, 4), final=False), + _event("hello", usage=_usage(3, 2, 5)), + ] + result = await create_google_adk_agents_handler()(BASE_CONFIG, "hello") + assert result["output"] == "hello" + assert result["usage"]["total"] == 5 + + async def test_blocking_call_sets_output_schema(self) -> None: + config = { + **BASE_CONFIG, + "outputFormat": { + "type": "object", + "properties": {"ok": {"type": "boolean"}}, + "required": ["ok"], + }, + } + await create_google_adk_agents_handler()(config, "hello") + assert handler_mod.LlmAgent.call_args.kwargs.get("output_schema") is not None + + async def test_history_is_seeded_before_run(self) -> None: + history = [ + {"role": "user", "content": "earlier"}, + {"role": "assistant", "content": "ack"}, + ] + await create_google_adk_agents_handler()( + BASE_CONFIG, "now", None, None, history + ) + runner = _FakeRunner.last + assert runner is not None + assert runner.session_service.append_event.await_count == 2 + assert _message_text(runner.run_calls[0]["new_message"]) == "now" + + +class TestHistoryContents: + def test_image_history_uses_inline_data_bytes(self) -> None: + contents = history_contents( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "see"}, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "aGVsbG8=", + }, + }, + ], + } + ] + ) + image = contents[0].parts[1].inline_data + assert image.mime_type == "image/png" + assert image.data == b"hello" + + +class TestTools: + async def test_only_tools_with_handlers_are_forwarded(self) -> None: + config = { + **BASE_CONFIG, + "tools": { + "lookup": { + "name": "lookup", + "description": "find", + "parameters": {"type": "object"}, + }, + "other": { + "name": "other", + "description": "nope", + "parameters": {"type": "object"}, + }, + "extra": { + "name": "extra", + "description": "also", + "parameters": {"type": "object"}, + }, + }, + } + + def lookup(q: str) -> str: + return f"found {q}" + + def extra() -> str: + return "extra" + + await create_google_adk_agents_handler()( + config, + "hello", + {"lookup": lookup, "extra": extra}, + ) + names = {tool.name for tool in _FakeFunctionTool.created} + assert names == {"lookup", "extra"} + lookup_tool = next( + tool for tool in _FakeFunctionTool.created if tool.name == "lookup" + ) + assert lookup_tool.description == "find" + assert lookup_tool.func("ada") == "found ada" + + async def test_tool_error_is_returned_to_the_model(self) -> None: + await create_google_adk_agents_handler()(BASE_CONFIG, "hello") + plugin = _FakeRunner.last.plugins[0] + result = plugin.on_tool_error_callback( + tool=SimpleNamespace(name="lookup"), + tool_args={"q": "x"}, + tool_context=SimpleNamespace(), + error=RuntimeError("tool broke"), + ) + if inspect.isawaitable(result): + result = await result + assert isinstance(result, dict) + assert "tool broke" in str(result) + spans_mod.fail_span.assert_called() + + async def test_mapping_handler_receives_one_argument_object( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + class _StrictTool: + def __init__( + self, func: Any, *, require_confirmation: bool = False + ) -> None: + self.func = func + self.name = getattr(func, "__name__", "") + + monkeypatch.setattr(handler_mod, "FunctionTool", _StrictTool) + config = { + **BASE_CONFIG, + "tools": { + "prefs": { + "name": "get_user_preferences", + "description": "prefs", + "parameters": {"type": "object"}, + } + }, + } + + def get_preferences(args: dict[str, Any]) -> str: + return str(args["user_id"]) + + await create_google_adk_agents_handler()( + config, "hello", {"get_user_preferences": get_preferences} + ) + agent_tools = handler_mod.LlmAgent.call_args.kwargs["tools"] + built = agent_tools[0] + assert built.name == "get_user_preferences" + assert built.func(user_id="ada") == "ada" + + async def test_native_tool_is_not_wrapped_as_a_function(self) -> None: + from launchdarkly_ai_server import NATIVE_TOOL_KEY, NativeTool + + tracked: list[str] = [] + + def stub() -> None: + tracked.append("google_search") + + setattr(stub, NATIVE_TOOL_KEY, NativeTool("google_search")) + config = { + **BASE_CONFIG, + "tools": { + "search": { + "name": "google_search", + "description": "web", + "parameters": {"type": "object"}, + } + }, + } + await create_google_adk_agents_handler()( + config, "hello", {"google_search": stub} + ) + assert _FakeFunctionTool.created == [] + plugin = _FakeRunner.last.plugins[0] + started = plugin.before_tool_callback( + tool=SimpleNamespace(name="google_search"), + tool_args={}, + tool_context=SimpleNamespace(function_call_id="call-1"), + ) + if inspect.isawaitable(started): + await started + assert tracked == ["google_search"] + + async def test_output_schema_marks_properties_required(self) -> None: + schema = {"type": "object", "properties": {"ok": {"type": "boolean"}}} + await create_google_adk_agents_handler()( + {**BASE_CONFIG, "outputFormat": schema}, "hello" + ) + passed = handler_mod.LlmAgent.call_args.kwargs["output_schema"] + assert passed["required"] == ["ok"] + assert passed["additionalProperties"] is False + assert "required" not in schema + + +class TestTelemetry: + async def test_model_turn_uses_span_helpers(self) -> None: + await create_google_adk_agents_handler()(BASE_CONFIG, "hello") + spans_mod.start_root_span.assert_called() + spans_mod.start_model_span.assert_called() + spans_mod.finish_model_span.assert_called() + + async def test_tool_callback_uses_start_tool_span(self) -> None: + await create_google_adk_agents_handler()(BASE_CONFIG, "hello") + plugin = _FakeRunner.last.plugins[0] + started = plugin.before_tool_callback( + tool=SimpleNamespace(name="lookup"), + tool_args={"q": "x"}, + tool_context=SimpleNamespace(), + ) + if inspect.isawaitable(started): + await started + finished = plugin.after_tool_callback( + tool=SimpleNamespace(name="lookup"), + tool_args={"q": "x"}, + tool_context=SimpleNamespace(), + result={"ok": True}, + ) + if inspect.isawaitable(finished): + await finished + spans_mod.start_tool_span.assert_called() + assert spans_mod.start_tool_span.call_args.args[0] == "lookup" + + async def test_overlapping_calls_keep_distinct_ids(self) -> None: + await create_google_adk_agents_handler()(BASE_CONFIG, "hello") + plugin = _FakeRunner.last.plugins[0] + for call_id in ("call-1", "call-2"): + started = plugin.before_tool_callback( + tool=SimpleNamespace(name="lookup"), + tool_args={"q": call_id}, + tool_context=SimpleNamespace(function_call_id=call_id), + ) + if inspect.isawaitable(started): + await started + ids = [call.args[1] for call in spans_mod.start_tool_span.call_args_list] + assert ids == ["call-1", "call-2"] + + async def test_capture_content_writes_messages_and_tool_arguments(self) -> None: + span = MagicMock() + span.attributes = {} + span.set_attribute = lambda key, value: span.attributes.__setitem__(key, value) + spans_mod.start_model_span.return_value = span + spans_mod.start_tool_span.return_value = span + await create_google_adk_agents_handler(capture_content=True)( + BASE_CONFIG, "hello" + ) + assert "gen_ai.input.messages" in span.attributes + plugin = _FakeRunner.last.plugins[0] + started = plugin.before_tool_callback( + tool=SimpleNamespace(name="lookup"), + tool_args={"q": "ada"}, + tool_context=SimpleNamespace(function_call_id="call-1"), + ) + if inspect.isawaitable(started): + await started + finished = plugin.after_tool_callback( + tool=SimpleNamespace(name="lookup"), + tool_args={"q": "ada"}, + tool_context=SimpleNamespace(function_call_id="call-1"), + result="found", + ) + if inspect.isawaitable(finished): + await finished + assert "gen_ai.tool.call.arguments" in span.attributes + + async def test_runner_error_fails_and_ends_the_root_span(self) -> None: + _FakeRunner.explode = RuntimeError("runner broke") + with pytest.raises(RuntimeError, match="runner broke"): + await create_google_adk_agents_handler()(BASE_CONFIG, "hello") + spans_mod.fail_span.assert_called() + + async def test_abandon_open_spans_does_not_fail_them(self) -> None: + await create_google_adk_agents_handler()(BASE_CONFIG, "hello") + plugin = _FakeRunner.last.plugins[0] + spans_mod.fail_span.reset_mock() + plugin.abandon_open_spans(set()) + spans_mod.fail_span.assert_not_called() + + +class TestStreaming: + async def test_stream_yields_chunks_then_one_done(self) -> None: + _FakeRunner.events = [ + _event("hel", partial=True, final=False), + _event("hello", usage=_usage(4, 2, 6)), + ] + handler = create_google_adk_agents_handler() + events = [event async for event in await handler.stream(BASE_CONFIG, "hello")] + assert events[0] == {"type": "chunk", "text": "hel"} + assert events[-1]["type"] == "done" + assert events[-1]["output"] == "hello" + assert events[-1]["usage"] == {"input": 4, "output": 2, "total": 6} + assert sum(1 for event in events if event["type"] == "done") == 1 + + async def test_streaming_ignores_output_format(self) -> None: + config = {**BASE_CONFIG, "outputFormat": {"type": "object"}} + handler = create_google_adk_agents_handler() + async for _ in await handler.stream(config, "hello"): + pass + assert handler_mod.LlmAgent.call_args.kwargs.get("output_schema") is None + + async def test_early_close_abandons_tool_spans( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + abandoned: list[bool] = [] + plugin_cls = handler_mod.LaunchDarklyTelemetryPlugin + + class _SpyPlugin(plugin_cls): + def abandon_open_spans(self, *args: Any, **kwargs: Any) -> None: + abandoned.append(True) + super().abandon_open_spans(*args, **kwargs) + + monkeypatch.setattr(handler_mod, "LaunchDarklyTelemetryPlugin", _SpyPlugin) + handler = create_google_adk_agents_handler() + generator = await handler.stream(BASE_CONFIG, "hello") + await generator.__anext__() + spans_mod.fail_span.reset_mock() + await generator.aclose() + assert abandoned == [True] + spans_mod.fail_span.assert_not_called() diff --git a/packages/google-adk-agents/tests/test_native_graph.py b/packages/google-adk-agents/tests/test_native_graph.py new file mode 100644 index 00000000..934f284d --- /dev/null +++ b/packages/google-adk-agents/tests/test_native_graph.py @@ -0,0 +1,268 @@ +"""Native ADK workflow adapter. Reference: TESTING.md §2.2 and §2.x.4.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any, ClassVar +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import launchdarkly_ai_google_adk_agents.native_graph as native_mod +from launchdarkly_ai_google_adk_agents.native_graph import to_adk_agents + + +class _FakeTool: + created: ClassVar[list[_FakeTool]] = [] + + def __init__(self, func: Any, **kwargs: Any) -> None: + self.func = func + self.name = kwargs.get("name", "") + _FakeTool.created.append(self) + + +class _FakeRunner: + created: ClassVar[list[_FakeRunner]] = [] + + def __init__(self, *args: Any, **kwargs: Any) -> None: + self.plugins = kwargs.get("plugins", []) + self.session_service = SimpleNamespace( + create_session=AsyncMock(return_value=SimpleNamespace(id="sess")), + append_event=AsyncMock(), + ) + _FakeRunner.created.append(self) + + async def run_async(self, **kwargs: Any) -> Any: + yield SimpleNamespace( + content=SimpleNamespace(parts=[SimpleNamespace(text="leaf-answer")]), + partial=False, + usage_metadata=SimpleNamespace( + prompt_token_count=2, + candidates_token_count=3, + total_token_count=5, + ), + is_final_response=lambda: True, + ) + + +def _node( + key: str, *, terminal: bool, edges: list[Any] | None = None +) -> SimpleNamespace: + return SimpleNamespace( + key=key, + config={ + "model": {"name": "gemini-2.5-flash"}, + "provider": {"name": "Google"}, + "instructions": f"instructions-{key}", + }, + meta={"variationKey": "v", "version": 1}, + edges=edges or [], + is_terminal=terminal, + ) + + +def _graph(enabled: bool = True, root: Any = None) -> SimpleNamespace: + leaf = _node("leaf", terminal=True) + root_node = root or _node( + "root", + terminal=False, + edges=[SimpleNamespace(target_key="leaf", key="root-leaf", source_key="root")], + ) + nodes = {"root": root_node, "leaf": leaf} + + def get_node(key: str) -> Any: + return nodes[key] + + return SimpleNamespace( + enabled=enabled, + key="graph-flag", + root=root_node if enabled else None, + get_node=get_node, + ) + + +@pytest.fixture(autouse=True) +def _mocks(monkeypatch: pytest.MonkeyPatch) -> None: + _FakeTool.created = [] + _FakeRunner.created = [] + monkeypatch.setattr( + native_mod, + "LlmAgent", + MagicMock(side_effect=lambda **kw: SimpleNamespace(**kw)), + ) + monkeypatch.setattr(native_mod, "FunctionTool", _FakeTool) + monkeypatch.setattr( + native_mod, + "Workflow", + MagicMock(side_effect=lambda **kw: SimpleNamespace(**kw)), + ) + monkeypatch.setattr(native_mod, "InMemoryRunner", _FakeRunner) + monkeypatch.setattr(native_mod, "START", "START") + + +class TestToAdkAgents: + async def test_disabled_graph_throws(self) -> None: + with pytest.raises(Exception, match="graph-flag"): + await to_adk_agents(_graph(enabled=False)).invoke("hi") + + async def test_missing_root_throws(self) -> None: + graph = _graph() + graph.root = None + with pytest.raises(Exception, match="root"): + await to_adk_agents(graph).invoke("hi") + + async def test_two_node_graph_builds_two_agents_and_a_transfer_tool(self) -> None: + result = await to_adk_agents(_graph()).invoke("hi") + assert native_mod.LlmAgent.call_count == 2 + instructions = [ + call.kwargs["instruction"] for call in native_mod.LlmAgent.call_args_list + ] + assert instructions == ["instructions-root", "instructions-leaf"] or set( + instructions + ) == { + "instructions-root", + "instructions-leaf", + } + names = [tool.name for tool in _FakeTool.created] + assert names.count("transfer_to_leaf") == 1 + assert not any( + name.startswith("transfer_to_") and name != "transfer_to_leaf" + for name in names + ) + edges = native_mod.Workflow.call_args.kwargs["edges"] + assert edges[0][0] == "START" + assert result["response"] == "leaf-answer" + assert result["usage"]["total"] == 5 + + async def test_transfer_tool_selects_the_target(self) -> None: + await to_adk_agents(_graph()).invoke("hi") + tool = next( + item for item in _FakeTool.created if item.name == "transfer_to_leaf" + ) + selected = tool.func() + assert "leaf" in str(selected) + + async def test_no_context_does_not_track(self) -> None: + tracker = MagicMock() + with pytest.MonkeyPatch.context() as patch: + patch.setattr(native_mod, "get_client", lambda: tracker) + await to_adk_agents(_graph()).invoke("hi") + tracker.track.assert_not_called() + + async def test_history_is_seeded_on_the_root_session_only(self) -> None: + await to_adk_agents(_graph()).invoke( + "hi", history=[{"role": "user", "content": "earlier"}] + ) + assert _FakeRunner.created[0].session_service.append_event.await_count == 1 + + async def test_node_local_tools_are_forwarded(self) -> None: + graph = _graph() + graph.root.config["tools"] = { + "lookup": { + "name": "lookup", + "description": "find", + "parameters": {"type": "object"}, + } + } + + def lookup(args: dict[str, Any]) -> str: + return str(args["q"]) + + await to_adk_agents(graph, tool_handlers={"lookup": lookup}).invoke("hi") + names = [tool.name for tool in _FakeTool.created] + assert "lookup" in names + assert names.count("transfer_to_leaf") == 1 + + async def test_transfer_during_the_run_visits_the_target( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + class _HandoffRunner(_FakeRunner): + hops: ClassVar[int] = 0 + + async def run_async(self, **kwargs: Any) -> Any: + if _HandoffRunner.hops == 0: + tool = next( + item + for item in _FakeTool.created + if item.name == "transfer_to_leaf" + ) + tool.func() + _HandoffRunner.hops += 1 + yield SimpleNamespace( + content=SimpleNamespace(parts=[SimpleNamespace(text="root")]), + partial=False, + usage_metadata=SimpleNamespace( + prompt_token_count=1, + candidates_token_count=1, + total_token_count=2, + ), + is_final_response=lambda: True, + ) + return + yield SimpleNamespace( + content=SimpleNamespace( + parts=[SimpleNamespace(text="leaf-answer")] + ), + partial=False, + usage_metadata=SimpleNamespace( + prompt_token_count=2, + candidates_token_count=3, + total_token_count=5, + ), + is_final_response=lambda: True, + ) + + _HandoffRunner.hops = 0 + monkeypatch.setattr(native_mod, "InMemoryRunner", _HandoffRunner) + result = await to_adk_agents(_graph()).invoke("hi") + assert _HandoffRunner.hops == 1 + assert len(_FakeRunner.created) == 2 + assert result["response"] == "leaf-answer" + assert result["usage"]["total"] == 7 + + async def test_context_emits_graph_span_and_events(self) -> None: + tracker = MagicMock() + span = MagicMock() + tracer = MagicMock() + tracer.start_span.return_value = span + with pytest.MonkeyPatch.context() as patch: + patch.setattr(native_mod, "get_client", lambda: tracker) + patch.setattr( + native_mod, "trace", SimpleNamespace(get_tracer=lambda _name: tracer) + ) + await to_adk_agents(_graph()).invoke( + "hi", context={"kind": "user", "key": "user-1"} + ) + assert tracer.start_span.call_args.args[0] == "ld.ai.graph" + events = [call.args[0] for call in tracker.track.call_args_list] + assert "$ld:ai:graph:invocation_success" in events + assert "$ld:ai:graph:duration:total" in events + assert "$ld:ai:graph:total_tokens" in events + span.end.assert_called() + + async def test_runner_error_tracks_invocation_failure_and_ends_the_span( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + class _Boom(_FakeRunner): + async def run_async(self, **kwargs: Any) -> Any: + if False: + yield None + raise RuntimeError("runner broke") + + tracker = MagicMock() + span = MagicMock() + tracer = MagicMock() + tracer.start_span.return_value = span + monkeypatch.setattr(native_mod, "InMemoryRunner", _Boom) + with pytest.MonkeyPatch.context() as patch: + patch.setattr(native_mod, "get_client", lambda: tracker) + patch.setattr( + native_mod, "trace", SimpleNamespace(get_tracer=lambda _name: tracer) + ) + with pytest.raises(RuntimeError, match="runner broke"): + await to_adk_agents(_graph()).invoke( + "hi", context={"kind": "user", "key": "user-1"} + ) + events = [call.args[0] for call in tracker.track.call_args_list] + assert events == ["$ld:ai:graph:invocation_failure"] + span.end.assert_called() diff --git a/packages/google-adk-agents/tests/test_spans.py b/packages/google-adk-agents/tests/test_spans.py new file mode 100644 index 00000000..0a7b4357 --- /dev/null +++ b/packages/google-adk-agents/tests/test_spans.py @@ -0,0 +1,98 @@ +"""Span helpers for the ADK telemetry plugin. Reference: TESTING.md §2.x.3.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock + +import launchdarkly_ai_google_adk_agents.spans as spans_mod +from launchdarkly_ai_google_adk_agents.spans import ( + abandon_open_spans, + start_model_span, + start_root_span, + start_tool_span, +) + +CONFIG = { + "model": {"name": "gemini-2.5-flash"}, + "provider": {"name": "Google"}, + "instructions": "Be helpful.", +} +VARIABLES = { + "__ld": { + "configKey": "cfg", + "variationKey": "var", + "runId": "run-1", + }, + "ldContext": {"kind": "user", "key": "user-1"}, +} + + +class _Span: + def __init__(self, name: str) -> None: + self.name = name + self.attributes: dict[str, Any] = {} + self.ended = False + self.failed = False + + def set_attribute(self, key: str, value: Any) -> None: + self.attributes[key] = value + + def add_event(self, name: str, attributes: dict[str, Any] | None = None) -> None: + self.attributes.setdefault("events", []).append(name) + + def end(self) -> None: + self.ended = True + + def set_status(self, *args: Any, **kwargs: Any) -> None: + return None + + def record_exception(self, exc: BaseException) -> None: + self.failed = True + + +def test_root_span_identity(monkeypatch: Any) -> None: + spans: list[_Span] = [] + + def start_span(name: str, **kwargs: Any) -> _Span: + span = _Span(name) + spans.append(span) + return span + + tracer = SimpleNamespace(start_span=start_span) + monkeypatch.setattr(spans_mod.trace, "get_tracer", lambda name: tracer) + root = start_root_span(CONFIG, VARIABLES) + assert root is spans[0] + assert root.name == "invoke_agent" + assert root.attributes["gen_ai.system"] == "google_adk" + assert root.attributes["gen_ai.provider.name"] == "gcp.gemini" + assert root.attributes["gen_ai.request.model"] == "gemini-2.5-flash" + assert root.attributes["launchdarkly.config.key"] == "cfg" + assert "feature_flag" in root.attributes["events"] + + +def test_child_spans_do_not_repeat_launchdarkly_identity(monkeypatch: Any) -> None: + def start_span(name: str, **kwargs: Any) -> _Span: + return _Span(name) + + monkeypatch.setattr( + spans_mod.trace, + "get_tracer", + lambda name: SimpleNamespace(start_span=start_span), + ) + model = start_model_span(CONFIG, parent=MagicMock()) + tool = start_tool_span("lookup", "call-1", parent=MagicMock()) + assert model.name == "chat gemini-2.5-flash" + assert tool.name == "execute_tool lookup" + assert "launchdarkly.config.key" not in model.attributes + assert "launchdarkly.config.key" not in tool.attributes + + +def test_abandon_does_not_fail_the_tool_span(monkeypatch: Any) -> None: + span = _Span("execute_tool lookup") + fail = MagicMock() + monkeypatch.setattr(spans_mod, "fail_span", fail) + abandon_open_spans([span], ended=set()) + assert span.ended is True + fail.assert_not_called() diff --git a/pyproject.toml b/pyproject.toml index 83e3de63..63b046cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,14 @@ pythonpath = ["."] strict = true python_version = "3.12" +# google-adk 2.9.2 caps opentelemetry-sdk<=1.42.1, which this workspace cannot +# solve next to opentelemetry-exporter-otlp-proto-http>=1.43. The handler +# imports ADK lazily and unit tests patch those names. Install google-adk in +# the application environment; typechecking here does not require it. +[[tool.mypy.overrides]] +module = ["google", "google.*"] +ignore_missing_imports = true + [tool.ruff] target-version = "py312" line-length = 88 @@ -61,6 +69,7 @@ known-first-party = [ "launchdarkly_ai_server", "launchdarkly_ai_claude_agents", "launchdarkly_ai_claude_messages", + "launchdarkly_ai_google_adk_agents", "launchdarkly_ai_langchain_agents", "launchdarkly_ai_langchain_messages", "launchdarkly_ai_openai_agents", diff --git a/release-please-config.json b/release-please-config.json index 9d2a04e3..eb389e99 100644 --- a/release-please-config.json +++ b/release-please-config.json @@ -79,6 +79,17 @@ ], "component": "launchdarkly-ai-langchain-agents" }, + "packages/google-adk-agents": { + "release-type": "python", + "versioning": "default", + "bump-minor-pre-major": true, + "bump-patch-for-minor-pre-major": true, + "include-v-in-tag": false, + "extra-files": [ + "src/launchdarkly_ai_google_adk_agents/__init__.py" + ], + "component": "launchdarkly-ai-google-adk-agents" + }, "packages/langchain-messages": { "release-type": "python", "versioning": "default", diff --git a/tests/test_cross_handler_parity.py b/tests/test_cross_handler_parity.py index 3ab7513a..6cd6180a 100644 --- a/tests/test_cross_handler_parity.py +++ b/tests/test_cross_handler_parity.py @@ -1,11 +1,11 @@ -"""Cross-handler invariants: the six handlers must agree with each other. +"""Cross-handler invariants: the handlers must agree with each other. -Every handler package tests its own spans. Nothing tested that the six agree, and that is exactly +Every handler package tests its own spans. Nothing tested that they agree, and that is exactly how they drifted apart: each was correct on its own terms while a single run emitted `chat` spans that disagreed about what a finish reason or a cached token was. These tests are the oracle for that. They live outside the packages because no package can own an -invariant about all six. +invariant about every handler. There are two kinds of check here. @@ -41,6 +41,7 @@ "openai-agents": "launchdarkly_ai_openai_agents.spans", "langchain-messages": "launchdarkly_ai_langchain_messages.spans", "langchain-agents": "launchdarkly_ai_langchain_agents.spans", + "google-adk-agents": "launchdarkly_ai_google_adk_agents.spans", } #: `claude-agents` builds its `chat` span inside an inference tracker rather than in a standalone @@ -519,6 +520,7 @@ def test_the_carrier_stays_behind_the_capture_gate(self) -> None: "claude-agents": "launchdarkly_ai_claude_agents.handler", "openai-agents": "launchdarkly_ai_openai_agents.handler", "langchain-agents": "launchdarkly_ai_langchain_agents.handler", + "google-adk-agents": "launchdarkly_ai_google_adk_agents.handler", } diff --git a/uv.lock b/uv.lock index c1c96777..11142f69 100644 --- a/uv.lock +++ b/uv.lock @@ -15,6 +15,7 @@ resolution-markers = [ members = [ "launchdarkly-ai-claude-agents", "launchdarkly-ai-claude-messages", + "launchdarkly-ai-google-adk-agents", "launchdarkly-ai-langchain-agents", "launchdarkly-ai-langchain-messages", "launchdarkly-ai-openai-agents", @@ -790,7 +791,7 @@ wheels = [ [[package]] name = "launchdarkly-ai-claude-agents" -version = "0.2.2" +version = "0.2.3" source = { editable = "packages/claude-agents" } dependencies = [ { name = "anthropic" }, @@ -809,7 +810,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-claude-messages" -version = "0.2.2" +version = "0.2.3" source = { editable = "packages/claude-messages" } dependencies = [ { name = "anthropic" }, @@ -824,9 +825,24 @@ requires-dist = [ { name = "opentelemetry-api", specifier = ">=1.25" }, ] +[[package]] +name = "launchdarkly-ai-google-adk-agents" +version = "0.1.0" +source = { editable = "packages/google-adk-agents" } +dependencies = [ + { name = "launchdarkly-ai-server" }, + { name = "opentelemetry-api" }, +] + +[package.metadata] +requires-dist = [ + { name = "launchdarkly-ai-server", editable = "packages/client" }, + { name = "opentelemetry-api", specifier = ">=1.25" }, +] + [[package]] name = "launchdarkly-ai-langchain-agents" -version = "0.2.2" +version = "0.2.3" source = { editable = "packages/langchain-agents" } dependencies = [ { name = "langchain-core" }, @@ -845,7 +861,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-langchain-messages" -version = "0.2.2" +version = "0.2.3" source = { editable = "packages/langchain-messages" } dependencies = [ { name = "langchain-core" }, @@ -862,7 +878,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-openai-agents" -version = "0.2.2" +version = "0.2.3" source = { editable = "packages/openai-agents" } dependencies = [ { name = "launchdarkly-ai-server" }, @@ -881,7 +897,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-openai-messages" -version = "0.2.2" +version = "0.2.3" source = { editable = "packages/openai-messages" } dependencies = [ { name = "launchdarkly-ai-server" }, @@ -898,7 +914,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-python" -version = "0.1.6" +version = "0.1.7" source = { editable = "packages/ai" } dependencies = [ { name = "launchdarkly-ai-server" }, @@ -918,7 +934,7 @@ provides-extras = ["otel"] [[package]] name = "launchdarkly-ai-server" -version = "0.2.2" +version = "0.2.3" source = { editable = "packages/client" } dependencies = [ { name = "opentelemetry-api" }, From 7a0a889b77c968d68008a43800c174ef429bbfdb Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Fri, 25 Sep 2026 12:38:17 -0400 Subject: [PATCH 2/2] refactor(google-adk-agents): rename the graph span to launchdarkly.graph --- .../src/launchdarkly_ai_google_adk_agents/native_graph.py | 6 +++--- packages/google-adk-agents/tests/test_native_graph.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/native_graph.py b/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/native_graph.py index 02a36631..bd550564 100644 --- a/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/native_graph.py +++ b/packages/google-adk-agents/src/launchdarkly_ai_google_adk_agents/native_graph.py @@ -105,7 +105,7 @@ async def invoke( ) current = graph_def.get_node(target) if target else None if span is not None: - span.set_attribute("ld.ai.graph.path", "->".join(path)) + span.set_attribute("launchdarkly.graph.path", "->".join(path)) if SpanStatusCode is not None: span.set_status(SpanStatusCode.OK) if ld_context is not None: @@ -145,8 +145,8 @@ def to_adk_agents(graph_def: Any, **options: Any) -> _AdkAgents: def _graph_span(key: str) -> Any: if trace is None: return None - span = trace.get_tracer(_TRACER).start_span("ld.ai.graph") - span.set_attribute("ld.ai.graph.key", key) + span = trace.get_tracer(_TRACER).start_span("launchdarkly.graph") + span.set_attribute("launchdarkly.graph.key", key) return span diff --git a/packages/google-adk-agents/tests/test_native_graph.py b/packages/google-adk-agents/tests/test_native_graph.py index 934f284d..b3657012 100644 --- a/packages/google-adk-agents/tests/test_native_graph.py +++ b/packages/google-adk-agents/tests/test_native_graph.py @@ -233,7 +233,7 @@ async def test_context_emits_graph_span_and_events(self) -> None: await to_adk_agents(_graph()).invoke( "hi", context={"kind": "user", "key": "user-1"} ) - assert tracer.start_span.call_args.args[0] == "ld.ai.graph" + assert tracer.start_span.call_args.args[0] == "launchdarkly.graph" events = [call.args[0] for call in tracker.track.call_args_list] assert "$ld:ai:graph:invocation_success" in events assert "$ld:ai:graph:duration:total" in events