[OPEN-7119] Add support for Google ADK

examples/tracing/google-adk/google_adk_tracing.ipynb

@@ -0,0 +1,293 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/openlayer-ai/openlayer-python/blob/main/examples/tracing/google-adk/google_adk_tracing.ipynb)\n",
+    "\n",
+    "# Google ADK Tracing with Openlayer\n",
+    "\n",
+    "This notebook demonstrates how to trace Google Agent Development Kit (ADK) agents with Openlayer.\n",
+    "\n",
+    "## Prerequisites\n",
+    "\n",
+    "Install the required packages:\n",
+    "```bash\n",
+    "pip install openlayer google-adk wrapt\n",
+    "```\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup\n",
+    "\n",
+    "First, configure your Openlayer credentials and Google Cloud credentials:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "# Openlayer configuration\n",
+    "os.environ[\"OPENLAYER_API_KEY\"] = \"your-api-key-here\"\n",
+    "os.environ[\"OPENLAYER_INFERENCE_PIPELINE_ID\"] = \"your-pipeline-id-here\"\n",
+    "\n",
+    "# Google AI API configuration (Option 1: Using Google AI Studio)\n",
+    "# Get your API key from: https://aistudio.google.com/apikey\n",
+    "os.environ[\"GOOGLE_API_KEY\"] = \"your-google-ai-api-key-here\"\n",
+    "\n",
+    "# Google Cloud Vertex AI configuration (Option 2: Using Google Cloud)\n",
+    "# Uncomment these if you're using Vertex AI instead of Google AI\n",
+    "# os.environ[\"GOOGLE_APPLICATION_CREDENTIALS\"] = \"path/to/your/service-account-key.json\"\n",
+    "# os.environ[\"GOOGLE_CLOUD_PROJECT\"] = \"your-project-id\"\n",
+    "# os.environ[\"GOOGLE_CLOUD_LOCATION\"] = \"us-central1\"\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Enable Google ADK Tracing\n",
+    "\n",
+    "Enable tracing before creating any agents. This patches Google ADK globally to send traces to Openlayer:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from openlayer.lib.integrations import trace_google_adk\n",
+    "\n",
+    "# Enable tracing (must be called before creating agents)\n",
+    "trace_google_adk()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Example 1: Basic Agent with LLM Calls\n",
+    "\n",
+    "Create a simple agent that responds to user queries:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "from google.genai import types\n",
+    "from google.adk.agents import LlmAgent\n",
+    "from google.adk.runners import Runner\n",
+    "from google.adk.sessions import InMemorySessionService\n",
+    "\n",
+    "# Setup constants\n",
+    "APP_NAME = \"openlayer_demo\"\n",
+    "USER_ID = \"user_123\"\n",
+    "SESSION_ID = \"session_123\"\n",
+    "\n",
+    "# Create session service (shared across examples)\n",
+    "session_service = InMemorySessionService()\n",
+    "\n",
+    "# Create a basic agent\n",
+    "agent = LlmAgent(\n",
+    "    model=\"gemini-2.0-flash-exp\",\n",
+    "    name=\"Assistant\",\n",
+    "    instruction=\"You are a helpful assistant. Provide concise and accurate responses.\"\n",
+    ")\n",
+    "\n",
+    "# Create runner\n",
+    "runner = Runner(\n",
+    "    agent=agent,\n",
+    "    app_name=APP_NAME,\n",
+    "    session_service=session_service\n",
+    ")\n",
+    "\n",
+    "# Define async function to run the agent\n",
+    "async def run_basic_agent():\n",
+    "    # Create session\n",
+    "    await session_service.create_session(\n",
+    "        app_name=APP_NAME,\n",
+    "        user_id=USER_ID,\n",
+    "        session_id=SESSION_ID\n",
+    "    )\n",
+    "    \n",
+    "    # Run the agent\n",
+    "    query = \"What is the capital of France?\"\n",
+    "    content = types.Content(role='user', parts=[types.Part(text=query)])\n",
+    "    \n",
+    "    # Process events and get response\n",
+    "    async for event in runner.run_async(user_id=USER_ID, session_id=SESSION_ID, new_message=content):\n",
+    "        if event.is_final_response() and event.content:\n",
+    "            final_answer = event.content.parts[0].text.strip()\n",
+    "\n",
+    "# Run the async function\n",
+    "await run_basic_agent()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Example 2: Agent with Tools/Functions\n",
+    "\n",
+    "Create an agent with custom tools that can be called during execution:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "\n",
+    "# Define custom tools as regular Python functions\n",
+    "def get_weather(city: str) -> str:\n",
+    "    \"\"\"Retrieves the current weather report for a specified city.\n",
+    "    \n",
+    "    Args:\n",
+    "        city: The name of the city for which to retrieve the weather report.\n",
+    "    \n",
+    "    Returns:\n",
+    "        str: Weather report or error message.\n",
+    "    \"\"\"\n",
+    "    if city.lower() == \"san francisco\":\n",
+    "        return \"The weather in San Francisco is sunny with a temperature of 72°F (22°C).\"\n",
+    "    else:\n",
+    "        return f\"Sorry, weather information for '{city}' is not available.\"\n",
+    "\n",
+    "def calculate(expression: str) -> str:\n",
+    "    \"\"\"Evaluates a mathematical expression.\n",
+    "    \n",
+    "    Args:\n",
+    "        expression: A mathematical expression to evaluate.\n",
+    "    \n",
+    "    Returns:\n",
+    "        str: Calculation result or error message.\n",
+    "    \"\"\"\n",
+    "    try:\n",
+    "        result = eval(expression)\n",
+    "        return f\"The result is {result}\"\n",
+    "    except Exception as e:\n",
+    "        return f\"Error: {str(e)}\"\n",
+    "\n",
+    "# Use different session IDs for tool agent\n",
+    "TOOL_USER_ID = \"user_456\"\n",
+    "TOOL_SESSION_ID = \"session_456\"\n",
+    "\n",
+    "# Create agent with tools (pass functions directly)\n",
+    "tool_agent = LlmAgent(\n",
+    "    model=\"gemini-2.0-flash-exp\",\n",
+    "    name=\"ToolAgent\",\n",
+    "    instruction=\"You are a helpful assistant with access to weather and calculation tools. Use them when appropriate.\",\n",
+    "    tools=[get_weather, calculate]\n",
+    ")\n",
+    "\n",
+    "# Create runner for tool agent (reuse the session_service)\n",
+    "tool_runner = Runner(\n",
+    "    agent=tool_agent,\n",
+    "    app_name=APP_NAME,\n",
+    "    session_service=session_service\n",
+    ")\n",
+    "\n",
+    "# Define async function to run the tool agent\n",
+    "async def run_tool_agent():\n",
+    "    # Create session\n",
+    "    await session_service.create_session(\n",
+    "        app_name=APP_NAME,\n",
+    "        user_id=TOOL_USER_ID,\n",
+    "        session_id=TOOL_SESSION_ID\n",
+    "    )\n",
+    "    \n",
+    "    # Run the agent with a query that requires tool use\n",
+    "    query = \"What's the weather in San Francisco? Also, what is 15 * 24?\"\n",
+    "    content = types.Content(role='user', parts=[types.Part(text=query)])\n",
+    "    \n",
+    "    # Process events and get response\n",
+    "    async for event in tool_runner.run_async(\n",
+    "        user_id=TOOL_USER_ID,\n",
+    "        session_id=TOOL_SESSION_ID,\n",
+    "        new_message=content\n",
+    "    ):\n",
+    "        if event.is_final_response() and event.content:\n",
+    "            final_answer = event.content.parts[0].text.strip()\n",
+    "\n",
+    "# Run the async function\n",
+    "await run_tool_agent()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## View Traces in Openlayer\n",
+    "\n",
+    "After running these examples, you can view the traces in your Openlayer dashboard:\n",
+    "\n",
+    "1. Go to https://app.openlayer.com\n",
+    "2. Navigate to your inference pipeline\n",
+    "3. View the traces tab to see:\n",
+    "   - Agent execution steps\n",
+    "   - LLM calls with token counts\n",
+    "   - Tool executions with inputs and outputs\n",
+    "   - Latency for each operation\n",
+    "   - Complete execution hierarchy\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Disable Tracing\n",
+    "\n",
+    "When you're done, you can disable tracing to restore ADK's default behavior:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from openlayer.lib.integrations import unpatch_google_adk\n",
+    "\n",
+    "# Disable tracing\n",
+    "unpatch_google_adk()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.18"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

pyproject.toml

@@ -19,7 +19,8 @@ dependencies = [
     "pyarrow>=18.0.0; python_version >= '3.9'",
     "pyyaml>=6.0",
     "requests_toolbelt>=1.0.0",
-    "tqdm"
+    "tqdm",
+    "wrapt>=1.14.0"
 ]
 requires-python = ">= 3.8"
 classifiers = [

src/openlayer/lib/__init__.py

@@ -14,6 +14,8 @@
     "trace_oci_genai",
     "trace_oci",  # Alias for backward compatibility
     "trace_litellm",
+    "trace_google_adk",
+    "unpatch_google_adk",
     "update_current_trace",
     "update_current_step",
     # Offline buffer management functions
@@ -188,3 +190,48 @@ def trace_litellm():
     from .integrations import litellm_tracer
 
     return litellm_tracer.trace_litellm()
+
+
+# ------------------------------ Google ADK ---------------------------------- #
+def trace_google_adk():
+    """Enable tracing for Google Agent Development Kit (ADK).
+
+    This function patches Google ADK to automatically trace agent execution,
+    LLM calls, and tool calls made through the ADK framework.
+
+    Requirements:
+        Google ADK and wrapt must be installed:
+        pip install google-adk wrapt
+
+    Example:
+        >>> import os
+        >>> os.environ["OPENLAYER_API_KEY"] = "your-api-key"
+        >>> os.environ["OPENLAYER_INFERENCE_PIPELINE_ID"] = "your-pipeline-id"
+        >>> from openlayer.lib import trace_google_adk
+        >>> # Enable tracing (must be called before creating agents)
+        >>> trace_google_adk()
+        >>> # Now create and run your ADK agents
+        >>> from google.adk.agents import Agent
+        >>> agent = Agent(name="Assistant", model="gemini-2.0-flash-exp")
+        >>> result = await agent.run_async(...)
+    """
+    # pylint: disable=import-outside-toplevel
+    from .integrations import google_adk_tracer
+
+    return google_adk_tracer.trace_google_adk()
+
+
+def unpatch_google_adk():
+    """Remove Google ADK tracing patches.
+
+    This function restores Google ADK's original behavior by removing all
+    Openlayer instrumentation.
+
+    Example:
+        >>> from openlayer.lib import unpatch_google_adk
+        >>> unpatch_google_adk()
+    """
+    # pylint: disable=import-outside-toplevel
+    from .integrations import google_adk_tracer
+
+    return google_adk_tracer.unpatch_google_adk()

src/openlayer/lib/integrations/__init__.py

@@ -24,3 +24,10 @@
     __all__.extend(["trace_oci_genai"])
 except ImportError:
     pass
+
+try:
+    from .google_adk_tracer import trace_google_adk, unpatch_google_adk
+
+    __all__.extend(["trace_google_adk", "unpatch_google_adk"])
+except ImportError:
+    pass