RAG Application Building Blocks: Flavor, Embeddings, Retriever, Signature

1.The elements a RAG chain needs

Each element in this objective comes from a step in the RAG chain. Databricks describes the simplest RAG agent in three steps. Retrieval queries an outside knowledge base. Augmentation combines what was retrieved with the user's request in a prompt template. Generation sends that prompt to an LLM. Retrieval in turn needs two things: an embedding model to turn the question into a vector, and a retriever to look up similar chunks. The chain can also have optional pre-processing of the user query and post-processing of the LLM's answer.

To deploy the chain you package all of it as an MLflow Model. The package needs a flavor, a list of dependencies, an input example and a signature. The table below shows which step each element serves.

Sources12

2.Choosing the model flavor: built-in (LangChain) or custom pyfunc

For a framework that has a built-in flavor, log the chain through that flavor's module. MLflow then handles the framework's specifics for you, including the main point of this objective: it reliably records the framework's own dependencies. The example below logs a LangChain chain. It also passes an input example and some parameters, which come up again later in this lesson.

model_info = mlflow.langchain.log_model(
  lc_model=chain,
  name="basic_chain",
  params={
    "temperature": 0.1,
    "max_tokens": 2000,
    "prompt_template": str(prompt)
  },
  model_type="agent",
  input_example={"messages": "What is MLflow?"},
)

Use the generic python-function (pyfunc) flavor when no built-in flavor fits, or when you need your own pre-processing and post-processing around the chain. Subclass mlflow.pyfunc.PythonModel. Its load_context method sets up the clients and resources, and its predict method runs the logic. Writing a whole custom flavor is also possible, but the docs say it takes more work. Every Python MLflow model, whatever its flavor, can be loaded as a generic Python function with mlflow.pyfunc.load_model().

Sources345

3.Choosing the embedding model

No. Retrieval compares the query vector with the chunk vectors, and that comparison only works if both came from the same embedding model. The cookbook says the retrieval query is translated into an embedding using the same embedding model that was used to embed the document chunks during data preparation. So the embedding model is chosen once, for both sides of the comparison.

When you create a Databricks AI Search index (formerly Vector Search), you also decide who computes the embeddings. That decision limits what you can change later.

If none of the hosted embedding endpoints suits you, you can bring your own model, packaged with the same MLflow mechanics as the chain. Log it as a custom Python model. Its input is a single string column, and its output is a tensor of embeddings returned as a NumPy array.

Sources267

4.Configuring the retriever

The retriever is the component the chain calls to fetch supporting chunks. For Databricks AI Search indexes, the AI Bridge packages databricks-langchain and databricks-openai provide VectorSearchRetrieverTool, plus helpers such as from_vector_search and from_uc_function that build retrievers from existing Databricks resources. On the tool you set how many results to return, which columns to return, filters, and the query type: ANN or HYBRID. Hybrid search helps when the source data contains exact keywords such as SKUs or identifiers.

The embedding model comes up again here. With a managed-embedding index, Databricks embeds the query for you. With a direct-access index, or a Delta Sync index using self-managed embeddings, the retriever cannot know which model to use, so you have to pass it the embedding model and the text column.

embedding_model = DatabricksEmbeddings(
    endpoint="databricks-bge-large-en",
)

vs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.index_name", # Index name in the format 'catalog.schema.index'
  num_results=5, # Max number of documents to return
  columns=["primary_key", "text_column"], # List of columns to include in the search
  filters={"text_column LIKE": "Databricks"}, # Filters to apply to the query
  query_type="ANN", # Query type ("ANN" or "HYBRID").
  tool_name="name of the tool", # Used by the LLM to understand the purpose of the tool
  tool_description="Purpose of the tool", # Used by the LLM to understand the purpose of the tool
  text_column="text_column", # Specify text column for embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.
  embedding=embedding_model # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

If you write your own retriever, for example in a pyfunc-flavored agent that calls a vector store hosted outside Databricks, what it returns matters too. To fit the MLflow retriever schema, it should return a list of MLflow Document objects. Extra attributes such as doc_uri or a similarity score go in each Document's metadata field.

Sources86

5.Declaring dependencies

During deployment, Model Serving builds a production container from what the MLflow model declares. The base image may include some system-level packages, but your application-level dependencies must be declared explicitly in the model. Logging a model automatically writes requirements.txt and conda.yaml. With mlflow.pyfunc.log_model, MLflow infers the requirements using mlflow.models.infer_pip_requirements. Inference doesn't always catch everything, so MLflow gives you several ways to fill the gaps.

Pin exact versions, for example nltk==<version> rather than just nltk. One trap is specific to Databricks. Databricks Runtime ML includes mlflow-skinny rather than the full mlflow package. If you log a pyfunc model there without pip_requirements, the model's conda.yaml records mlflow-skinny, and Model Serving cannot build the container image because it needs mlflow. The fix is to pin mlflow==<version> in pip_requirements. Databricks also recommends running inference on the same runtime version you used to create the model.

mlflow.pyfunc.log_model(
   name=name,
    ? =[filename.py],
   data_path=data_path,
   conda_env=conda_env,
)

Sources94

6.Input examples and model signatures

The last two elements describe the model's interface. An input example is a sample request saved with the model, such as the {"messages": "What is MLflow?"} passed to the LangChain log_model call earlier. A signature formally declares the input and output schema. One way to get a signature is to infer it from sample data and the model's output, then pass it to log_model, as the custom models overview shows.

You can also build the signature by hand. The custom embedding model from earlier does this, and the same call brings together the elements this lesson has covered. The input schema is one string column (ColSpec) and the output is a float32 tensor (TensorSpec). The model uses the pyfunc flavor, its requirements are pinned and include mlflow==<version> as discussed above, and it carries both a signature and an input example.

with mlflow.start_run() as run:
    input_schema = Schema([ColSpec(DataType.string)])
    output_schema = Schema([TensorSpec(np.dtype(np.float32), [-1, 1024])])
    signature = ModelSignature(inputs=input_schema, outputs=output_schema)
    mlflow.pyfunc.log_model(
        artifact_path="model",
        python_model=CustomEmbeddingModel(),
        pip_requirements=["mlflow==2.21.3", "openai==1.69.0"],
        signature=signature,
        input_example=input_example
    )

The signature is not optional if the model is going to Unity Catalog. The documentation states plainly that Models in Unity Catalog require a signature, and model versions registered without one have limitations. Registering and serving the model are covered in neighbouring lessons. For this objective, what matters is to log the signature with the chain from the start.

Sources107

Exam traps

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

1. 1.You can embed user queries with any embedding model, even a different one from the model that built the index.Why is that wrong?

   Similarity is only meaningful when the query and the chunks were embedded by the same model, so the query must use the index's embedding model.

   Covered in Choosing the embedding model

2. 2.MLflow's automatic inference captures every dependency, so logging a pyfunc model on Databricks Runtime ML with default settings is always servable.Why is that wrong?

   On ML runtimes MLflow records mlflow-skinny, and Model Serving cannot build the image unless you pin mlflow in pip_requirements. Non-Python dependencies are never picked up automatically either.

   Covered in Declaring dependencies

3. 3.A model signature is optional metadata. An input example alone is enough to register the chain in Unity Catalog.Why is that wrong?

   Unity Catalog requires a signature, and versions registered without one have limitations.

   Covered in Input examples and model signatures

4. 4.VectorSearchRetrieverTool always embeds the query for you, whatever kind of index it points at.Why is that wrong?

   For direct-access indexes and Delta Sync indexes with self-managed embeddings, you must supply the embedding model and the text column yourself.

   Covered in Configuring the retriever