Show the trace ID inline with messages

Author: aabmass

instrumentation-genai/opentelemetry-instrumentation-vertexai/examples/langgraph-chatbot-demo/run_streamlit.py.lock

@@ -331,7 +331,7 @@ wheels = [
 
 [[package]]
 name = "google-cloud-aiplatform"
-version = "1.79.0"
+version = "1.80.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "docstring-parser" },
@@ -347,9 +347,9 @@ dependencies = [
     { name = "shapely" },
     { name = "typing-extensions" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/a5/8e/93e9f5a7059883c21a82adf8687248c6615d4b833b3bf665631a768b8ebd/google_cloud_aiplatform-1.79.0.tar.gz", hash = "sha256:362bfd16716dcfb6c131736f25246790002b29c99a246fcf4c08a7c71bd2301f", size = 8455732 }
+sdist = { url = "https://files.pythonhosted.org/packages/f0/88/d36384280cc4653e190a4a30025e66b285fbaef06024f68a4264cc588a33/google_cloud_aiplatform-1.80.0.tar.gz", hash = "sha256:bcaa4570a6fb56d3d29cb6b8f92588d4d1a1931de5f90cf07761853dab4c76fd", size = 8459480 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/d9/df/a7629fc1c405ead82249a70903068992932cc5a8c494c396e22995b4429d/google_cloud_aiplatform-1.79.0-py2.py3-none-any.whl", hash = "sha256:e52d518c386ce2b4ce57f1b73b46c57531d9a6ccd70c21a37b349f428bfc1c3f", size = 7086167 },
+    { url = "https://files.pythonhosted.org/packages/b5/57/5e761e7a8b03efc8e7faa4c0b2775991177bbd4dae7a6656a60dfd092ca8/google_cloud_aiplatform-1.80.0-py2.py3-none-any.whl", hash = "sha256:45d2a170f22431dae977551eccb740400bdb899807d0c8d4c16c53b2c1dbc6a5", size = 7089949 },
 ]
 
 [[package]]
@@ -818,7 +818,7 @@ wheels = [
 [[package]]
 name = "langgraph-chatbot-demo"
 version = "0.1.0"
-source = { git = "https://github.com/aabmass/opentelemetry-python-contrib.git?subdirectory=instrumentation-genai%2Fopentelemetry-instrumentation-vertexai%2Fexamples%2Flanggraph-chatbot-demo&branch=vertex-langgraph#441d9f6440e895957ebb982fd5fc1ae6d2731f46" }
+source = { git = "https://github.com/aabmass/opentelemetry-python-contrib.git?subdirectory=instrumentation-genai%2Fopentelemetry-instrumentation-vertexai%2Fexamples%2Flanggraph-chatbot-demo&branch=vertex-langgraph#833be299733a585349eca5bbc3f4977a6760e0be" }
 dependencies = [
     { name = "duckduckgo-search" },
     { name = "langchain-community" },
@@ -1268,7 +1268,7 @@ wheels = [
 [[package]]
 name = "opentelemetry-instrumentation-vertexai"
 version = "2.1b0.dev0"
-source = { git = "https://github.com/aabmass/opentelemetry-python-contrib.git?subdirectory=instrumentation-genai%2Fopentelemetry-instrumentation-vertexai&branch=vertex-langgraph#441d9f6440e895957ebb982fd5fc1ae6d2731f46" }
+source = { git = "https://github.com/aabmass/opentelemetry-python-contrib.git?subdirectory=instrumentation-genai%2Fopentelemetry-instrumentation-vertexai&branch=vertex-langgraph#833be299733a585349eca5bbc3f4977a6760e0be" }
 dependencies = [
     { name = "opentelemetry-api" },
     { name = "opentelemetry-instrumentation" },

instrumentation-genai/opentelemetry-instrumentation-vertexai/examples/langgraph-chatbot-demo/src/langgraph_chatbot_demo/_streamlit_helpers.py

@@ -1,12 +1,21 @@
+from __future__ import annotations
+
 from textwrap import dedent
 
 import google.auth
 import pandas as pd
 import streamlit as st
+from langchain_core.messages.base import BaseMessage
 from langgraph.checkpoint.memory import InMemorySaver
 from sqlalchemy import Engine, inspect
 
 
+@st.cache_data
+def _get_project_id() -> str:
+    project_id: str = google.auth.default()[1]  # type: ignore
+    return project_id
+
+
 def render_intro() -> None:
     st.markdown("""\
     This demo allows you to chat with an Agent that has full access to an
@@ -90,8 +99,41 @@ def on_create() -> None:
         on_change=on_create,
     )
 
-    _, project_id = google.auth.default()  # type: ignore
     st.link_button(
         "View in Google Cloud Trace",
-        f"https://console.cloud.google.com/traces/explorer;query=%7B%22plotType%22:%22HEATMAP%22,%22targetAxis%22:%22Y1%22,%22traceQuery%22:%7B%22resourceContainer%22:%22projects%2F{project_id}%2Flocations%2Fglobal%2FtraceScopes%2F_Default%22,%22spanDataValue%22:%22SPAN_DURATION%22,%22spanFilters%22:%7B%22attributes%22:%5B%5D,%22displayNames%22:%5B%22chain%20invoke%22%5D,%22isRootSpan%22:true,%22kinds%22:%5B%5D,%22maxDuration%22:%22%22,%22minDuration%22:%22%22,%22services%22:%5B%5D,%22status%22:%5B%5D%7D%7D%7D;duration=PT30M",
+        _trace_url(),
+        icon=":material/account_tree:",
     )
+
+
+def styles() -> None:
+    st.html("""\
+    <style>
+        .stChatMessage { padding: 16px; }
+        .stChatMessage .stLinkButton { text-align: right; }
+    </style>
+    """)
+
+
+def render_message(message: BaseMessage, trace_id: str | None) -> None:
+    # Filter out tool calls
+    if not (message.type in ("human", "ai") and message.content):
+        return
+
+    with st.chat_message(message.type):
+        col1, col2 = st.columns([0.9, 0.1])
+        col1.markdown(message.content)
+        if trace_id:
+            col2.link_button(
+                "",
+                _trace_url(trace_id),
+                icon=":material/account_tree:",
+                help="Open in Cloud Trace",
+            )
+
+
+def _trace_url(trace_id: str | None = None) -> str:
+    url = f"https://console.cloud.google.com/traces/explorer;query=%7B%22plotType%22:%22HEATMAP%22,%22targetAxis%22:%22Y1%22,%22traceQuery%22:%7B%22resourceContainer%22:%22projects%2F{_get_project_id()}%2Flocations%2Fglobal%2FtraceScopes%2F_Default%22,%22spanDataValue%22:%22SPAN_DURATION%22,%22spanFilters%22:%7B%22attributes%22:%5B%5D,%22displayNames%22:%5B%22chain%20invoke%22%5D,%22isRootSpan%22:true,%22kinds%22:%5B%5D,%22maxDuration%22:%22%22,%22minDuration%22:%22%22,%22services%22:%5B%5D,%22status%22:%5B%5D%7D%7D%7D;duration=PT30M"
+    if not trace_id:
+        return url
+    return f"{url};traceId={trace_id}"

instrumentation-genai/opentelemetry-instrumentation-vertexai/examples/langgraph-chatbot-demo/src/langgraph_chatbot_demo/langchain_history.py

@@ -1,13 +1,14 @@
 """Adapted from https://github.com/langchain-ai/streamlit-agent/blob/main/streamlit_agent/basic_memory.py"""
 
+import sqlite3
 import tempfile
 from os import environ
 from random import getrandbits
 
 import streamlit as st
 from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
 from langchain_community.utilities.sql_database import SQLDatabase
-from langchain_core.messages import SystemMessage
+from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.messages.base import BaseMessage
 from langchain_core.runnables.config import (
     RunnableConfig,
@@ -35,6 +36,7 @@
 title = "LangGraph SQL Agent Demo"
 st.set_page_config(page_title=title, page_icon="📖", layout="wide")
 st.title(title)
+_streamlit_helpers.styles()
 
 
 model = ChatVertexAI(
@@ -75,12 +77,19 @@ def search(query: str):
 
 
 system_prompt = SystemMessage(
-    content="You are a helpful AI assistant with a mastery of database management and querying. You have "
-    "access to an ephemeral sqlite database that you can query and modify through some tools. Help "
-    "answer questions and perform actions."
-    "\n"
-    "Make sure you always use QuerySQLCheckerTool to validate queries before running them! If you "
-    "make a mistake, try to recover."
+    content=f"""\
+You are a careful and helpful AI assistant with a mastery of database design and querying. You
+have access to an ephemeral sqlite3 database that you can query and modify through some tools.
+Help answer questions and perform actions. Follow these rules:
+
+- Make sure you always use sql_db_query_checker to validate SQL statements **before** running
+  them! In pseudocode: `checked_query = sql_db_query_checker(query);
+  sql_db_query(checked_query)`.
+- The sqlite version is {sqlite3.sqlite_version} which supports multiple row inserts.
+- Always prefer to insert multiple rows in a single call to the sql_db_query tool, if possible.
+- You may request to execute multiple sql_db_query tool calls which will be run in parallel.
+
+If you make a mistake, try to recover."""
 )
 
 
@@ -119,35 +128,49 @@ def get_db(thread_id: str) -> SQLDatabase:
 else:
     messages = []
 
+
+@st.cache_resource
+def get_trace_ids(thread_id: str) -> "dict[str, str]":
+    # Stores the trace IDs. Unfortunately I can't find a way to easily retrieve this from the
+    # checkpointer, so just store it separately.
+    return {}
+
+
+trace_ids = get_trace_ids(st.query_params.thread_id)
+
 col1, col2 = st.columns([0.6, 0.4])
 with col1:
     _streamlit_helpers.render_intro()
     st.divider()
 
     # Add system message
-    st.chat_message(
-        "human", avatar=":material/precision_manufacturing:"
-    ).markdown(f"**System Instructions**\n> {system_prompt.content}")
+    st.expander(
+        "System Instructions", icon=":material/precision_manufacturing:"
+    ).markdown(system_prompt.content)
 
     # Render current messages
-    for msg in messages:
-        # Filter out tool calls
-        if msg.type in ("human", "ai") and msg.content:
-            st.chat_message(msg.type).write(msg.content)
+    for message in messages:
+        trace_id = trace_ids.get(message.id or "")
+        _streamlit_helpers.render_message(message, trace_id)
 
 # If user inputs a new prompt, generate and draw a new response
 if prompt := st.chat_input():
+    message = HumanMessage(prompt)
     with col1:
-        st.chat_message("human").write(prompt)
         with tracer.start_as_current_span(
             "chain invoke",
             attributes={"thread_id": st.query_params.thread_id},
-        ) as span, st.spinner("Thinking..."):
-            st.toast(
-                f"Trace ID {format_trace_id(span.get_span_context().trace_id)}"
-            )
+        ) as span:
+            trace_id = format_trace_id(span.get_span_context().trace_id)
+            _streamlit_helpers.render_message(message, trace_id=trace_id)
+
             # Invoke the agent
-            app.invoke({"messages": [prompt]}, config=config)
+            with st.spinner("Thinking..."):
+                res = app.invoke({"messages": [message]}, config=config)
+
+            # Store trace ID for rendering
+            trace_ids[message.id or ""] = trace_id
+            trace_ids[res["messages"][-1].id] = trace_id
 
     st.rerun()
 
@@ -160,6 +183,3 @@ def get_db(thread_id: str) -> SQLDatabase:
 
     with st.expander("View the message contents in session state"):
         st.json(messages)
-
-# https://pantheon.corp.google.com/traces/explorer;traceId=8a69802b84077b162277f175e0f15276;duration=PT30M?project=otlp-test-deleteme
-# https://pantheon.corp.google.com/traces/explorer;query=%7B%22plotType%22:%22HEATMAP%22,%22targetAxis%22:%22Y1%22,%22traceQuery%22:%7B%22resourceContainer%22:%22projects%2Fotlp-test-deleteme%2Flocations%2Fglobal%2FtraceScopes%2F_Default%22,%22spanDataValue%22:%22SPAN_DURATION%22,%22spanFilters%22:%7B%22attributes%22:%5B%5D,%22displayNames%22:%5B%22chain%20invoke%22%5D,%22isRootSpan%22:true,%22kinds%22:%5B%5D,%22maxDuration%22:%22%22,%22minDuration%22:%22%22,%22services%22:%5B%5D,%22status%22:%5B%5D%7D%7D%7D;traceId=8a69802b84077b162277f175e0f15276;duration=PT30M

instrumentation-genai/opentelemetry-instrumentation-vertexai/examples/langgraph-chatbot-demo/src/langgraph_chatbot_demo/run_streamlit.py

@@ -35,7 +35,10 @@ def run_streamlit() -> None:
     setenv_default("OTEL_RESOURCE_ATTRIBUTES", f"gcp.project_id={project_id}")
 
     # Hide metadata and token refreshes
-    setenv_default("OTEL_PYTHON_EXCLUDED_URLS", "computeMetadata")
+    setenv_default(
+        "OTEL_PYTHON_EXCLUDED_URLS",
+        "computeMetadata,oauth2.googleapis.com",
+    )
 
     langchain_app_spec = importlib.util.find_spec(
         "langgraph_chatbot_demo.langchain_history"