Coding a Chain with LangChain and MLflow pyfunc

1.Composing a chain with LangChain

LangChain links chain components with the pipe operator. Each component's output becomes the next component's input. There are three parts: a prompt template with named placeholders, a chat model, and a final step that converts the model's message into a plain string. The example below comes from the Databricks MLflow tracing documentation and uses an OpenAI chat model. In a Databricks chain, the llm slot is usually filled by ChatDatabricks(endpoint=...) from the databricks-langchain package, which is covered in the next section.

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7, max_tokens=1000)

prompt_template = PromptTemplate.from_template(
    "Answer the question as if you are {person}, fully embodying their style, wit, personality, and habits of speech. "
    "Emulate their quirks and mannerisms to the best of your ability, embracing their traits—even if they aren't entirely "
    "constructive or inoffensive. The question is: {question}"
)

chain = prompt_template | llm |  ? ()

# Let's test another call
chain.invoke(
    {
        "person": "Linus Torvalds",
        "question": "Can I just set everyone's access to sudo to make things easier?",
    }
)

The input to invoke is a dictionary whose keys match the template's placeholders. A missing key means the prompt cannot be filled. The same page also warns against hardcoding provider API keys in chain code. For production, use AI Gateway or Databricks secrets.

Sources1

2.Adding the retrieval step

A RAG chain also needs a retriever. Databricks AI Bridge packages such as databricks-langchain include VectorSearchRetrieverTool, which wraps an existing index, and ChatDatabricks, which calls a Databricks serving endpoint. You can test the tool by calling invoke on it directly.

from databricks_langchain import VectorSearchRetrieverTool, ChatDatabricks

# Initialize the retriever tool.
vs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.my_databricks_docs_index",
  tool_name="databricks_docs_retriever",
  tool_description="Retrieves information about Databricks products from official Databricks documentation."
)

# Run a query against the vector search index locally for testing
vs_tool.invoke("Databricks Agent Framework?")

# Bind the retriever tool to your Langchain LLM of choice
llm = ChatDatabricks(endpoint="databricks-claude-sonnet-4-5")
llm_with_tools = llm.bind_tools([vs_tool])

The last two lines matter for design. bind_tools lets the LLM decide whether to call the retriever, which is how a single agent works. In a deterministic chain, your own code calls the retriever every time and passes the results into the prompt. If the index uses self-managed embeddings, or is a direct-access index, the tool also needs a text_column and an embedding model, for example DatabricksEmbeddings(endpoint="databricks-bge-large-en").

A deterministic RAG chain follows the same fixed order for every request: retrieve the top-k results from the index, augment a prompt by combining the user request with the retrieved context, then generate a response by sending the augmented prompt to an LLM. In LangChain terms, your code fetches the chunks first, supplies them together with the question as the placeholders of a prompt template, and then pipes that template into the chat model and the output parser, exactly as in the previous section. The LLM never decides whether to retrieve. That fixed flow is what makes a deterministic chain predictable and easy to audit.

You can also skip LangChain and query the index yourself with index.similarity_search. Which argument you pass depends on how the index got its embeddings.

# Delta Sync Index with pre-calculated embeddings
results2 = index.similarity_search(
     ? =[0.9] * 1024,
    columns=["id", "text"],
    num_results=2
    )

Sources23

3.Wrapping custom logic as a pyfunc model

If a requirement calls for logic that a framework does not provide, such as reformatting inputs before the model or cleaning raw outputs afterwards, MLflow's pyfunc can package any Python code as a model. The Databricks guidance lists cases where pyfunc fits: the model needs preprocessing, its raw outputs need post-processing, it has branching logic per request, or you want to deploy fully custom code.

class CustomModel(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        self.model = torch.load(context.artifacts["model-weights"])
        from preprocessing_utils.my_custom_tokenizer import CustomTokenizer
        self.tokenizer = CustomTokenizer(context.artifacts["tokenizer_cache"])

    def format_inputs(self, model_input):
        # insert some code that formats your inputs
        pass

    def format_outputs(self, outputs):
        predictions = (torch.sigmoid(outputs)).data.numpy()
        return predictions

    def predict(self, context, model_input):
        model_input = self.format_inputs(model_input)
        outputs = self.model.predict(model_input)
        return self.format_outputs(outputs)

The two required methods split the work. load_context holds everything that only needs loading once, such as weights and tokenizers, so that each request has less to load. predict runs on every request. In this example it calls the preprocessing helper format_inputs, then the model, then the post-processing helper format_outputs. A chain fits the same shape: the optional query preprocessing step of a RAG chain maps to the input formatting, and the optional post-processing step, such as adding citations or applying business rules to the LLM's response, maps to the output formatting. Before deploying, you can check that the model can be served by using mlflow.models.predict. After logging it, you can register it and serve it from a Model Serving endpoint.

Logging matters for serving. Databricks Runtime ML includes mlflow-skinny by default, and if you log without specifying pip_requirements, MLflow captures mlflow-skinny in the model's conda.yaml. Model Serving requires mlflow instead and cannot build the container image otherwise, so always pin mlflow explicitly when you call mlflow.pyfunc.log_model().

# DBR ML ships with mlflow-skinny by default, so specify mlflow explicitly
# to ensure Model Serving compatibility.
mlflow.pyfunc.log_model(
    name="model",
    python_model=your_model,
    pip_requirements=["mlflow==3.8.1"],  # use mlflow, not mlflow-skinny
    registered_model_name="catalog.schema.model_name",
)

Sources45

Exam traps

Each one states something that sounds right. Open it to see what is actually true.

1. 1.In a pyfunc model, load_context runs on every request alongside predict.Why is that wrong?

   load_context is for things loaded once. predict contains the logic that runs on each request.

   Covered in Wrapping custom logic as a pyfunc model

2. 2.VectorSearchRetrieverTool needs only an index_name, whatever kind of index it points to.Why is that wrong?

   Direct-access indexes and Delta Sync indexes with self-managed embeddings also need a custom embedding model and a text column.

   Covered in Adding the retrieval step

Practise it for real

Prototype the retrieval component of a chain against an existing index, then pipe a prompt into a Databricks-served LLM so the chain retrieves on every request

1. 1.Run %pip install --upgrade databricks-langchain in a notebook

   Why: This package contains Databricks AI Bridge, including VectorSearchRetrieverTool and ChatDatabricks

   You should see: The package installs and both classes can be imported

2. 2.Create VectorSearchRetrieverTool with your index_name (catalog.schema.index), a tool_name and a descriptive tool_description

   Why: The retriever wraps an existing index. A self-managed embeddings index also needs text_column and embedding

   You should see: The tool object is created without errors

3. 3.Call vs_tool.invoke("<a question your documents answer>")

   Why: This tests retrieval on its own before any LLM is involved

   You should see: Relevant chunks come back from the index

4. 4.Create ChatDatabricks(endpoint="databricks-claude-sonnet-4-5"), a PromptTemplate with placeholders for the question and the retrieved context, and compose prompt_template | llm | StrOutputParser()

   Why: This is the pipe-operator chain from the first section, now with the LLM served by Databricks. Your code, not the LLM, decides that retrieval happens first

   You should see: The chain object is created without errors

5. 5.Pass the result of vs_tool.invoke and your question into chain.invoke as the template's placeholder values

   Why: Retrieve, augment, generate in a fixed order is a deterministic RAG chain. Calling bind_tools instead would let the LLM decide, which is the agent pattern

   You should see: A plain string answer that draws on the retrieved documentation