-
Notifications
You must be signed in to change notification settings - Fork 116
AI-472 Add replay-safe Google ADK metrics sample #355
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
9c39154
a058b20
0b92811
e2e0254
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,33 @@ | ||
| # Google ADK replay-safe metrics | ||
|
|
||
| This sample exports Google ADK's OpenTelemetry metrics to a local Prometheus endpoint while preventing Workflow replay from recording the same observations again. The default scripted model is deterministic and makes no network model calls, so no API key is needed. | ||
|
|
||
| Start a local Temporal development server: | ||
|
|
||
| ```shell | ||
| temporal server start-dev | ||
| ``` | ||
|
|
||
| In another terminal, start the worker from the repository root: | ||
|
|
||
| ```shell | ||
| uv run python -m google_adk_agents.metrics.run_worker | ||
| ``` | ||
|
|
||
| Then run the Workflow: | ||
|
|
||
| ```shell | ||
| uv run python -m google_adk_agents.metrics.run_metrics_workflow | ||
| ``` | ||
|
|
||
| The starter prints `Replay-safe metrics are ready.` Inspect the metrics exposed by the worker: | ||
|
|
||
| ```shell | ||
| curl -s http://127.0.0.1:9464/metrics | grep gen_ai | ||
| ``` | ||
|
|
||
| The output includes `gen_ai.invoke_agent`, `gen_ai.client.operation.duration`, and `gen_ai.client.token.usage` metrics. Prometheus replaces dots with underscores, so an exported line looks like `gen_ai_invoke_agent_duration_seconds_count{gen_ai_agent_name="metrics_agent"} 1.0`. `ReplaySafeMeterProvider` drops observations made while replaying, so replay does not multiply the recorded counts. | ||
|
|
||
| Recordings are first-execution-only rather than exactly-once. Replay is suppressed, but a Workflow Task retry re-executes live and can record again, so treat these metrics as at-least-once usage signals. | ||
|
|
||
| OpenTelemetry's global meter provider can be installed only once per process. `run_worker.py` installs the replay-safe provider before importing Google ADK or the Workflow. Applications embedding this setup must likewise make it the first and only global meter provider installation in that process. | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1 @@ | ||
|
|
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1 @@ | ||
|
|
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| @@ -0,0 +1,33 @@ | ||||||||||||||||||||||
| from collections.abc import AsyncGenerator | ||||||||||||||||||||||
|
|
||||||||||||||||||||||
| from google.adk.models import BaseLlm | ||||||||||||||||||||||
| from google.adk.models.llm_request import LlmRequest | ||||||||||||||||||||||
| from google.adk.models.llm_response import LlmResponse | ||||||||||||||||||||||
| from google.genai import types | ||||||||||||||||||||||
|
|
||||||||||||||||||||||
| MODEL_NAME = "local-metrics-model" | ||||||||||||||||||||||
|
|
||||||||||||||||||||||
|
|
||||||||||||||||||||||
| class LocalMetricsModel(BaseLlm): | ||||||||||||||||||||||
| @classmethod | ||||||||||||||||||||||
| def supported_models(cls) -> list[str]: | ||||||||||||||||||||||
| return [MODEL_NAME] | ||||||||||||||||||||||
|
|
||||||||||||||||||||||
| async def generate_content_async( | ||||||||||||||||||||||
| self, llm_request: LlmRequest, stream: bool = False | ||||||||||||||||||||||
| ) -> AsyncGenerator[LlmResponse, None]: | ||||||||||||||||||||||
|
Comment on lines
+16
to
+18
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
Minor, separately:
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Done. Streaming now fails explicitly, and the workflow imports MODEL_NAME. |
||||||||||||||||||||||
| if stream: | ||||||||||||||||||||||
| raise NotImplementedError( | ||||||||||||||||||||||
| "LocalMetricsModel does not implement streaming responses." | ||||||||||||||||||||||
| ) | ||||||||||||||||||||||
| yield LlmResponse( | ||||||||||||||||||||||
| content=types.Content( | ||||||||||||||||||||||
| role="model", | ||||||||||||||||||||||
| parts=[types.Part(text="Replay-safe metrics are ready.")], | ||||||||||||||||||||||
| ), | ||||||||||||||||||||||
| usage_metadata=types.GenerateContentResponseUsageMetadata( | ||||||||||||||||||||||
| prompt_token_count=8, | ||||||||||||||||||||||
| candidates_token_count=5, | ||||||||||||||||||||||
| total_token_count=13, | ||||||||||||||||||||||
| ), | ||||||||||||||||||||||
| ) | ||||||||||||||||||||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,21 @@ | ||
| import asyncio | ||
|
|
||
| from temporalio.client import Client | ||
| from temporalio.contrib.google_adk_agents import GoogleAdkPlugin | ||
|
|
||
| from google_adk_agents.metrics.workflows.metrics_workflow import MetricsWorkflow | ||
|
|
||
|
|
||
| async def main() -> None: | ||
| client = await Client.connect("localhost:7233", plugins=[GoogleAdkPlugin()]) | ||
| result = await client.execute_workflow( | ||
| MetricsWorkflow.run, | ||
| "Explain replay-safe metrics.", | ||
| id="google-adk-agents-metrics-workflow-id", | ||
| task_queue="google-adk-agents-metrics", | ||
| ) | ||
| print(f"Result: {result}") | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| asyncio.run(main()) |
| Original file line number | Diff line number | Diff line change | |||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| @@ -0,0 +1,33 @@ | |||||||||||||||||||||||||||||||||
| import asyncio | |||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||
| from opentelemetry.exporter.prometheus import PrometheusMetricReader | |||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||
| from google_adk_agents.metrics.telemetry import install_meter_provider | |||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||
| async def main() -> None: | |||||||||||||||||||||||||||||||||
| install_meter_provider(PrometheusMetricReader()) | |||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||
| from google.adk.models import LLMRegistry | |||||||||||||||||||||||||||||||||
| from prometheus_client import start_http_server | |||||||||||||||||||||||||||||||||
| from temporalio.client import Client | |||||||||||||||||||||||||||||||||
| from temporalio.contrib.google_adk_agents import GoogleAdkPlugin | |||||||||||||||||||||||||||||||||
| from temporalio.worker import Worker | |||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||
| from google_adk_agents.metrics.models.local_metrics_model import LocalMetricsModel | |||||||||||||||||||||||||||||||||
| from google_adk_agents.metrics.workflows.metrics_workflow import MetricsWorkflow | |||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||
| LLMRegistry.register(LocalMetricsModel) | |||||||||||||||||||||||||||||||||
| start_http_server(port=9464, addr="127.0.0.1") | |||||||||||||||||||||||||||||||||
| plugin = GoogleAdkPlugin() | |||||||||||||||||||||||||||||||||
| client = await Client.connect("localhost:7233", plugins=[plugin]) | |||||||||||||||||||||||||||||||||
| worker = Worker( | |||||||||||||||||||||||||||||||||
| client, | |||||||||||||||||||||||||||||||||
| task_queue="google-adk-agents-metrics", | |||||||||||||||||||||||||||||||||
| workflows=[MetricsWorkflow], | |||||||||||||||||||||||||||||||||
| ) | |||||||||||||||||||||||||||||||||
|
Comment on lines
+24
to
+28
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
191x and 240x under-report, in the exact configuration a reader copies. It's also the only non-test worker in samples-python that disables the cache — The head commit's
Suggested change
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Done. Removed max_cached_workflows from the shipped worker. |
|||||||||||||||||||||||||||||||||
| await worker.run() | |||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||
| if __name__ == "__main__": | |||||||||||||||||||||||||||||||||
| asyncio.run(main()) | |||||||||||||||||||||||||||||||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,12 @@ | ||
| import opentelemetry.metrics | ||
| from opentelemetry.sdk.metrics import MeterProvider | ||
| from opentelemetry.sdk.metrics.export import MetricReader | ||
| from temporalio.contrib.opentelemetry import ReplaySafeMeterProvider | ||
|
|
||
|
|
||
| def install_meter_provider(reader: MetricReader) -> ReplaySafeMeterProvider: | ||
| provider = ReplaySafeMeterProvider(MeterProvider(metric_readers=[reader])) | ||
| opentelemetry.metrics.set_meter_provider(provider) | ||
| if opentelemetry.metrics.get_meter_provider() is not provider: | ||
| raise RuntimeError("The global OpenTelemetry meter provider is already set") | ||
| return provider |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1 @@ | ||
|
|
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,38 @@ | ||
| from google.adk import Agent | ||
| from google.adk.runners import InMemoryRunner | ||
| from google.adk.utils.context_utils import Aclosing | ||
| from google.genai import types | ||
| from temporalio import workflow | ||
| from temporalio.contrib.google_adk_agents import TemporalModel | ||
|
|
||
| from google_adk_agents.metrics.models.local_metrics_model import MODEL_NAME | ||
|
|
||
|
|
||
| @workflow.defn | ||
| class MetricsWorkflow: | ||
| @workflow.run | ||
| async def run(self, prompt: str) -> str: | ||
| agent = Agent( | ||
| name="metrics_agent", | ||
| model=TemporalModel(MODEL_NAME), | ||
| instruction="Answer the user briefly.", | ||
| ) | ||
| runner = InMemoryRunner(agent=agent, app_name="metrics_app") | ||
| session = await runner.session_service.create_session( | ||
| app_name="metrics_app", user_id="sample-user" | ||
| ) | ||
|
|
||
| final_text = "" | ||
| async with Aclosing( | ||
| runner.run_async( | ||
| user_id="sample-user", | ||
| session_id=session.id, | ||
| new_message=types.Content(role="user", parts=[types.Part(text=prompt)]), | ||
| ) | ||
| ) as events: | ||
| async for event in events: | ||
| if event.content and event.content.parts: | ||
| for part in event.content.parts: | ||
| if part.text: | ||
| final_text = part.text | ||
| return final_text |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,83 @@ | ||
| import uuid | ||
|
|
||
| import pytest | ||
| from google.adk.models import BaseLlm, LLMRegistry | ||
| from opentelemetry.sdk.metrics.export import HistogramDataPoint, InMemoryMetricReader | ||
| from temporalio.client import Client | ||
| from temporalio.contrib.google_adk_agents import GoogleAdkPlugin | ||
| from temporalio.worker import Replayer, Worker | ||
|
|
||
| from google_adk_agents.metrics.models.local_metrics_model import ( | ||
| MODEL_NAME, | ||
| LocalMetricsModel, | ||
| ) | ||
| from google_adk_agents.metrics.telemetry import install_meter_provider | ||
| from google_adk_agents.metrics.workflows.metrics_workflow import MetricsWorkflow | ||
|
|
||
| ADK_METER_SCOPE = "gcp.vertex.agent" | ||
|
|
||
|
|
||
| async def test_metrics_are_not_inflated_by_replay( | ||
| client: Client, monkeypatch: pytest.MonkeyPatch | ||
| ) -> None: | ||
| reader = InMemoryMetricReader() | ||
| install_meter_provider(reader) | ||
|
|
||
| original_new_llm = LLMRegistry.new_llm | ||
|
|
||
| def new_llm(model: str) -> BaseLlm: | ||
| if model == MODEL_NAME: | ||
| return LocalMetricsModel(model=model) | ||
| return original_new_llm(model) | ||
|
|
||
| monkeypatch.setattr(LLMRegistry, "new_llm", staticmethod(new_llm)) | ||
|
|
||
| plugin = GoogleAdkPlugin() | ||
| config = client.config() | ||
| config["plugins"] = [*config["plugins"], plugin] | ||
| client = Client(**config) | ||
| task_queue = f"google-adk-agents-metrics-{uuid.uuid4()}" | ||
| async with Worker( | ||
| client, | ||
| task_queue=task_queue, | ||
| workflows=[MetricsWorkflow], | ||
| ): | ||
| handle = await client.start_workflow( | ||
| MetricsWorkflow.run, | ||
| "Explain replay-safe metrics.", | ||
| id=f"google-adk-agents-metrics-{uuid.uuid4()}", | ||
| task_queue=task_queue, | ||
| ) | ||
| result = await handle.result() | ||
| history = await handle.fetch_history() | ||
|
|
||
| assert result == "Replay-safe metrics are ready." | ||
| counts_before_replay = metric_counts(reader) | ||
| assert counts_before_replay["gen_ai.invoke_agent.duration"] > 0 | ||
| assert counts_before_replay["gen_ai.invoke_agent.inference_calls"] > 0 | ||
| assert counts_before_replay["gen_ai.client.operation.duration"] > 0 | ||
| assert counts_before_replay["gen_ai.client.token.usage"] > 0 | ||
|
|
||
| await Replayer(workflows=[MetricsWorkflow], plugins=[plugin]).replay_workflow( | ||
| history | ||
| ) | ||
|
|
||
| assert metric_counts(reader) == counts_before_replay | ||
|
|
||
|
|
||
| def metric_counts(reader: InMemoryMetricReader) -> dict[str, int]: | ||
| counts: dict[str, int] = {} | ||
| data = reader.get_metrics_data() | ||
| if data is not None: | ||
| for resource_metrics in data.resource_metrics: | ||
| for scope_metrics in resource_metrics.scope_metrics: | ||
| if scope_metrics.scope.name != ADK_METER_SCOPE: | ||
| continue | ||
| for metric in scope_metrics.metrics: | ||
| count = 0 | ||
| for point in metric.data.data_points: | ||
| if not isinstance(point, HistogramDataPoint): | ||
| raise TypeError(f"Unexpected metric point: {type(point)}") | ||
| count += point.count | ||
| counts[metric.name] = count | ||
| return counts |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Companion to the
pyproject.toml:8comment — once the nested project is gone,--projecthas nothing to point at, and the suite README already documents the house form at:48-53.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.