chore: import upstream snapshot with attribution

tests/langchain/sample_code/chain.py

@@ -0,0 +1,89 @@
+import dbutils
+
+dbutils.library.restartPython()
+
+import os
+from typing import Any
+
+# `databricks-langchain` versions older than ~0.9 eagerly construct a
+# `WorkspaceClient` inside `ChatDatabricks.__init__`, which requires
+# Databricks credentials. Cross-version test jobs pin those older releases
+# (e.g. 0.8.2 with `langchain==0.3.30`), so set fake creds before the fake
+# chat model is instantiated below.
+os.environ.setdefault("DATABRICKS_HOST", "https://fake-host")
+os.environ.setdefault("DATABRICKS_TOKEN", "fake-token")
+
+from databricks_langchain import ChatDatabricks
+from langchain_community.document_loaders import TextLoader
+from langchain_community.embeddings.fake import FakeEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain_core.callbacks.manager import CallbackManagerForLLMRun
+from langchain_core.messages import AIMessage, BaseMessage
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.outputs import ChatGeneration, ChatResult
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnablePassthrough
+from langchain_text_splitters.character import CharacterTextSplitter
+
+import mlflow
+from mlflow.models import ModelConfig, set_model, set_retriever_schema
+
+base_config = ModelConfig(development_config="tests/langchain/config.yml")
+
+
+def get_fake_chat_model(endpoint="fake-endpoint"):
+    class FakeChatModel(ChatDatabricks):
+        """Fake Chat Model wrapper for testing purposes."""
+
+        endpoint: str = "fake-endpoint"
+
+        def _generate(
+            self,
+            messages: list[BaseMessage],
+            stop: list[str] | None = None,
+            run_manager: CallbackManagerForLLMRun | None = None,
+            **kwargs: Any,
+        ) -> ChatResult:
+            message = AIMessage(content=str(base_config.get("response")))
+            return ChatResult(generations=[ChatGeneration(message=message)])
+
+        @property
+        def _llm_type(self) -> str:
+            return "fake chat model"
+
+    return FakeChatModel(endpoint=endpoint)
+
+
+# No need to define the model, but simulating common practice in dev notebooks
+mlflow.langchain.autolog()
+
+text_path = "tests/langchain/state_of_the_union.txt"
+loader = TextLoader(text_path)
+documents = loader.load()
+text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
+docs = text_splitter.split_documents(documents)
+
+embeddings = FakeEmbeddings(size=base_config.get("embedding_size"))
+vectorstore = FAISS.from_documents(docs, embeddings)
+retriever = vectorstore.as_retriever()
+
+prompt = ChatPromptTemplate.from_template(base_config.get("llm_prompt_template"))
+retrieval_chain = (
+    {
+        "context": retriever,
+        "question": RunnablePassthrough(),
+    }
+    | prompt
+    | get_fake_chat_model()
+    | StrOutputParser()
+)
+
+set_model(retrieval_chain)
+set_retriever_schema(
+    primary_key="primary-key",
+    text_column="text-column",
+    doc_uri="doc-uri",
+    other_columns=["column1", "column2"],
+)
+
+retrieval_chain.invoke({"question": "What is the capital of Japan?"})

tests/langchain/sample_code/chain_with_mlflow_prompt.py

@@ -0,0 +1,13 @@
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+
+import mlflow
+from mlflow.models import set_model
+
+prompt = ChatPromptTemplate.from_template(
+    mlflow.load_prompt("prompts:/qa_prompt@production").to_single_brace_format()
+)
+chain = prompt | ChatOpenAI(temperature=0) | StrOutputParser()
+
+set_model(chain)

tests/langchain/sample_code/config.yml

@@ -0,0 +1,7 @@
+llm_prompt_template: "Answer the following question based on the context: {context}\nQuestion: {question}"
+embedding_size: 5
+response: "Databricks"
+not_used_array:
+  - 1
+  - 2
+  - 3

tests/langchain/sample_code/langchain_chat_agent.py

@@ -0,0 +1,103 @@
+from operator import itemgetter
+from typing import Any, Generator
+
+from langchain_core.messages import AIMessage, AIMessageChunk
+from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnableLambda
+from langchain_core.runnables.base import Runnable
+from langchain_openai import ChatOpenAI
+
+import mlflow
+from mlflow.langchain.output_parsers import ChatAgentOutputParser
+from mlflow.pyfunc.model import ChatAgent
+from mlflow.types.agent import ChatAgentChunk, ChatAgentMessage, ChatAgentResponse, ChatContext
+
+
+class FakeOpenAI(ChatOpenAI, extra="allow"):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self._responses = iter([AIMessage(content="1")])
+        self._stream_responses = iter([
+            AIMessageChunk(content="1"),
+            AIMessageChunk(content="2"),
+            AIMessageChunk(content="3"),
+        ])
+
+    def _generate(self, *args, **kwargs):
+        return ChatResult(generations=[ChatGeneration(message=next(self._responses))])
+
+    def _stream(self, *args, **kwargs):
+        for r in self._stream_responses:
+            yield ChatGenerationChunk(message=r)
+
+
+mlflow.langchain.autolog()
+
+
+# Helper functions
+def extract_user_query_string(messages):
+    return messages[-1]["content"]
+
+
+def extract_chat_history(messages):
+    return messages[:-1]
+
+
+# Define components
+prompt = ChatPromptTemplate.from_template(
+    """Previous conversation:
+{chat_history}
+
+User's question:
+{question}"""
+)
+
+model = FakeOpenAI()
+output_parser = ChatAgentOutputParser()
+
+# Chain definition
+chain = (
+    {
+        "question": itemgetter("messages") | RunnableLambda(extract_user_query_string),
+        "chat_history": itemgetter("messages") | RunnableLambda(extract_chat_history),
+    }
+    | prompt
+    | model
+    | output_parser
+)
+
+
+class LangChainChatAgent(ChatAgent):
+    """
+    Helper class to wrap a LangChain runnable as a :py:class:`ChatAgent <mlflow.pyfunc.ChatAgent>`.
+    Use this class with
+    :py:class:`ChatAgentOutputParser <mlflow.langchain.output_parsers.ChatAgentOutputParser>`.
+    """
+
+    def __init__(self, agent: Runnable):
+        self.agent = agent
+
+    def predict(
+        self,
+        messages: list[ChatAgentMessage],
+        context: ChatContext | None = None,
+        custom_inputs: dict[str, Any] | None = None,
+    ) -> ChatAgentResponse:
+        response = self.agent.invoke({"messages": self._convert_messages_to_dict(messages)})
+        return ChatAgentResponse(**response)
+
+    def predict_stream(
+        self,
+        messages: list[ChatAgentMessage],
+        context: ChatContext | None = None,
+        custom_inputs: dict[str, Any] | None = None,
+    ) -> Generator[ChatAgentChunk, None, None]:
+        for event in self.agent.stream({"messages": self._convert_messages_to_dict(messages)}):
+            yield ChatAgentChunk(**event)
+
+
+chat_agent = LangChainChatAgent(chain)
+
+mlflow.models.set_model(chat_agent)

tests/langchain/sample_code/model_with_config.py

@@ -0,0 +1,18 @@
+from langchain_core.prompts import PromptTemplate
+from langchain_core.runnables import ConfigurableField
+from langchain_openai import ChatOpenAI
+
+from mlflow.models import set_model
+
+model = ChatOpenAI(temperature=0).configurable_fields(
+    temperature=ConfigurableField(
+        id="temperature",
+        name="LLM temperature",
+        description="The temperature of the LLM",
+    )
+)
+
+prompt = PromptTemplate.from_template("Pick a random number above {x}")
+chain = prompt | model
+
+set_model(chain)

tests/langchain/sample_code/no_config/chain.py

@@ -0,0 +1,67 @@
+import os
+from typing import Any
+
+# See `tests/langchain/sample_code/chain.py` for why fake creds are set.
+os.environ.setdefault("DATABRICKS_HOST", "https://fake-host")
+os.environ.setdefault("DATABRICKS_TOKEN", "fake-token")
+
+from databricks_langchain import ChatDatabricks
+from langchain_community.document_loaders import TextLoader
+from langchain_community.embeddings import FakeEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain_core.callbacks.manager import CallbackManagerForLLMRun
+from langchain_core.messages import AIMessage, BaseMessage
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.outputs import ChatGeneration, ChatResult
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnablePassthrough
+from langchain_text_splitters.character import CharacterTextSplitter
+
+from mlflow.models import set_model
+
+
+def get_fake_chat_model(endpoint="fake-endpoint"):
+    class FakeChatModel(ChatDatabricks):
+        """Fake Chat Model wrapper for testing purposes."""
+
+        endpoint: str = "fake-endpoint"
+
+        def _generate(
+            self,
+            messages: list[BaseMessage],
+            stop: list[str] | None = None,
+            run_manager: CallbackManagerForLLMRun | None = None,
+            **kwargs: Any,
+        ) -> ChatResult:
+            return ChatResult(generations=[ChatGeneration(message=AIMessage(content="Databricks"))])
+
+        @property
+        def _llm_type(self) -> str:
+            return "fake chat model"
+
+    return FakeChatModel(endpoint=endpoint)
+
+
+text_path = "tests/langchain/state_of_the_union.txt"
+loader = TextLoader(text_path)
+documents = loader.load()
+text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
+docs = text_splitter.split_documents(documents)
+embeddings = FakeEmbeddings(size=5)
+vectorstore = FAISS.from_documents(docs, embeddings)
+retriever = vectorstore.as_retriever()
+
+prompt = ChatPromptTemplate.from_template(
+    "Answer the following question based on the context: {context}\nQuestion: {question}"
+)
+retrieval_chain = (
+    {
+        "context": retriever,
+        "question": RunnablePassthrough(),
+    }
+    | prompt
+    | get_fake_chat_model()
+    | StrOutputParser()
+)
+
+set_model(retrieval_chain)

tests/langchain/sample_code/openai_agent.py

@@ -0,0 +1,51 @@
+import itertools
+
+from langchain.agents import create_agent
+from langchain.tools import tool
+from langchain_core.messages import AIMessageChunk, ToolCall
+from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
+from langchain_openai import ChatOpenAI
+
+import mlflow
+
+
+class FakeOpenAI(ChatOpenAI, extra="allow"):
+    # In normal LangChain tests, we use the fake OpenAI server to mock the OpenAI REST API.
+    # The fake server returns the input payload as it is. However, for agent tests, the
+    # response should be a specific format so that the agent can parse it correctly.
+    # Also, mocking with mock.patch does not work for testing model serving (as the server
+    # will run in a separate process).
+    # Therefore, we mock the OpenAI client in the model definition here.
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Using itertools.cycle to create an infinite iterator
+        self._responses = itertools.cycle([
+            AIMessageChunk(
+                content="",
+                tool_calls=[ToolCall(name="multiply", args={"a": 2, "b": 3}, id="123")],
+            ),
+            AIMessageChunk(content="The result of 2 * 3 is 6."),
+        ])
+
+    def _generate(self, *args, **kwargs):
+        return ChatResult(generations=[ChatGeneration(message=next(self._responses))])
+
+    def _stream(self, *args, **kwargs):
+        yield ChatGenerationChunk(message=next(self._responses))
+
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add two numbers."""
+    return a + b
+
+
+@tool
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers."""
+    return a * b
+
+
+llm = FakeOpenAI()
+agent = create_agent(llm, [add, multiply], system_prompt="You are a helpful assistant")
+mlflow.models.set_model(agent)

tests/langchain/sample_code/simple_runnable.py

@@ -0,0 +1,14 @@
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import PromptTemplate
+from langchain_openai import ChatOpenAI
+
+import mlflow
+
+prompt = PromptTemplate(
+    input_variables=["product"],
+    template="What is {product}?",
+)
+llm = ChatOpenAI(temperature=0.1, stream_usage=True)
+chain = prompt | llm | StrOutputParser()
+
+mlflow.models.set_model(chain)

tests/langchain/sample_code/workflow.py

@@ -0,0 +1,119 @@
+import json
+import os
+from typing import Any, Sequence
+
+from langchain_core.language_models import LanguageModelLike
+from langchain_core.messages import AIMessage, ToolCall
+from langchain_core.outputs import ChatGeneration, ChatResult
+from langchain_core.runnables import RunnableConfig, RunnableLambda
+from langchain_core.tools import BaseTool, tool
+from langchain_openai import ChatOpenAI
+from langgraph.graph import END, StateGraph
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.prebuilt import ToolNode
+
+import mlflow
+from mlflow.langchain.chat_agent_langgraph import (
+    ChatAgentState,
+    ChatAgentToolNode,
+)
+
+os.environ["OPENAI_API_KEY"] = "test"
+
+
+class FakeOpenAI(ChatOpenAI, extra="allow"):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self._responses = iter([
+            AIMessage(
+                content="",
+                tool_calls=[ToolCall(name="uc_tool_format", args={}, id="123")],
+            ),
+            AIMessage(
+                content="",
+                tool_calls=[ToolCall(name="lc_tool_format", args={}, id="456")],
+            ),
+            AIMessage(content="Successfully generated", id="789"),
+        ])
+
+    def _generate(self, *args, **kwargs):
+        return ChatResult(generations=[ChatGeneration(message=next(self._responses))])
+
+
+@tool
+def uc_tool_format() -> str:
+    """Returns uc tool format"""
+    return json.dumps({
+        "format": "SCALAR",
+        "value": '{"content":"hi","attachments":{"a":"b"},"custom_outputs":{"c":"d"}}',
+        "truncated": False,
+    })
+
+
+@tool
+def lc_tool_format() -> dict[str, Any]:
+    """Returns lc tool format"""
+    nums = [1, 2]
+    return {
+        "content": f"Successfully generated array of 2 random ints: {nums}.",
+        "attachments": {"key1": "attach1", "key2": "attach2"},
+        "custom_outputs": {"random_nums": nums},
+    }
+
+
+tools = [uc_tool_format, lc_tool_format]
+
+
+def create_tool_calling_agent(
+    model: LanguageModelLike,
+    tools: ToolNode | Sequence[BaseTool],
+    agent_prompt: str | None = None,
+) -> CompiledStateGraph:
+    model = model.bind_tools(tools)
+
+    def should_continue(state: ChatAgentState):
+        messages = state["messages"]
+        last_message = messages[-1]
+        # If there are function calls, continue. else, end
+        if last_message.get("tool_calls"):
+            return "continue"
+        else:
+            return "end"
+
+    preprocessor = RunnableLambda(lambda state: state["messages"])
+    model_runnable = preprocessor | model
+
+    @mlflow.trace
+    def call_model(
+        state: ChatAgentState,
+        config: RunnableConfig,
+    ):
+        response = model_runnable.invoke(state, config)
+
+        return {"messages": [response]}
+
+    workflow = StateGraph(ChatAgentState)
+
+    workflow.add_node("agent", RunnableLambda(call_model))
+    workflow.add_node("tools", ChatAgentToolNode(tools))
+
+    workflow.set_entry_point("agent")
+    workflow.add_conditional_edges(
+        "agent",
+        should_continue,
+        {
+            "continue": "tools",
+            "end": END,
+        },
+    )
+    workflow.add_edge("tools", "agent")
+
+    return workflow.compile()
+
+
+mlflow.langchain.autolog()
+llm = FakeOpenAI()
+graph = create_tool_calling_agent(llm, tools)
+
+mlflow.models.set_model(graph)