WIP multimodal

Author: aabmass

instrumentation-genai/opentelemetry-instrumentation-vertexai/examples/langgraph-chatbot-demo/pyproject.toml

@@ -6,6 +6,7 @@ readme = "README.md"
 requires-python = ">=3.9"
 dependencies = [
     "duckduckgo-search>=7.3.0",
+    "google-cloud-storage[tracing]>=2.19.0",
     "langchain-community>=0.3.16",
     "langchain-core>=0.3.31",
     "langchain-google-vertexai>=2.0.7",

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

@@ -11,7 +11,7 @@
 
 
 @st.cache_data
-def _get_project_id() -> str:
+def get_project_id() -> str:
     project_id: str = google.auth.default()[1]  # type: ignore
     return project_id
 
@@ -117,12 +117,20 @@ def styles() -> None:
 
 def render_message(message: BaseMessage, trace_id: str | None) -> None:
     # Filter out tool calls
-    if not (message.type in ("human", "ai") and message.content):
+    if message.type not in ("human", "ai"):
+        return
+
+    content = (
+        message.content
+        if isinstance(message.content, str)
+        else message.content[-1]["text"]
+    ).strip()
+    if not content:
         return
 
     with st.chat_message(message.type):
         col1, col2 = st.columns([0.9, 0.1])
-        col1.markdown(message.content)
+        col1.markdown(content)
         if trace_id:
             col2.link_button(
                 "",
@@ -133,7 +141,7 @@ def render_message(message: BaseMessage, trace_id: str | None) -> None:
 
 
 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"
+    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

@@ -2,8 +2,11 @@
 
 import sqlite3
 import tempfile
+import pathlib
 from os import environ
 from random import getrandbits
+from typing import cast
+from google.cloud.exceptions import NotFound
 
 import streamlit as st
 from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
@@ -23,6 +26,8 @@
 from opentelemetry import trace
 from opentelemetry.trace.span import format_trace_id
 
+from google.cloud import storage
+
 _ = """
 Ideas for things to add:
 
@@ -39,10 +44,7 @@
 _streamlit_helpers.styles()
 
 
-model = ChatVertexAI(
-    model="gemini-1.5-flash",
-    project=environ.get("GOOGLE_CLOUD_PROJECT", None),
-)
+model = ChatVertexAI(model="gemini-1.5-flash")
 
 if not st.query_params.get("thread_id"):
     result = model.invoke(
@@ -52,6 +54,8 @@
         seed=getrandbits(31),
     )
     st.query_params.thread_id = str(result.content).strip()
+if "upload_key" not in st.session_state:
+    st.session_state.upload_key = 0
 
 
 # Initialize memory to persist state between graph runs
@@ -65,6 +69,21 @@ def get_checkpointer() -> InMemorySaver:
     _streamlit_helpers.render_sidebar(checkpointer)
 
 
+@st.cache_resource
+def get_storage_bucket() -> storage.Bucket:
+    storage_client = storage.Client()
+    bucket_name = (
+        f"{_streamlit_helpers.get_project_id()}-langgraph-chatbot-storage"
+    )
+    try:
+        return storage_client.get_bucket(bucket_name)
+    except NotFound:
+        return storage_client.create_bucket(bucket_name)
+
+
+bucket = get_storage_bucket()
+
+
 # Define the tools for the agent to use
 @tool
 @tracer.start_as_current_span("tool search")
@@ -154,8 +173,30 @@ def get_trace_ids(thread_id: str) -> "dict[str, str]":
         _streamlit_helpers.render_message(message, trace_id)
 
 # If user inputs a new prompt, generate and draw a new response
+# TODO: see if st.form() looks better
+file_upload = st.file_uploader(
+    "Upload an image",
+    type=["png", "jpg", "jpeg", "pdf", "webp"],
+    # Hack to clear the upload
+    key=f"file_uploader_{st.session_state.upload_key}",
+)
 if prompt := st.chat_input():
-    message = HumanMessage(prompt)
+    content = []
+
+    # Put the image first https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/image-understanding#best-practices
+    if file_upload:
+        filename: str = file_upload.name
+        blob = bucket.blob(filename)
+        blob.upload_from_file(file_upload, content_type=file_upload.type)
+        st.session_state.upload_key += 1
+
+        uri = f"gs://{bucket.name}/{blob.name}"
+        content.append({"type": "image_url", "image_url": {"url": uri}})
+
+    content.append({"type": "text", "text": prompt})
+
+    message = HumanMessage(content)
+
     with col1:
         with tracer.start_as_current_span(
             "chain invoke",

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

@@ -33,6 +33,7 @@ def run_streamlit() -> None:
         "OTEL_PYTHON_EXCLUDED_URLS",
         "computeMetadata,oauth2.googleapis.com",
     )
+    # subprocess.run(["opentelemetry-instrument", "ipython"])
 
     langchain_app_spec = importlib.util.find_spec(
         "langgraph_chatbot_demo.langchain_history"