gemini/reasoning-engine/tracing_agents_in_reasoning_engine.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Debugging and Optimizing Agents: A Guide to Tracing in Reasoning Engine\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/reasoning-engine/tracing_agents_in_reasoning_engine.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Freasoning-engine%2Ftracing_agents_in_reasoning_engine.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>    \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/reasoning-engine/tracing_agents_in_reasoning_engine.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/reasoning-engine/tracing_agents_in_reasoning_engine.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/reasoning-engine/tracing_agents_in_reasoning_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/reasoning-engine/tracing_agents_in_reasoning_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/reasoning-engine/tracing_agents_in_reasoning_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/53/X_logo_2023_original.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/reasoning-engine/tracing_agents_in_reasoning_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/reasoning-engine/tracing_agents_in_reasoning_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>            "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84f0f73a0f76"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "| Author(s) | [Kristopher Overholt](https://github.com/koverholt) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "[Reasoning Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/reasoning-engine/overview) (LangChain on Vertex AI) helps you build and deploy agent-based AI applications that use LLMs and custom tools. Understanding your agent's decision-making process is essential for debugging and optimization, and [Cloud Trace](https://cloud.google.com/trace) is a great tool for exploring this tracing data to get insights.\n",
        "\n",
        "<img src=\"https://storage.googleapis.com/github-repo/generative-ai/gemini/reasoning-engine/images/cloud-trace-agent.png\">\n",
        "\n",
        "This notebook demonstrates how to:\n",
        "\n",
        "- **Learn Key Concepts**: Learn about the fundamental building blocks of tracing.\n",
        "- **Deploy Your Agent**: Make your tracing-enabled agent available in a production-like environment on Reasoning Engine.\n",
        "- **Enable Tracing**: Enable tracing in a simple agent\n",
        "- **Examine Traces**: Use the Cloud Console and Cloud Trace SDK to access and analyze a specific trace.\n",
        "\n",
        "By the end of this notebook, you'll be able to leverage tracing to build more robust and efficient AI agents on Vertex AI."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d38b3d918526"
      },
      "source": [
        "## Concepts\n",
        "\n",
        "Here are some of the key concepts and terminology related to tracing, which will be helpful to understand as we explore traces generated by an agent in Reasoning Engine:\n",
        "\n",
        "Below is an example of a trace in JSON format, showing a single span. This span represents a call to a large language model (LLM). Notice how the trace data captures important details:\n",
        "\n",
        "### Example trace\n",
        "\n",
        "```json\n",
        "{\n",
        "   \"name\": \"llm\",\n",
        "   \"context\": {\n",
        "       \"trace_id\": \"ed7b336d-e71a-46f0-a334-5f2e87cb6cfc\",\n",
        "       \"span_id\": \"ad67332a-38bd-428e-9f62-538ba2fa90d4\"\n",
        "   },\n",
        "   \"span_kind\": \"LLM\",\n",
        "   \"parent_id\": \"f89ebb7c-10f6-4bf8-8a74-57324d2556ef\",\n",
        "   \"start_time\": \"2023-09-07T12:54:47.597121-06:00\",\n",
        "   \"end_time\": \"2023-09-07T12:54:49.321811-06:00\",\n",
        "   \"status_code\": \"OK\",\n",
        "   \"status_message\": \"\",\n",
        "   \"attributes\": {\n",
        "       \"llm.input_messages\": [\n",
        "           {\n",
        "               \"message.role\": \"system\",\n",
        "               \"message.content\": \"You are an expert Q&A system that is trusted around the world.\\nAlways answer the query using the provided context information, and not prior knowledge.\\nSome rules to follow:\\n1. Never directly reference the given context in your answer.\\n2. Avoid statements like 'Based on the context, ...' or 'The context information ...' or anything along those lines.\"\n",
        "           },\n",
        "           {\n",
        "               \"message.role\": \"user\",\n",
        "               \"message.content\": \"Hello?\"\n",
        "           }\n",
        "       ],\n",
        "       \"output.value\": \"assistant: Yes I am here\",\n",
        "       \"output.mime_type\": \"text/plain\"\n",
        "   },\n",
        "   \"events\": [],\n",
        "}\n",
        "```\n",
        "\n",
        "### Trace\n",
        "\n",
        "You can think of a [trace](https://opentelemetry.io/docs/concepts/signals/traces/) like a timeline of requests as they travel through your application. A trace is composed of individual spans, with the first span representing the overall request. Each span provides details about a specific operation within the request.\n",
        "\n",
        "### Span\n",
        "\n",
        "A [span](https://opentelemetry.io/docs/concepts/signals/traces/#spans) represents a single unit of work, like a function call or an interaction with an LLM. It captures information such as the operation's name, start and end times, and any relevant attributes (metadata). Spans can be nested, showing parent-child relationships between operations.\n",
        "\n",
        "### Span Attribute\n",
        "\n",
        "[Span attributes](https://opentelemetry.io/docs/concepts/signals/traces/#attributes) are key-value pairs that provide additional context about a span. For instance, an LLM span might have attributes like the model name, prompt text, and token count.\n",
        "\n",
        "### Span Kind\n",
        "\n",
        "[Span kind](https://opentelemetry.io/docs/concepts/signals/traces/#span-kind) categorizes the type of operation a span represents. Common kinds include:\n",
        "\n",
        "- `CHAIN`: Links between LLM application steps or the start of a request.\n",
        "- `LLM`: A call to a large language model.\n",
        "- `TOOL`: An interaction with an external tool (API, database, etc.).\n",
        "- `AGENT`: A reasoning block that combines LLM and tool interactions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK and other required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --user --quiet \\\n",
        "    \"google-cloud-aiplatform[langchain,reasoningengine]\" \\\n",
        "    cloudpickle==3.0.0 \\\n",
        "    pydantic==2.7.4 \\\n",
        "    google-cloud-trace"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n",
        "\n",
        "The restart might take a minute or longer. After it's restarted, continue to the next step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "STAGING_BUCKET = \"gs://[your-staging-bucket]\"  # @param {type:\"string\"}\n",
        "\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=STAGING_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user(project_id=PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdvJRUWRNGHE"
      },
      "source": [
        "## Build and deploy an agent\n",
        "\n",
        "Let's dive into building a simple agent that utilizes tracing. This agent will use a few custom tools to demonstrate how tracing can provide insights into its workflow."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4c5e77496259"
      },
      "source": [
        "### Import libraries\n",
        "\n",
        "Before you start building your agent, you'll import the necessary libraries. These include the Vertex AI SDK, pandas for data analysis, and the Cloud Trace SDK for working with trace data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "5a2198c40f52"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime, timedelta\n",
        "\n",
        "from google.cloud import trace_v1 as trace\n",
        "import pandas as pd\n",
        "from vertexai.preview import reasoning_engines\n",
        "from vertexai.reasoning_engines._reasoning_engines import _utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6f73957911c3"
      },
      "source": [
        "### Define tools\n",
        "\n",
        "You'll define a few Python functions to act as tools for your agent. These tools will simulate actions or API calls that a real-world agent might perform. For this example, you'll create tools to classify a customer support ticket, query a knowledge base, and escalate a ticket to a human agent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "0e7c7732f26f"
      },
      "outputs": [],
      "source": [
        "def classify_ticket(ticket_text: str) -> str:\n",
        "    \"\"\"Classifies a support ticket into a category.\"\"\"\n",
        "    # Simulate a call to a classification model\n",
        "    categories = {\n",
        "        \"general\": \"Questions and information\",\n",
        "        \"billing\": \"Payment and invoices\",\n",
        "        \"technical\": \"API and SDK developer documentation\",\n",
        "    }\n",
        "    if \"payment\" in ticket_text:\n",
        "        category = \"billing\"\n",
        "        description = categories[category]\n",
        "    elif \"settings\" in ticket_text:\n",
        "        category = \"technical\"\n",
        "        description = categories[category]\n",
        "    else:\n",
        "        category = \"general\"\n",
        "        description = categories[category]\n",
        "\n",
        "    return f\"This ticket is in the {category} category for questions related to {description}\"\n",
        "\n",
        "\n",
        "def search_knowledge_base(category: str) -> list[dict]:\n",
        "    \"\"\"Searches a knowledge base for relevant articles and documentation links.\"\"\"\n",
        "    # Simulate a knowledge base search\n",
        "    articles = {\n",
        "        \"general\": [\n",
        "            {\n",
        "                \"title\": \"Contacting support\",\n",
        "                \"url\": \"https://example.com/contact\",\n",
        "            }\n",
        "        ],\n",
        "        \"billing\": [\n",
        "            {\n",
        "                \"title\": \"How to update your payment information\",\n",
        "                \"url\": \"https://example.com/billing/update\",\n",
        "            },\n",
        "        ],\n",
        "        \"technical\": [\n",
        "            {\n",
        "                \"title\": \"Troubleshooting common login issues\",\n",
        "                \"url\": \"https://example.com/technical/help\",\n",
        "            },\n",
        "        ],\n",
        "    }\n",
        "    return articles.get(category, [])\n",
        "\n",
        "\n",
        "def escalate_to_human(ticket_text: str) -> str:\n",
        "    \"\"\"Initiates escalation to a human agent for outage reports.\"\"\"\n",
        "    return \"Your ticket has been escalated to a human agent. Please expect a response within 1-2 hours.\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bd55dbf42903"
      },
      "source": [
        "### Define agent and enable tracing\n",
        "\n",
        "Now, let's define your agent using the LangChain template in Reasoning Engine and the Vertex AI SDK. Enable tracing by setting the `enable_tracing` parameter to `True`, which allows you to capture detailed information about the agent's execution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "78fb7ab21e3a"
      },
      "outputs": [],
      "source": [
        "agent = reasoning_engines.LangchainAgent(\n",
        "    model=\"gemini-1.5-pro\",\n",
        "    model_kwargs={\"temperature\": 0},\n",
        "    tools=[classify_ticket, search_knowledge_base, escalate_to_human],\n",
        "    enable_tracing=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6d714e7d8432"
      },
      "source": [
        "### Test your agent locally (with traces!)\n",
        "\n",
        "Let's test your agent locally by sending it a query. Since you've enabled tracing, you'll be able to see how the agent processes this request and interacts with its tools."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "d2e9ac4ab9e9"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'input': '\\n    Classify the following ticket into a category and give me a relevant documentation link.\\n    \\n    Support ticket text:\\n    I need to update my billing information since my payment method has expired.\\n    ',\n",
              " 'output': \"This ticket is about billing. Here's a link to our documentation about updating your payment information: https://example.com/billing/update. \\n\"}"
            ]
          },
          "execution_count": 8,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "agent.query(\n",
        "    input=\"\"\"\n",
        "    Classify the following ticket into a category and give me a relevant documentation link.\n",
        "    \n",
        "    Support ticket text:\n",
        "    I need to update my billing information since my payment method has expired.\n",
        "    \"\"\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "82b1fbdafe7c"
      },
      "source": [
        "### Get your first trace\n",
        "\n",
        "Before diving deeper into trace analysis, let's use the Cloud Trace SDK to retrieve a specific trace generated by your local agent. This will give you a concrete example to work with."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "d4894bb6a2d1"
      },
      "outputs": [],
      "source": [
        "client = trace.TraceServiceClient()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "ea72cdb247ec"
      },
      "outputs": [],
      "source": [
        "result = [\n",
        "    r\n",
        "    for r in client.list_traces(\n",
        "        request=trace.types.ListTracesRequest(\n",
        "            project_id=PROJECT_ID,\n",
        "            # Return all traces containing `labels {key: \"openinference.span.kind\" value: \"AGENT\"}`\n",
        "            filter=\"openinference.span.kind:AGENT\",\n",
        "        )\n",
        "    )\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "c56a32df5305"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "span_id: 1514793182933739503\n",
              "name: \"RunnableParallel<input,agent_scratchpad>\"\n",
              "start_time {\n",
              "  seconds: 1719635620\n",
              "  nanos: 435836160\n",
              "}\n",
              "end_time {\n",
              "  seconds: 1719635620\n",
              "  nanos: 816612096\n",
              "}\n",
              "parent_span_id: 10997585570109326486\n",
              "labels {\n",
              "  key: \"output.value\"\n",
              "  value: \"{\\\"input\\\": \\\"\\\\n    Classify the following ticket into a category and suggest a documentation link.\\\\n    \\\\n    Support ticket text:\\\\n    I need to update my billing information since my payment method has expired.\\\\n    \\\", \\\"agent_scratchpad\\\": []}\"\n",
              "}\n",
              "labels {\n",
              "  key: \"output.mime_type\"\n",
              "  value: \"application/json\"\n",
              "}\n",
              "labels {\n",
              "  key: \"openinference.span.kind\"\n",
              "  value: \"AGENT\"\n",
              "}\n",
              "labels {\n",
              "  key: \"input.value\"\n",
              "  value: \"{\\\"input\\\": \\\"\\\\n    Classify the following ticket into a category and suggest a documentation link.\\\\n    \\\\n    Support ticket text:\\\\n    I need to update my billing information since my payment method has expired.\\\\n    \\\", \\\"intermediate_steps\\\": []}\"\n",
              "}\n",
              "labels {\n",
              "  key: \"input.mime_type\"\n",
              "  value: \"application/json\"\n",
              "}\n",
              "labels {\n",
              "  key: \"g.co/agent\"\n",
              "  value: \"opentelemetry-python 1.25.0; google-cloud-trace-exporter 1.6.0\"\n",
              "}"
            ]
          },
          "execution_count": 11,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "trace_data = client.get_trace(project_id=PROJECT_ID, trace_id=result[0].trace_id).spans[\n",
        "    0\n",
        "]\n",
        "trace_data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "19aa87abae0a"
      },
      "source": [
        "After you deploy your agent and make remote queries in the following sections, you'll dive into the details for working with trace data in the Cloud Console or using the Python SDK for Cloud Trace."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NdGcHqUv8THp"
      },
      "source": [
        "### Deploy your agent\n",
        "\n",
        "Now that you've seen how tracing works locally, let's deploy your agent to Reasoning Engine. This will allow you to send it queries in a production-like environment and observe its behavior through traces."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "NrTI0_1j8E7w"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Using bucket your-bucket-name\n",
            "Writing to gs://your-bucket-name/reasoning_engine/reasoning_engine.pkl\n",
            "Writing to gs://your-bucket-name/reasoning_engine/requirements.txt\n",
            "Creating in-memory tarfile of extra_packages\n",
            "Writing to gs://your-bucket-name/reasoning_engine/dependencies.tar.gz\n",
            "Creating ReasoningEngine\n",
            "Create ReasoningEngine backing LRO: projects/your-project-number/locations/us-central1/reasoningEngines/3146221736555446272/operations/2750583127503011840\n",
            "ReasoningEngine created. Resource name: projects/your-project-number/locations/us-central1/reasoningEngines/3146221736555446272\n",
            "To use this ReasoningEngine in another session:\n",
            "reasoning_engine = vertexai.preview.reasoning_engines.ReasoningEngine('projects/your-project-number/locations/us-central1/reasoningEngines/3146221736555446272')\n"
          ]
        }
      ],
      "source": [
        "remote_agent = reasoning_engines.ReasoningEngine.create(\n",
        "    agent,\n",
        "    requirements=[\n",
        "        \"google-cloud-aiplatform[langchain,reasoningengine]\",\n",
        "        \"cloudpickle==3.0.0\",\n",
        "        \"pydantic==2.7.4\",\n",
        "    ],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "019c774a18fd"
      },
      "source": [
        "### Query your deployed agent\n",
        "\n",
        "With your agent deployed, you can interact with it remotely. Let's send a query and generate some trace data to explore."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "zZjdFQ_Z_J43"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'input': '\\n    Classify the following ticket into a category and route the customer accordingly:\\n    \\n    Support ticket text:\\n    I am unable to make any API calls and I need to report an outage in the system\\n    ',\n",
              " 'output': 'The ticket has been escalated to a human agent. Please expect a response within 1-2 hours. \\n'}"
            ]
          },
          "execution_count": 13,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "remote_agent.query(\n",
        "    input=\"\"\"\n",
        "    Classify the following ticket into a category and route the customer accordingly:\n",
        "    \n",
        "    Support ticket text:\n",
        "    I am unable to make any API calls and I need to report an outage in the system\n",
        "    \"\"\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RYwmTcdgIlza"
      },
      "source": [
        "## Exploring traces in the Cloud Console\n",
        "\n",
        "The Cloud Trace console provides a powerful and intuitive visual interface for exploring trace data, including visualizing, filtering, and analyzing your traces.\n",
        "\n",
        "Accessing the Trace Console:\n",
        "\n",
        "- **Project-Level View**: To see all traces for your Google Cloud project (replace `your-project-id` with your actual project ID), go to: https://console.cloud.google.com/traces/list?project=your-project-id\n",
        "\n",
        "- **Specific Trace**: If you know the unique Trace ID for a specific trace you want to examine, you can view it directly (replace your-trace-id with the actual Trace ID): https://console.cloud.google.com/traces/list?project=your-project-id&tid=your-trace-id\n",
        "\n",
        "Features to Explore in the Console:\n",
        "\n",
        "- **Trace List**: View a list of traces, sorted by start time, along with summary information (duration, number of spans).\n",
        "- **Waterfall View**: Visualize the spans within a trace as a timeline, showing the duration of each operation and their relationships.\n",
        "- **Span Details**: Click on a span to view its attributes, including the input and output data, and any custom metadata you've added.\n",
        "- **Filtering and Search**: The console provides powerful options for filtering traces by time range, service, span name, and other criteria. You can also search for specific traces using keywords or attributes.\n",
        "\n",
        "For a detailed guide to working with traces in the console, refer to the [Cloud Trace documentation on finding traces](https://cloud.google.com/trace/docs/finding-traces). Experiment with the Cloud Trace console to gain a deeper understanding of your agent's behavior and how it's executing within Reasoning Engine."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "761e2e6121f6"
      },
      "source": [
        "## Working with traces using `pandas`\n",
        "\n",
        "For more programmatic analysis, you can use the pandas library to work with trace data. You'll fetch traces, convert them to DataFrames, and then use pandas' functionality to explore the trace data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "e68cf66f2b2c"
      },
      "outputs": [],
      "source": [
        "result = [\n",
        "    r\n",
        "    for r in client.list_traces(\n",
        "        request=trace.types.ListTracesRequest(\n",
        "            project_id=PROJECT_ID,\n",
        "            # Return all traces containing `labels {key: \"openinference.span.kind\" value: \"AGENT\"}`\n",
        "            filter=\"openinference.span.kind:AGENT\",\n",
        "        )\n",
        "    )\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "0ee427fa9cb3"
      },
      "outputs": [],
      "source": [
        "trace_data = client.get_trace(project_id=PROJECT_ID, trace_id=result[0].trace_id)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "73a61c86fc7c"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>spanId</th>\n",
              "      <th>name</th>\n",
              "      <th>startTime</th>\n",
              "      <th>endTime</th>\n",
              "      <th>labels</th>\n",
              "      <th>parentSpanId</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>12108864055932459456</td>\n",
              "      <td>AgentExecutor</td>\n",
              "      <td>2024-06-29T04:33:40.426806016Z</td>\n",
              "      <td>2024-06-29T04:33:49.886738176Z</td>\n",
              "      <td>{'g.co/agent': 'opentelemetry-python 1.25.0; g...</td>\n",
              "      <td>NaN</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>1514793182933739503</td>\n",
              "      <td>RunnableParallel&lt;input,agent_scratchpad&gt;</td>\n",
              "      <td>2024-06-29T04:33:40.435836160Z</td>\n",
              "      <td>2024-06-29T04:33:40.816612096Z</td>\n",
              "      <td>{'g.co/agent': 'opentelemetry-python 1.25.0; g...</td>\n",
              "      <td>10997585570109326486</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>2419854474582761855</td>\n",
              "      <td>ChatPromptTemplate</td>\n",
              "      <td>2024-06-29T04:33:40.988899840Z</td>\n",
              "      <td>2024-06-29T04:33:40.990212864Z</td>\n",
              "      <td>{'g.co/agent': 'opentelemetry-python 1.25.0; g...</td>\n",
              "      <td>10997585570109326486</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>883609395272937079</td>\n",
              "      <td>ChatVertexAI</td>\n",
              "      <td>2024-06-29T04:33:41.174941184Z</td>\n",
              "      <td>2024-06-29T04:33:43.206272Z</td>\n",
              "      <td>{'llm.invocation_parameters': '{\"model_name\": ...</td>\n",
              "      <td>10997585570109326486</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>10997585570109326486</td>\n",
              "      <td>RunnableSequence</td>\n",
              "      <td>2024-06-29T04:33:40.431934208Z</td>\n",
              "      <td>2024-06-29T04:33:43.667400960Z</td>\n",
              "      <td>{'g.co/agent': 'opentelemetry-python 1.25.0; g...</td>\n",
              "      <td>12108864055932459456</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "                 spanId                                      name  \\\n",
              "0  12108864055932459456                             AgentExecutor   \n",
              "1   1514793182933739503  RunnableParallel<input,agent_scratchpad>   \n",
              "2   2419854474582761855                        ChatPromptTemplate   \n",
              "3    883609395272937079                              ChatVertexAI   \n",
              "4  10997585570109326486                          RunnableSequence   \n",
              "\n",
              "                        startTime                         endTime  \\\n",
              "0  2024-06-29T04:33:40.426806016Z  2024-06-29T04:33:49.886738176Z   \n",
              "1  2024-06-29T04:33:40.435836160Z  2024-06-29T04:33:40.816612096Z   \n",
              "2  2024-06-29T04:33:40.988899840Z  2024-06-29T04:33:40.990212864Z   \n",
              "3  2024-06-29T04:33:41.174941184Z     2024-06-29T04:33:43.206272Z   \n",
              "4  2024-06-29T04:33:40.431934208Z  2024-06-29T04:33:43.667400960Z   \n",
              "\n",
              "                                              labels          parentSpanId  \n",
              "0  {'g.co/agent': 'opentelemetry-python 1.25.0; g...                   NaN  \n",
              "1  {'g.co/agent': 'opentelemetry-python 1.25.0; g...  10997585570109326486  \n",
              "2  {'g.co/agent': 'opentelemetry-python 1.25.0; g...  10997585570109326486  \n",
              "3  {'llm.invocation_parameters': '{\"model_name\": ...  10997585570109326486  \n",
              "4  {'g.co/agent': 'opentelemetry-python 1.25.0; g...  12108864055932459456  "
            ]
          },
          "execution_count": 16,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "spans = pd.DataFrame.from_records([_utils.to_dict(span) for span in trace_data.spans])\n",
        "spans.head()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "a5fdff97b29f"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>spanId</th>\n",
              "      <th>name</th>\n",
              "      <th>startTime</th>\n",
              "      <th>endTime</th>\n",
              "      <th>labels</th>\n",
              "      <th>parentSpanId</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>883609395272937079</td>\n",
              "      <td>ChatVertexAI</td>\n",
              "      <td>2024-06-29T04:33:41.174941184Z</td>\n",
              "      <td>2024-06-29T04:33:43.206272Z</td>\n",
              "      <td>{'llm.invocation_parameters': '{\"model_name\": ...</td>\n",
              "      <td>10997585570109326486</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>14</th>\n",
              "      <td>1267786026295004320</td>\n",
              "      <td>ChatVertexAI</td>\n",
              "      <td>2024-06-29T04:33:44.592414976Z</td>\n",
              "      <td>2024-06-29T04:33:46.278656Z</td>\n",
              "      <td>{'g.co/agent': 'opentelemetry-python 1.25.0; g...</td>\n",
              "      <td>8170159531257482779</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>22</th>\n",
              "      <td>15517842396867663478</td>\n",
              "      <td>ChatVertexAI</td>\n",
              "      <td>2024-06-29T04:33:47.494297856Z</td>\n",
              "      <td>2024-06-29T04:33:49.470009088Z</td>\n",
              "      <td>{'g.co/agent': 'opentelemetry-python 1.25.0; g...</td>\n",
              "      <td>4922396732366218475</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "                  spanId          name                       startTime  \\\n",
              "3     883609395272937079  ChatVertexAI  2024-06-29T04:33:41.174941184Z   \n",
              "14   1267786026295004320  ChatVertexAI  2024-06-29T04:33:44.592414976Z   \n",
              "22  15517842396867663478  ChatVertexAI  2024-06-29T04:33:47.494297856Z   \n",
              "\n",
              "                           endTime  \\\n",
              "3      2024-06-29T04:33:43.206272Z   \n",
              "14     2024-06-29T04:33:46.278656Z   \n",
              "22  2024-06-29T04:33:49.470009088Z   \n",
              "\n",
              "                                               labels          parentSpanId  \n",
              "3   {'llm.invocation_parameters': '{\"model_name\": ...  10997585570109326486  \n",
              "14  {'g.co/agent': 'opentelemetry-python 1.25.0; g...   8170159531257482779  \n",
              "22  {'g.co/agent': 'opentelemetry-python 1.25.0; g...   4922396732366218475  "
            ]
          },
          "execution_count": 17,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "spans[spans[\"name\"] == \"ChatVertexAI\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "691fe8935660"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>llm.invocation_parameters</th>\n",
              "      <th>g.co/agent</th>\n",
              "      <th>output.mime_type</th>\n",
              "      <th>llm.input_messages.0.message.content</th>\n",
              "      <th>llm.output_messages.0.message.function_call_arguments_json</th>\n",
              "      <th>output.value</th>\n",
              "      <th>input.mime_type</th>\n",
              "      <th>llm.input_messages.0.message.role</th>\n",
              "      <th>input.value</th>\n",
              "      <th>llm.output_messages.0.message.function_call_name</th>\n",
              "      <th>...</th>\n",
              "      <th>llm.input_messages.1.message.role</th>\n",
              "      <th>llm.input_messages.2.message.role</th>\n",
              "      <th>llm.input_messages.1.message.function_call_arguments_json</th>\n",
              "      <th>llm.input_messages.1.message.function_call_name</th>\n",
              "      <th>llm.input_messages.4.message.content</th>\n",
              "      <th>llm.input_messages.3.message.role</th>\n",
              "      <th>llm.input_messages.4.message.role</th>\n",
              "      <th>llm.output_messages.0.message.content</th>\n",
              "      <th>llm.input_messages.3.message.function_call_name</th>\n",
              "      <th>llm.input_messages.3.message.function_call_arguments_json</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>{\"model_name\": \"gemini-1.5-pro-001\", \"temperature...</td>\n",
              "      <td>opentelemetry-python 1.25.0; google-cloud-trac...</td>\n",
              "      <td>application/json</td>\n",
              "      <td>\\n    Classify the following ticket into a cat...</td>\n",
              "      <td>{\"ticket_text\": \"I need to update my billing i...</td>\n",
              "      <td>{\"generations\": [[{\"text\": \"\", \"generation_inf...</td>\n",
              "      <td>application/json</td>\n",
              "      <td>user</td>\n",
              "      <td>{\"messages\": [[{\"lc\": 1, \"type\": \"constructor\"...</td>\n",
              "      <td>classify_ticket</td>\n",
              "      <td>...</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>14</th>\n",
              "      <td>{\"model_name\": \"gemini-1.5-pro-001\", \"temperature...</td>\n",
              "      <td>opentelemetry-python 1.25.0; google-cloud-trac...</td>\n",
              "      <td>application/json</td>\n",
              "      <td>\\n    Classify the following ticket into a cat...</td>\n",
              "      <td>{\"category\": \"billing\"}</td>\n",
              "      <td>{\"generations\": [[{\"text\": \"\", \"generation_inf...</td>\n",
              "      <td>application/json</td>\n",
              "      <td>user</td>\n",
              "      <td>{\"messages\": [[{\"lc\": 1, \"type\": \"constructor\"...</td>\n",
              "      <td>query_knowledge_base</td>\n",
              "      <td>...</td>\n",
              "      <td>assistant</td>\n",
              "      <td>tool</td>\n",
              "      <td>{\"ticket_text\": \"I need to update my billing i...</td>\n",
              "      <td>classify_ticket</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "      <td>NaN</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>22</th>\n",
              "      <td>{\"model_name\": \"gemini-1.5-pro-001\", \"temperature...</td>\n",
              "      <td>opentelemetry-python 1.25.0; google-cloud-trac...</td>\n",
              "      <td>application/json</td>\n",
              "      <td>\\n    Classify the following ticket into a cat...</td>\n",
              "      <td>NaN</td>\n",
              "      <td>{\"generations\": [[{\"text\": \"The ticket is related...</td>\n",
              "      <td>application/json</td>\n",
              "      <td>user</td>\n",
              "      <td>{\"messages\": [[{\"lc\": 1, \"type\": \"constructor\"...</td>\n",
              "      <td>NaN</td>\n",
              "      <td>...</td>\n",
              "      <td>assistant</td>\n",
              "      <td>tool</td>\n",
              "      <td>{\"ticket_text\": \"I need to update my billing i...</td>\n",
              "      <td>classify_ticket</td>\n",
              "      <td>[{\"title\": \"How to update your payment information...</td>\n",
              "      <td>assistant</td>\n",
              "      <td>tool</td>\n",
              "      <td>The ticket is related to **billing**. Here's a...</td>\n",
              "      <td>query_knowledge_base</td>\n",
              "      <td>{\"category\": \"billing\"}</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "<p>3 rows × 24 columns</p>\n",
              "</div>"
            ],
            "text/plain": [
              "                            llm.invocation_parameters  \\\n",
              "3   {\"model_name\": \"gemini-1.5-pro-001\", \"temperature...   \n",
              "14  {\"model_name\": \"gemini-1.5-pro-001\", \"temperature...   \n",
              "22  {\"model_name\": \"gemini-1.5-pro-001\", \"temperature...   \n",
              "\n",
              "                                           g.co/agent  output.mime_type  \\\n",
              "3   opentelemetry-python 1.25.0; google-cloud-trac...  application/json   \n",
              "14  opentelemetry-python 1.25.0; google-cloud-trac...  application/json   \n",
              "22  opentelemetry-python 1.25.0; google-cloud-trac...  application/json   \n",
              "\n",
              "                 llm.input_messages.0.message.content  \\\n",
              "3   \\n    Classify the following ticket into a cat...   \n",
              "14  \\n    Classify the following ticket into a cat...   \n",
              "22  \\n    Classify the following ticket into a cat...   \n",
              "\n",
              "   llm.output_messages.0.message.function_call_arguments_json  \\\n",
              "3   {\"ticket_text\": \"I need to update my billing i...           \n",
              "14                            {\"category\": \"billing\"}           \n",
              "22                                                NaN           \n",
              "\n",
              "                                         output.value   input.mime_type  \\\n",
              "3   {\"generations\": [[{\"text\": \"\", \"generation_inf...  application/json   \n",
              "14  {\"generations\": [[{\"text\": \"\", \"generation_inf...  application/json   \n",
              "22  {\"generations\": [[{\"text\": \"The ticket is related...  application/json   \n",
              "\n",
              "   llm.input_messages.0.message.role  \\\n",
              "3                               user   \n",
              "14                              user   \n",
              "22                              user   \n",
              "\n",
              "                                          input.value  \\\n",
              "3   {\"messages\": [[{\"lc\": 1, \"type\": \"constructor\"...   \n",
              "14  {\"messages\": [[{\"lc\": 1, \"type\": \"constructor\"...   \n",
              "22  {\"messages\": [[{\"lc\": 1, \"type\": \"constructor\"...   \n",
              "\n",
              "   llm.output_messages.0.message.function_call_name  ...  \\\n",
              "3                                   classify_ticket  ...   \n",
              "14                             query_knowledge_base  ...   \n",
              "22                                              NaN  ...   \n",
              "\n",
              "   llm.input_messages.1.message.role llm.input_messages.2.message.role  \\\n",
              "3                                NaN                               NaN   \n",
              "14                         assistant                              tool   \n",
              "22                         assistant                              tool   \n",
              "\n",
              "   llm.input_messages.1.message.function_call_arguments_json  \\\n",
              "3                                                 NaN          \n",
              "14  {\"ticket_text\": \"I need to update my billing i...          \n",
              "22  {\"ticket_text\": \"I need to update my billing i...          \n",
              "\n",
              "   llm.input_messages.1.message.function_call_name  \\\n",
              "3                                              NaN   \n",
              "14                                 classify_ticket   \n",
              "22                                 classify_ticket   \n",
              "\n",
              "                 llm.input_messages.4.message.content  \\\n",
              "3                                                 NaN   \n",
              "14                                                NaN   \n",
              "22  [{\"title\": \"How to update your payment information...   \n",
              "\n",
              "   llm.input_messages.3.message.role llm.input_messages.4.message.role  \\\n",
              "3                                NaN                               NaN   \n",
              "14                               NaN                               NaN   \n",
              "22                         assistant                              tool   \n",
              "\n",
              "                llm.output_messages.0.message.content  \\\n",
              "3                                                 NaN   \n",
              "14                                                NaN   \n",
              "22  The ticket is related to **billing**. Here's a...   \n",
              "\n",
              "   llm.input_messages.3.message.function_call_name  \\\n",
              "3                                              NaN   \n",
              "14                                             NaN   \n",
              "22                            query_knowledge_base   \n",
              "\n",
              "   llm.input_messages.3.message.function_call_arguments_json  \n",
              "3                                                 NaN         \n",
              "14                                                NaN         \n",
              "22                            {\"category\": \"billing\"}         \n",
              "\n",
              "[3 rows x 24 columns]"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "spans[spans[\"name\"] == \"ChatVertexAI\"].labels.apply(pd.Series)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "450a154ab035"
      },
      "source": [
        "## Exploring traces with the Python SDK for Cloud Trace\n",
        "\n",
        "The Cloud Trace Python SDK provides even more flexibility for working with trace data. We'll use it to demonstrate how to filter traces by date, time, labels, and view types."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d1e6ec298a2e"
      },
      "source": [
        "**Filter by date and time**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "7593b0dd5ab1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "project_id: \"your-project-id\"\n",
            "trace_id: \"0001176b2bd051c5d83d896ed018d7de\"\n",
            "\n",
            "project_id: \"your-project-id\"\n",
            "trace_id: \"001e512a2ffb7611709dfe1de5f3e6cb\"\n",
            "\n",
            "project_id: \"your-project-id\"\n",
            "trace_id: \"0030a937ac00a0650013262c30749dd2\"\n",
            "\n",
            "project_id: \"your-project-id\"\n",
            "trace_id: \"00561cc26879f8eb5967cce3679053f7\"\n",
            "\n",
            "project_id: \"your-project-id\"\n",
            "trace_id: \"006a2e24dc29a02a8c7b02116a471c71\"\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Calculate the start and end times\n",
        "now = datetime.utcnow()\n",
        "yesterday = now - timedelta(hours=24)\n",
        "\n",
        "# Format the dates as ISO 8601 strings with 'Z' for UTC\n",
        "end_time = now.isoformat() + \"Z\"\n",
        "start_time = yesterday.isoformat() + \"Z\"\n",
        "\n",
        "# Request a filtered list of traces by date and time\n",
        "result = client.list_traces(\n",
        "    request=trace.types.ListTracesRequest(\n",
        "        project_id=PROJECT_ID,\n",
        "        start_time=start_time,\n",
        "        end_time=end_time,\n",
        "    )\n",
        ")\n",
        "\n",
        "for count, r in enumerate(result):\n",
        "    if count >= 5:\n",
        "        break\n",
        "    print(r)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0241dfa7e78e"
      },
      "source": [
        "**Filter by label**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "42f9490f60e5"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "project_id: \"your-project-id\"\n",
            "trace_id: \"04bb169965ea79fd1fdc0a31551d6c80\"\n",
            "\n",
            "project_id: \"your-project-id\"\n",
            "trace_id: \"7b3c5fe06c403e251b5da2e8118ffd90\"\n",
            "\n",
            "project_id: \"your-project-id\"\n",
            "trace_id: \"d542c8e017a0e34548174ef58caec674\"\n",
            "\n"
          ]
        }
      ],
      "source": [
        "result = client.list_traces(\n",
        "    request=trace.types.ListTracesRequest(\n",
        "        project_id=PROJECT_ID,\n",
        "        # Return traces where any root span's name starts with AgentExecutor\n",
        "        filter=\"root:AgentExecutor\",\n",
        "    )\n",
        ")\n",
        "\n",
        "for count, r in enumerate(result):\n",
        "    if count >= 5:\n",
        "        break\n",
        "    print(r)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b13e8cc93b5d"
      },
      "source": [
        "**Filter by view type**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "154dfc32ae35"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "project_id: \"your-project-id\"\n",
            "trace_id: \"017c70ff1768ffde39ce7b02d766bf39\"\n",
            "\n",
            "project_id: \"your-project-id\"\n",
            "trace_id: \"01e181b1c375655a861c45dcb048b29b\"\n",
            "\n",
            "project_id: \"your-project-id\"\n",
            "trace_id: \"04bb169965ea79fd1fdc0a31551d6c80\"\n",
            "\n",
            "project_id: \"your-project-id\"\n",
            "trace_id: \"0f4bb7af00a5c772750ad57ba7ead7ee\"\n",
            "\n",
            "project_id: \"your-project-id\"\n",
            "trace_id: \"1125cac78d7858da1f1108b9667c98a8\"\n",
            "\n"
          ]
        }
      ],
      "source": [
        "result = client.list_traces(\n",
        "    request=trace.types.ListTracesRequest(\n",
        "        project_id=PROJECT_ID,\n",
        "        # view=trace.types.ListTracesRequest.ViewType.ROOTSPAN,\n",
        "        view=trace.types.ListTracesRequest.ViewType.MINIMAL,\n",
        "        # view=trace.types.ListTracesRequest.ViewType.COMPLETE,\n",
        "    )\n",
        ")\n",
        "\n",
        "for count, r in enumerate(result):\n",
        "    if count >= 5:\n",
        "        break\n",
        "    print(r)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a4e033321ad"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "After you've finished experimenting, it's a good practice to clean up your cloud resources. You can delete the deployed Reasoning Engine instance to avoid any unexpected charges on your Google Cloud account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "6327bf9c6ca9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Deleting ReasoningEngine : projects/your-project-number/locations/us-central1/reasoningEngines/3146221736555446272\n",
            "ReasoningEngine deleted. . Resource name: projects/your-project-number/locations/us-central1/reasoningEngines/3146221736555446272\n",
            "Deleting ReasoningEngine resource: projects/your-project-number/locations/us-central1/reasoningEngines/3146221736555446272\n",
            "Delete ReasoningEngine backing LRO: projects/your-project-number/locations/us-central1/operations/4247467053650280448\n",
            "ReasoningEngine resource projects/your-project-number/locations/us-central1/reasoningEngines/3146221736555446272 deleted.\n"
          ]
        }
      ],
      "source": [
        "remote_agent.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "tracing_agents_in_reasoning_engine.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}