Repository: comet-ml/opik
Pull request: [OPIK-1501]: [SDK] Add cost and token tracking to ADK integration
## Details
Implemented support token usage tracking in ADK integration.

## Issues

Resolves #1955

## Testing

## Documentation


Problem statement:
[FR]: Add support for cost and token tracking in the ADK integration
### Proposal summary

When using Google's Agent Development Kit, costs or token usage are not logged to the platform. Cost and token tracking is supported in many other integrations, would be good to have it supported in ADK as well.

### Motivation

_No response_

Hints from issue discussion:
Hi @jverre, thank you for your request. We will put it on top of our list.
Supporting this use case. This is currently a blocker for me to use Opik in our product, since we moved away from Google generative-ai to google-adk for better support for multi-agent development.

Commit message: [OPIK-1501]: Started implementation of the ADK library integration test.

Patch to review:
diff --git a/.github/workflows/lib-adk-tests.yml b/.github/workflows/lib-adk-tests.yml
new file mode 100644
index 0000000000..3b6a1b87b1
--- /dev/null
+++ b/.github/workflows/lib-adk-tests.yml
@@ -0,0 +1,61 @@
+# Workflow to run ADK tests
+#
+# Please read inputs to provide correct values.
+#
+name: SDK Lib ADK (Google) Tests
+run-name: "SDK Lib ADK Tests ${{ github.ref_name }} by @${{ github.actor }}"
+env:
+  GCP_CREDENTIALS_JSON: ${{ secrets.GCP_CREDENTIALS_JSON }}
+  GOOGLE_CLOUD_LOCATION: us-east1
+  GOOGLE_CLOUD_PROJECT: opik-sdk-tests
+  GOOGLE_GENAI_USE_VERTEXAI: TRUE
+  OPIK_ENABLE_LITELLM_MODELS_MONITORING: False
+  OPIK_SENTRY_ENABLE: False
+on:
+  workflow_call:
+  pull_request:
+
+jobs:
+  tests:
+    name: ADK Python ${{matrix.python_version}}
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: sdks/python
+
+    strategy:
+      fail-fast: true
+      matrix:
+        python_version: [
+          "3.9",
+          "3.10",
+          "3.11",
+          "3.12"
+        ]
+
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+
+      - name: Setup Python ${{matrix.python_version}}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{matrix.python_version}}
+
+      - name: Install opik
+        run: pip install .
+
+      - name: Install test tools
+        run: |
+          cd ./tests
+          pip install --no-cache-dir --disable-pip-version-check -r test_requirements.txt
+
+      - name: Install lib
+        run: |
+          cd ./tests
+          pip install --no-cache-dir --disable-pip-version-check -r library_integration/adk/requirements.txt
+
+      - name: Run tests
+        run: |
+          cd ./tests/library_integration/adk/
+          python -m pytest  -vv .
\ No newline at end of file
diff --git a/.github/workflows/lib-integration-tests-runner.yml b/.github/workflows/lib-integration-tests-runner.yml
index db429c4e5b..b0e3b2fe9e 100644
--- a/.github/workflows/lib-integration-tests-runner.yml
+++ b/.github/workflows/lib-integration-tests-runner.yml
@@ -21,6 +21,7 @@ on:
           - dspy
           - crewai
           - genai
+          - adk
   schedule:
     - cron: "0 0 */1 * *"
   pull_request:
@@ -110,3 +111,8 @@ jobs:
     uses: ./.github/workflows/lib-genai-tests.yml
     secrets: inherit
 
+  adk_tests:
+    needs: [ init_environment ]
+    if: contains(fromJSON('["adk", "all"]'), needs.init_environment.outputs.LIBS)
+    uses: ./.github/workflows/lib-adk-tests.yml
+    secrets: inherit
diff --git a/sdks/python/src/opik/integrations/adk/decorators.py b/sdks/python/src/opik/integrations/adk/adk_decorators.py
similarity index 85%
rename from sdks/python/src/opik/integrations/adk/decorators.py
rename to sdks/python/src/opik/integrations/adk/adk_decorators.py
index 83546a7675..555343e6e3 100644
--- a/sdks/python/src/opik/integrations/adk/decorators.py
+++ b/sdks/python/src/opik/integrations/adk/adk_decorators.py
@@ -13,6 +13,8 @@
 from opik.api_objects import span
 from opik.decorator import arguments_helpers, base_track_decorator
 
+from . import llm_response_wrapper
+
 
 def convert_adk_base_models(arg: Any) -> Dict[str, Any]:
     """Most ADK objects are Pydantic Base Models"""
@@ -48,9 +50,19 @@ def _end_span_inputs_preprocessor(
         capture_output: bool,
         current_span_data: span.SpanData,
     ) -> arguments_helpers.EndSpanParameters:
-        result = arguments_helpers.EndSpanParameters(
-            output=convert_adk_base_models(output),
-        )
+        result_dict = convert_adk_base_models(output)
+        usage_data = llm_response_wrapper.pop_llm_usage_data(**result_dict)
+
+        if usage_data is not None:
+            result = arguments_helpers.EndSpanParameters(
+                output=result_dict,
+                usage=usage_data.opik_usage,
+                model=usage_data.model,
+            )
+        else:
+            result = arguments_helpers.EndSpanParameters(
+                output=result_dict,
+            )
 
         return result
 
diff --git a/sdks/python/src/opik/integrations/adk/llm_response_wrapper.py b/sdks/python/src/opik/integrations/adk/llm_response_wrapper.py
new file mode 100644
index 0000000000..2dd462aa64
--- /dev/null
+++ b/sdks/python/src/opik/integrations/adk/llm_response_wrapper.py
@@ -0,0 +1,81 @@
+import dataclasses
+import logging
+from typing import Callable, Optional, Dict, Any
+
+from google.adk import models as adk_models
+from google.genai import types as genai_types
+
+from ... import llm_usage, LLMProvider
+from ...llm_usage import opik_usage
+
+
+LOGGER = logging.Logger(__name__)
+
+
+class LlmResponseCreateWrapper:
+    def __init__(
+        self,
+        wrapped_response_create: Callable[
+            [genai_types.GenerateContentResponse], adk_models.LlmResponse
+        ],
+    ) -> None:
+        self.wrapped_response_create = wrapped_response_create
+
+    def __call__(
+        self, generate_content_response: genai_types.GenerateContentResponse
+    ) -> adk_models.LlmResponse:
+        return _wrap_llm_response_create(
+            generate_content_response, self.wrapped_response_create
+        )
+
+
+@dataclasses.dataclass
+class LLMUsageData:
+    opik_usage: opik_usage.OpikUsage
+    model: Optional[str]
+    provider: Optional[str]
+
+
+def pop_llm_usage_data(**result_dict: Dict[str, Any]) -> Optional[LLMUsageData]:
+    """Extracts Opik usage metadata from ADK output and removes it from the result dict."""
+    custom_metadata = result_dict.get("custom_metadata", None)
+    if custom_metadata is None:
+        return None
+
+    opik_usage_metadata = custom_metadata.pop("opik_usage", None)
+    if opik_usage_metadata is None:
+        return None
+
+    model = custom_metadata.pop("model_version", None)
+    provider = custom_metadata.pop("provider", None)
+    usage = llm_usage.try_build_opik_usage_or_log_error(
+        provider=LLMProvider(provider),
+        usage=opik_usage_metadata,
+        logger=LOGGER,
+        error_message="Failed to log token usage from ADK Gemini call",
+    )
+    if usage is None:
+        return None
+
+    return LLMUsageData(opik_usage=usage, model=model, provider=provider)
+
+
+def _wrap_llm_response_create(
+    generate_content_response: genai_types.GenerateContentResponse,
+    wrapped_response_create: Callable[
+        [genai_types.GenerateContentResponse], adk_models.LlmResponse
+    ],
+) -> adk_models.LlmResponse:
+    usage_metadata = generate_content_response.usage_metadata
+    response = wrapped_response_create(generate_content_response)
+    if usage_metadata is None:
+        return response
+
+    if response.custom_metadata is None:
+        response.custom_metadata = {}
+
+    response.custom_metadata["opik_usage"] = usage_metadata
+    response.custom_metadata["provider"] = LLMProvider.GOOGLE_AI
+    response.custom_metadata["model_version"] = generate_content_response.model_version
+
+    return response
diff --git a/sdks/python/src/opik/integrations/adk/opik_tracer.py b/sdks/python/src/opik/integrations/adk/opik_tracer.py
index ff437bc25b..df4c720e47 100644
--- a/sdks/python/src/opik/integrations/adk/opik_tracer.py
+++ b/sdks/python/src/opik/integrations/adk/opik_tracer.py
@@ -10,11 +10,8 @@
 from opik.decorator import arguments_helpers
 from opik.types import DistributedTraceHeadersDict
 
-from .decorators import (
-    ADKLLMTrackDecorator,
-    ADKToolTrackDecorator,
-    convert_adk_base_models,
-)
+from . import adk_decorators
+from . import llm_response_wrapper
 
 
 class OpikTracer:
@@ -36,8 +33,13 @@ def __init__(
 
         self._last_model_output: Optional[Dict[str, Any]] = None
 
-        self._llm_tracer = ADKLLMTrackDecorator()
-        self._tool_tracer = ADKToolTrackDecorator()
+        self._llm_tracer = adk_decorators.ADKLLMTrackDecorator()
+        self._tool_tracer = adk_decorators.ADKToolTrackDecorator()
+
+        # monkey patch LLMResponse to store usage_metadata
+        old_function = LlmResponse.create
+        create_wrapper = llm_response_wrapper.LlmResponseCreateWrapper(old_function)
+        LlmResponse.create = create_wrapper
 
     def before_agent_callback(
         self, callback_context: CallbackContext, *args: Any, **kwargs: Any
@@ -53,7 +55,9 @@ def before_agent_callback(
         trace_metadata = self.metadata.copy()
         trace_metadata["adk_invocation_id"] = callback_context.invocation_id
 
-        user_input = convert_adk_base_models(callback_context.user_content)
+        user_input = adk_decorators.convert_adk_base_models(
+            callback_context.user_content
+        )
 
         self.trace_data = trace.TraceData(
             id=helpers.generate_id(),
@@ -78,6 +82,10 @@ def after_agent_callback(
 
         output = self._last_model_output
 
+        # remove the custom metadata with opik usage we added
+        if output is not None:
+            llm_response_wrapper.pop_llm_usage_data(**output)
+
         self.trace_data.update(output=output).init_end_time()
 
         self._client.trace(**self.trace_data.__dict__)
@@ -118,8 +126,7 @@ def after_model_callback(
         *args: Any,
         **kwargs: Any,
     ) -> None:
-        self._last_model_output = convert_adk_base_models(llm_response)
-
+        self._last_model_output = adk_decorators.convert_adk_base_models(llm_response)
         self._llm_tracer._after_call(
             output=llm_response,
             error_info=None,
diff --git a/sdks/python/src/opik/integrations/adk/requirements.txt b/sdks/python/src/opik/integrations/adk/requirements.txt
new file mode 100644
index 0000000000..17f66e4e9c
--- /dev/null
+++ b/sdks/python/src/opik/integrations/adk/requirements.txt
@@ -0,0 +1 @@
+google-adk
diff --git a/sdks/python/tests/library_integration/adk/__init__.py b/sdks/python/tests/library_integration/adk/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/sdks/python/tests/library_integration/adk/sample_agent/__init__.py b/sdks/python/tests/library_integration/adk/sample_agent/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/sdks/python/tests/library_integration/adk/sample_agent/agent.py b/sdks/python/tests/library_integration/adk/sample_agent/agent.py
new file mode 100644
index 0000000000..da797f5c00
--- /dev/null
+++ b/sdks/python/tests/library_integration/adk/sample_agent/agent.py
@@ -0,0 +1,72 @@
+import datetime
+from zoneinfo import ZoneInfo
+from google.adk.agents import Agent
+
+from opik.integrations.adk import OpikTracer
+
+
+def get_weather(city: str) -> dict:
+    """Retrieves the current weather report for a specified city.
+
+    Args:
+        city (str): The name of the city for which to retrieve the weather report.
+
+    Returns:
+        dict: status and result or error msg.
+    """
+    if city.lower() == "new york":
+        return {
+            "status": "success",
+            "report": (
+                "The weather in New York is sunny with a temperature of 25 degrees"
+                " Celsius (77 degrees Fahrenheit)."
+            ),
+        }
+    else:
+        return {
+            "status": "error",
+            "error_message": f"Weather information for '{city}' is not available.",
+        }
+
+
+def get_current_time(city: str) -> dict:
+    """Returns the current time in a specified city.
+
+    Args:
+        city (str): The name of the city for which to retrieve the current time.
+
+    Returns:
+        dict: status and result or error msg.
+    """
+
+    if city.lower() == "new york":
+        tz_identifier = "America/New_York"
+    else:
+        return {
+            "status": "error",
+            "error_message": (f"Sorry, I don't have timezone information for {city}."),
+        }
+
+    tz = ZoneInfo(tz_identifier)
+    now = datetime.datetime.now(tz)
+    report = f'The current time in {city} is {now.strftime("%Y-%m-%d %H:%M:%S %Z%z")}'
+    return {"status": "success", "report": report}
+
+
+opik_tracer = OpikTracer()
+
+root_agent = Agent(
+    name="weather_time_agent",
+    model="gemini-2.0-flash",
+    description=("Agent to answer questions about the time and weather in a city."),
+    instruction=(
+        "You are a helpful agent who can answer user questions about the time and weather in a city."
+    ),
+    tools=[get_weather, get_current_time],
+    before_agent_callback=opik_tracer.before_agent_callback,
+    after_agent_callback=opik_tracer.after_agent_callback,
+    before_model_callback=opik_tracer.before_model_callback,
+    after_model_callback=opik_tracer.after_model_callback,
+    before_tool_callback=opik_tracer.before_tool_callback,
+    after_tool_callback=opik_tracer.after_tool_callback,
+)
diff --git a/sdks/python/tests/library_integration/adk/test_opik_tracer.py b/sdks/python/tests/library_integration/adk/test_opik_tracer.py
new file mode 100644
index 0000000000..69f3811570
--- /dev/null
+++ b/sdks/python/tests/library_integration/adk/test_opik_tracer.py
@@ -0,0 +1,64 @@
+import subprocess
+
+import pytest
+import requests
+
+from opik import synchronization
+
+ADK_SERVER_PORT = 21345
+
+
+def _check_server_running(base_url: str, user_id: str, session_id: str) -> bool:
+    try:
+        url = f"{base_url}/apps/sample_agent/users/{user_id}/session/{session_id}"
+        response = requests.get(url)
+        if response.status_code == 200:
+            return True
+    except requests.exceptions.ConnectionError:
+        return False
+    return False
+
+
+@pytest.fixture(scope="module")
+def api_server() -> str:
+    with subprocess.Popen(["adk", "api_server", str(ADK_SERVER_PORT)]) as proc:
+        base_url = f"http://localhost:{ADK_SERVER_PORT}"
+        user = "user_113"
+        session = "session_113"
+
+        # wait until the server is ready and session created
+        synchronization.wait_for_done(
+            check_function=lambda: _check_server_running(
+                base_url=base_url, user_id=user, session_id=session
+            ),
+            timeout=20,
+            sleep_time=5,
+        )
+
+        yield base_url, user, session
+
+        print(proc.stdout.read())
+
+
+def test_opik_tracer_with_sample_agent(api_server, fake_backend) -> None:
+    base_url, user_id, session_id = api_server
+
+    # send the request to the ADK API server
+    result = requests.post(
+        f"{base_url}/run",
+        json={
+            "app_name": "my_sample_agent",
+            "user_id": user_id,
+            "session_id": session_id,
+            "new_message": {
+                "role": "user",
+                "parts": [{"text": "Hey, whats the weather in New York today?"}],
+            },
+        },
+    )
+
+    print(result.json())
+    # assert result.status_code == 200
+
+    print(fake_backend.trace_trees)
+    assert len(fake_backend.trace_trees) == 1


Code review comments:
diff_hunk: @@ -0,0 +1,61 @@
+# Workflow to run ADK tests
+#
+# Please read inputs to provide correct values.
+#
+name: SDK Lib ADK (Google) Tests
+run-name: "SDK Lib ADK Tests ${{ github.ref_name }} by @${{ github.actor }}"
+env:
+  GCP_CREDENTIALS_JSON: ${{ secrets.GCP_CREDENTIALS_JSON }}
+  GOOGLE_CLOUD_LOCATION: us-east1
+  GOOGLE_CLOUD_PROJECT: opik-sdk-tests
+  GOOGLE_GENAI_USE_VERTEXAI: TRUE
+  OPIK_ENABLE_LITELLM_MODELS_MONITORING: False
+  OPIK_SENTRY_ENABLE: False
+on:
+  workflow_call:
+  pull_request:
+
+jobs:
+  tests:
+    name: ADK Python ${{matrix.python_version}}
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: sdks/python
+
+    strategy:
+      fail-fast: true
+      matrix:
+        python_version: [
+          "3.9",
+          "3.10",
+          "3.11",
+          "3.12"
+        ]
+
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+
+      - name: Setup Python ${{matrix.python_version}}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{matrix.python_version}}
+
+      - name: Install opik
+        run: pip install .
+
+      - name: Install test tools
+        run: |
+          cd ./tests
+          pip install --no-cache-dir --disable-pip-version-check -r test_requirements.txt
+
+      - name: Install lib
+        run: |
+          cd ./tests
+          pip install --no-cache-dir --disable-pip-version-check -r library_integration/adk/requirements.txt
+
+      - name: Run tests
+        run: |
+          cd ./tests/library_integration/adk/
+          python -m pytest  -vv .
path: .github/workflows/lib-adk-tests.yml
text: @user1:
## Workflow does not contain permissions

Actions job or workflow does not limit the permissions of the GITHUB_TOKEN. Consider setting an explicit permissions block, using the following as a minimal starting point: {{contents: read}}

[Show more details](https://github.com/comet-ml/opik/security/code-scanning/113)

diff_hunk: @@ -110,3 +111,8 @@
     uses: ./.github/workflows/lib-genai-tests.yml
     secrets: inherit
 
+  adk_tests:
+    needs: [ init_environment ]
+    if: contains(fromJSON('["adk", "all"]'), needs.init_environment.outputs.LIBS)
+    uses: ./.github/workflows/lib-adk-tests.yml
+    secrets: inherit
path: .github/workflows/lib-integration-tests-runner.yml
text: @user1:
## Workflow does not contain permissions

Actions job or workflow does not limit the permissions of the GITHUB_TOKEN. Consider setting an explicit permissions block, using the following as a minimal starting point: {{}}

[Show more details](https://github.com/comet-ml/opik/security/code-scanning/114)