Mention blocked responses in the UI and fix ChatVertexAI to handle the broken history

Author: aabmass

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

@@ -125,6 +125,11 @@ def render_message(message: BaseMessage, trace_id: str | None) -> None:
         if isinstance(message.content, str)
         else message.content[-1]["text"]
     ).strip()
+
+    # Response was probably blocked by a harm category, go check the trace for details
+    if message.response_metadata.get("is_blocked", False):
+        content = ":red[:material/error: Response blocked, try again]"
+
     if not content:
         return
 

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

@@ -4,6 +4,7 @@
 import tempfile
 import pathlib
 from os import environ
+import logging
 from random import getrandbits
 from typing import cast
 from google.cloud.exceptions import NotFound
@@ -21,13 +22,17 @@
 from langgraph.checkpoint.memory import InMemorySaver
 from langgraph.prebuilt import create_react_agent
 from langgraph_chatbot_demo import _streamlit_helpers
+from langgraph_chatbot_demo.patched_vertexai import PatchedChatVertexAI
 from sqlalchemy import Engine, create_engine
 
 from opentelemetry import trace
 from opentelemetry.trace.span import format_trace_id
 
 from google.cloud import storage
 
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
 _ = """
 Ideas for things to add:
 
@@ -44,7 +49,7 @@
 _streamlit_helpers.styles()
 
 
-model = ChatVertexAI(model="gemini-1.5-flash")
+model = PatchedChatVertexAI(model="gemini-2.0-flash")
 
 if not st.query_params.get("thread_id"):
     result = model.invoke(
@@ -208,6 +213,7 @@ def get_trace_ids(thread_id: str) -> "dict[str, str]":
             # Invoke the agent
             with st.spinner("Thinking..."):
                 res = app.invoke({"messages": [message]}, config=config)
+                logger.debug("agent response", extra={"response": str(res)})
 
             # Store trace ID for rendering
             trace_ids[message.id or ""] = trace_id

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

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from typing import Any
+
+from google.cloud.aiplatform_v1.types import (
+    GenerateContentRequest as v1GenerateContentRequest,
+)
+from google.cloud.aiplatform_v1beta1.types import (
+    GenerateContentRequest,
+)
+from langchain_core.messages import (
+    BaseMessage,
+)
+from langchain_google_vertexai import ChatVertexAI
+
+
+class PatchedChatVertexAI(ChatVertexAI):
+    def _prepare_request_gemini(
+        self, messages: list[BaseMessage], *args: Any, **kwargs: Any
+    ) -> v1GenerateContentRequest | GenerateContentRequest:
+        # Filter out any blocked messages with no content which can appear if you have a blocked
+        # message from finish_reason SAFETY:
+        #
+        # AIMessage(
+        #     content="",
+        #     additional_kwargs={},
+        #     response_metadata={
+        #         "is_blocked": True,
+        #         "safety_ratings": [ ... ],
+        #         "finish_reason": "SAFETY",
+        #     },
+        #     ...
+        # )
+        #
+        # These cause `google.api_core.exceptions.InvalidArgument: 400 Unable to submit request
+        # because it must include at least one parts field`
+
+        messages = [
+            message
+            for message in messages
+            if not message.response_metadata.get("is_blocked", False)
+        ]
+        return super()._prepare_request_gemini(messages, *args, **kwargs)