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?"})