Langchain examples for testing

and send with otlp and full chatbot demo

---

So... turns out Langchain uses the v1beta1 prediction service client
under the hood directly.. So we should probably instrument that after
all instead of the main wrapper API.

It also has a streaming option so we should try to support that as well,
and it has ainvoke() for asyncio.

Author: aabmass

instrumentation-genai/opentelemetry-instrumentation-vertexai/examples/langchain/chatbot.py

@@ -0,0 +1,259 @@
+# /// script
+# dependencies = [
+#     "langchain-core",
+#     "langchain-google-vertexai>=2.0.9",
+#     "langgraph>0.2.27",
+#     "opentelemetry-distro",
+#     "opentelemetry-exporter-otlp-proto-grpc",
+#     "opentelemetry-instrumentation-vertexai",
+# ]
+#
+# [tool.uv.sources]
+# opentelemetry-instrumentation-vertexai = { git = "https://github.com/aabmass/opentelemetry-python-contrib.git", subdirectory = "instrumentation-genai/opentelemetry-instrumentation-vertexai", branch = "vertex-langgraph" }
+#
+# ///
+
+# https://python.langchain.com/docs/tutorials/chatbot
+
+# Load auto instrumentation
+from opentelemetry.instrumentation.auto_instrumentation import sitecustomize  # noqa: F401, I001
+
+from typing import Sequence
+
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+    HumanMessage,
+    SystemMessage,
+    trim_messages,
+)
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_google_vertexai import ChatVertexAI
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import START, MessagesState, StateGraph
+from langgraph.graph.message import add_messages
+from typing_extensions import Annotated, TypedDict
+
+from opentelemetry import trace
+
+
+def main() -> None:
+    model = ChatVertexAI(model="gemini-1.5-flash")
+
+    # Define a new graph
+    workflow = StateGraph(state_schema=MessagesState)
+
+    # Define the function that calls the model
+    def call_model(state: MessagesState):
+        response = model.invoke(state["messages"])
+        return {"messages": response}
+
+    # Define the (single) node in the graph
+    workflow.add_edge(START, "model")
+    workflow.add_node("model", call_model)
+
+    # Add memory
+    memory = MemorySaver()
+    app = workflow.compile(checkpointer=memory)
+
+    config = {"configurable": {"thread_id": "abc123"}}
+
+    query = "Hi! I'm Bob."
+
+    input_messages = [HumanMessage(query)]
+    output = app.invoke({"messages": input_messages}, config)
+    output["messages"][
+        -1
+    ].pretty_print()  # output contains all messages in state
+
+    query = "What's my name?"
+
+    input_messages = [HumanMessage(query)]
+    output = app.invoke({"messages": input_messages}, config)
+    output["messages"][-1].pretty_print()
+
+    config = {"configurable": {"thread_id": "abc234"}}
+
+    input_messages = [HumanMessage(query)]
+    output = app.invoke({"messages": input_messages}, config)
+    output["messages"][-1].pretty_print()
+
+    config = {"configurable": {"thread_id": "abc123"}}
+
+    input_messages = [HumanMessage(query)]
+    output = app.invoke({"messages": input_messages}, config)
+    output["messages"][-1].pretty_print()
+
+    prompt_template = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                "You talk like a pirate. Answer all questions to the best of your ability.",
+            ),
+            MessagesPlaceholder(variable_name="messages"),
+        ]
+    )
+
+    workflow = StateGraph(state_schema=MessagesState)
+
+    def call_model(state: MessagesState):
+        # highlight-start
+        prompt = prompt_template.invoke(state)
+        response = model.invoke(prompt)
+        # highlight-end
+        return {"messages": response}
+
+    workflow.add_edge(START, "model")
+    workflow.add_node("model", call_model)
+
+    memory = MemorySaver()
+    app = workflow.compile(checkpointer=memory)
+
+    config = {"configurable": {"thread_id": "abc345"}}
+    query = "Hi! I'm Jim."
+
+    input_messages = [HumanMessage(query)]
+    output = app.invoke({"messages": input_messages}, config)
+    output["messages"][-1].pretty_print()
+
+    query = "What is my name?"
+
+    input_messages = [HumanMessage(query)]
+    output = app.invoke({"messages": input_messages}, config)
+    output["messages"][-1].pretty_print()
+
+    prompt_template = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                "You are a helpful assistant. Answer all questions to the best of your ability in {language}.",
+            ),
+            MessagesPlaceholder(variable_name="messages"),
+        ]
+    )
+
+    # highlight-next-line
+    class State(TypedDict):
+        # highlight-next-line
+        messages: Annotated[Sequence[BaseMessage], add_messages]
+        # highlight-next-line
+        language: str
+
+    workflow = StateGraph(state_schema=State)
+
+    def call_model(state: State):
+        prompt = prompt_template.invoke(state)
+        response = model.invoke(prompt)
+        return {"messages": [response]}
+
+    workflow.add_edge(START, "model")
+    workflow.add_node("model", call_model)
+
+    memory = MemorySaver()
+    app = workflow.compile(checkpointer=memory)
+
+    config = {"configurable": {"thread_id": "abc456"}}
+    query = "Hi! I'm Bob."
+    language = "Spanish"
+
+    input_messages = [HumanMessage(query)]
+    output = app.invoke(
+        # highlight-next-line
+        {"messages": input_messages, "language": language},
+        config,
+    )
+    output["messages"][-1].pretty_print()
+
+    query = "What is my name?"
+
+    input_messages = [HumanMessage(query)]
+    output = app.invoke(
+        {"messages": input_messages},
+        config,
+    )
+    output["messages"][-1].pretty_print()
+
+    trimmer = trim_messages(
+        max_tokens=65,
+        strategy="last",
+        token_counter=model,
+        include_system=True,
+        allow_partial=False,
+        start_on="human",
+    )
+
+    messages = [
+        SystemMessage(content="you're a good assistant"),
+        HumanMessage(content="hi! I'm bob"),
+        AIMessage(content="hi!"),
+        HumanMessage(content="I like vanilla ice cream"),
+        AIMessage(content="nice"),
+        HumanMessage(content="whats 2 + 2"),
+        AIMessage(content="4"),
+        HumanMessage(content="thanks"),
+        AIMessage(content="no problem!"),
+        HumanMessage(content="having fun?"),
+        AIMessage(content="yes!"),
+    ]
+
+    trimmer.invoke(messages)
+
+    workflow = StateGraph(state_schema=State)
+
+    def call_model(state: State):
+        # highlight-start
+        trimmed_messages = trimmer.invoke(state["messages"])
+        prompt = prompt_template.invoke(
+            {"messages": trimmed_messages, "language": state["language"]}
+        )
+        response = model.invoke(prompt)
+        # highlight-end
+        return {"messages": [response]}
+
+    workflow.add_edge(START, "model")
+    workflow.add_node("model", call_model)
+
+    memory = MemorySaver()
+    app = workflow.compile(checkpointer=memory)
+
+    config = {"configurable": {"thread_id": "abc567"}}
+    query = "What is my name?"
+    language = "English"
+
+    # highlight-next-line
+    input_messages = messages + [HumanMessage(query)]
+    output = app.invoke(
+        {"messages": input_messages, "language": language},
+        config,
+    )
+    output["messages"][-1].pretty_print()
+
+    config = {"configurable": {"thread_id": "abc678"}}
+    query = "What math problem did I ask?"
+    language = "English"
+
+    input_messages = messages + [HumanMessage(query)]
+    output = app.invoke(
+        {"messages": input_messages, "language": language},
+        config,
+    )
+    output["messages"][-1].pretty_print()
+
+    config = {"configurable": {"thread_id": "abc789"}}
+    query = "Hi I'm Todd, please tell me a joke."
+    language = "English"
+
+    input_messages = [HumanMessage(query)]
+    # highlight-next-line
+    for chunk, metadata in app.stream(
+        {"messages": input_messages, "language": language},
+        config,
+        # highlight-next-line
+        stream_mode="messages",
+    ):
+        if isinstance(chunk, AIMessage):  # Filter to just model responses
+            print(chunk.content, end="|")
+
+
+with trace.get_tracer(__name__).start_as_current_span("demo-root-span"):
+    main()