chore: import upstream snapshot with attribution

tests/langgraph/sample_code/langgraph_with_custom_span.py

@@ -0,0 +1,59 @@
+from typing import Literal
+
+from langchain_core.messages import AIMessage, ToolCall
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.outputs import ChatGeneration, ChatResult
+from langchain_core.prompts import PromptTemplate
+from langchain_core.tools import tool
+from langchain_openai import ChatOpenAI
+from langgraph.prebuilt import create_react_agent
+
+import mlflow
+from mlflow.entities.span import SpanType
+
+
+class FakeOpenAI(ChatOpenAI, extra="allow"):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self._responses = iter([
+            AIMessage(
+                content="",
+                tool_calls=[ToolCall(name="get_weather", args={"city": "sf"}, id="123")],
+            ),
+            AIMessage(content="The weather in San Francisco is always sunny!"),
+        ])
+
+    def _generate(self, *args, **kwargs):
+        return ChatResult(generations=[ChatGeneration(message=next(self._responses))])
+
+
+def get_inner_runnable():
+    llm = ChatOpenAI()
+    prompt = PromptTemplate.from_template("what is the weather in {city}?")
+    return prompt | llm | StrOutputParser()
+
+
+@tool
+def get_weather(city: Literal["nyc", "sf"]):
+    """Use this to get weather information."""
+    with mlflow.start_span(name="get_weather_inner", span_type=SpanType.CHAIN) as span:
+        span.set_inputs(city)
+
+        # Call another LangChain module
+        inner_runnable = get_inner_runnable()
+        inner_runnable.invoke({"city": city})
+
+        if city == "nyc":
+            output = "It might be cloudy in nyc"
+        elif city == "sf":
+            output = "It's always sunny in sf"
+        span.set_outputs(output)
+    return output
+
+
+llm = FakeOpenAI()
+tools = [get_weather]
+graph = create_react_agent(llm, tools)
+
+mlflow.models.set_model(graph)