databricks-logo

dspy-pyfunc-simple-agent

(Python)

Mosaic AI Agent Framework: Author and deploy a simple DSPy agent

This notebook demonstrates how to author a DSPy agent that's compatible with Mosaic AI Agent Framework features. You learn how to:

  • Author a DSPy agent with ChatAgent.
  • Manually test the agent's output.
  • Log and deploy the agent.

To learn more about authoring an agent using Mosaic AI Agent Framework, see Databricks documentation (AWS | Azure).

Prerequisites

  • Address all TODOs in this notebook.

Limitations

  • The agent in this example does not support streaming output.
2
%pip install -U -qqqq mlflow[databricks] dspy databricks-agents uv matplotlib
dbutils.library.restartPython()
Note: you may need to restart the kernel using %restart_python or dbutils.library.restartPython() to use updated packages.

Define the agent in code

Define the agent code in a single cell below. Then, you can write the agent code to a local Python file, using the %%writefile magic command, for subsequent logging and deployment.

4
%%writefile agent.py

from typing import Any, Generator, Optional
import mlflow
from databricks.sdk import WorkspaceClient
from mlflow.entities import SpanType
from mlflow.pyfunc.model import ChatAgent
from mlflow.types.agent import (
    ChatAgentMessage,
    ChatAgentResponse,
    ChatContext,
)
import dspy
import uuid

# Autolog DSPy traces to MLflow
mlflow.dspy.autolog()

# Set up DSPy with a Databricks-hosted LLM
LLM_ENDPOINT_NAME = "databricks-meta-llama-3-3-70b-instruct"
lm = dspy.LM(model=f"databricks/{LLM_ENDPOINT_NAME}", max_tokens=250)
dspy.settings.configure(lm=lm)


class DSPyChatAgent(ChatAgent):     
    def __init__(self):
        self.agent = dspy.ChainOfThought("question,history -> answer")


    def prepare_message_history(self, messages: list[ChatAgentMessage]):
        history_entries = []
        # Assume the last message in the input is the most recent user question.
        for i in range(0, len(messages) - 1, 2):
            history_entries.append({"question": messages[i].content, "answer": messages[i + 1].content})
        return dspy.History(messages=history_entries)

    @mlflow.trace(span_type=SpanType.AGENT)
    def predict(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> ChatAgentResponse:
        latest_question = messages[-1].content
        response = self.agent(question=latest_question, history=self.prepare_message_history(messages)).answer
        return ChatAgentResponse(
            messages=[ChatAgentMessage(role="assistant", content=response, id=uuid.uuid4().hex)]
        )

# Set model for logging or interactive testing
from mlflow.models import set_model
AGENT = DSPyChatAgent()
set_model(AGENT)
Overwriting agent.py

Test the agent

Interact with the agent to test its output.

Since you manually traced methods within ChatAgent, you can view the trace for each step the agent takes, with any LLM calls made via DSPy APIs automatically traced by autologging.

Replace this placeholder input with an appropriate domain-specific example for your agent.

6
dbutils.library.restartPython()
7
from agent import AGENT

AGENT.predict({"messages": [{"role": "user", "content": "What is 5+5?"}, {"role": "assistant", "content": "5+5=10"}, {"role":"user", "content": "What is the square root of that?"}]})
ChatAgentResponse(messages=[ChatAgentMessage(role='assistant', content='The square root of 10 is approximately 3.16227766.', name=None, id='4414239976264032b5d1aaa6dee74cb0', tool_calls=None, tool_call_id=None, attachments=None)], finish_reason=None, custom_outputs=None, usage=None)
Trace(request_id=tr-ef2fe52762bb41dbb7fcabde591bee4f)

Log the agent as an MLflow model

Log the agent as code from the agent.py file. See MLflow - Models from Code.

9
import mlflow
from agent import LLM_ENDPOINT_NAME
from mlflow.models.resources import DatabricksServingEndpoint

with mlflow.start_run():
    logged_agent_info = mlflow.pyfunc.log_model(
        artifact_path="agent",
        python_model="agent.py",
        pip_requirements=[
            "mlflow",
            "dspy",
            "databricks-sdk",
        ],
        resources=[DatabricksServingEndpoint(endpoint_name=LLM_ENDPOINT_NAME)],
    )
2025/03/27 23:21:02 INFO mlflow.pyfunc: Predicting on input example to validate output

Pre-deployment agent validation

Before registering and deploying the agent, perform pre-deployment checks using the mlflow.models.predict() API. See Databricks documentation (AWS | Azure).

11
mlflow.models.predict(
    model_uri=f"runs:/{logged_agent_info.run_id}/agent",
    input_data={"messages": [{"role": "user", "content": "Hello!"}]},
    env_manager="uv",
)
2025/03/27 23:21:08 INFO mlflow.models.flavor_backend_registry: Selected backend for flavor 'python_function'
2025/03/27 23:21:10 INFO mlflow.utils.virtualenv: Creating a new environment in /tmp/virtualenv_envs/mlflow-d0754190d85c75df3d4fb1a330acb20139572181 with python version 3.11.10 using uv Using CPython 3.11.10 interpreter at: /usr/bin/python3.11 Creating virtual environment at: /tmp/virtualenv_envs/mlflow-d0754190d85c75df3d4fb1a330acb20139572181 Activate with: source /tmp/virtualenv_envs/mlflow-d0754190d85c75df3d4fb1a330acb20139572181/bin/activate 2025/03/27 23:21:11 INFO mlflow.utils.virtualenv: Installing dependencies Using Python 3.11.10 environment at: /tmp/virtualenv_envs/mlflow-d0754190d85c75df3d4fb1a330acb20139572181 Resolved 3 packages in 132ms Downloading setuptools (1.2MiB) Downloading pip (1.7MiB) Downloaded setuptools Downloaded pip Prepared 3 packages in 123ms Installed 3 packages in 26ms + pip==24.2 + setuptools==75.1.0 + wheel==0.38.4 Using Python 3.11.10 environment at: /tmp/virtualenv_envs/mlflow-d0754190d85c75df3d4fb1a330acb20139572181 Resolved 118 packages in 1.24s Downloading tiktoken (1.1MiB) Downloading pillow (4.3MiB) Downloading pyarrow (40.1MiB) Downloading mlflow (26.9MiB) Downloading sqlalchemy (3.1MiB) Downloading aiohttp (1.6MiB) Downloading pydantic-core (1.9MiB) Downloading kiwisolver (1.4MiB) Downloading scikit-learn (12.9MiB) Downloading numpy (18.6MiB) Downloading pandas (12.5MiB) Downloading pygments (1.2MiB) Downloading matplotlib (8.2MiB) Downloading litellm (6.6MiB) Downloading scipy (35.9MiB) Downloading fonttools (4.7MiB) Downloading tokenizers (2.9MiB) Downloading mlflow-skinny (5.9MiB) Downloaded kiwisolver Downloaded tiktoken Downloaded aiohttp Downloaded pydantic-core Downloaded tokenizers Downloaded pygments Downloaded sqlalchemy Downloaded pillow Downloaded fonttools Downloaded matplotlib Downloaded litellm Downloaded mlflow-skinny Downloaded scikit-learn Downloaded pandas Downloaded mlflow Downloaded numpy Downloaded scipy Downloaded pyarrow Prepared 118 packages in 5.56s Installed 118 packages in 344ms + aiohappyeyeballs==2.6.1 + aiohttp==3.11.14 + aiosignal==1.3.2 + alembic==1.15.1 + annotated-types==0.7.0 + anyio==4.9.0 + asyncer==0.0.8 + attrs==25.3.0 + backoff==2.2.1 + blinker==1.9.0 + cachetools==5.5.2 + certifi==2025.1.31 + charset-normalizer==3.4.1 + click==8.1.8 + cloudpickle==3.1.1 + colorlog==6.9.0 + contourpy==1.3.1 + cycler==0.12.1 + databricks-sdk==0.48.0 + datasets==3.5.0 + deprecated==1.2.18 + dill==0.3.8 + diskcache==5.6.3 + distro==1.9.0 + docker==7.1.0 + dspy==2.6.15 + fastapi==0.115.12 + filelock==3.18.0 + flask==3.1.0 + fonttools==4.56.0 + frozenlist==1.5.0 + fsspec==2024.12.0 + gitdb==4.0.12 + gitpython==3.1.44 + google-auth==2.38.0 + graphene==3.4.3 + graphql-core==3.2.6 + graphql-relay==3.2.0 + greenlet==3.1.1 + gunicorn==23.0.0 + h11==0.14.0 + httpcore==1.0.7 + httpx==0.28.1 + huggingface-hub==0.30.0rc2 + idna==3.10 + importlib-metadata==8.6.1 + itsdangerous==2.2.0 + jinja2==3.1.6 + jiter==0.9.0 + joblib==1.4.2 + json-repair==0.40.0 + jsonschema==4.23.0 + jsonschema-specifications==2024.10.1 + kiwisolver==1.4.8 + litellm==1.63.7 + magicattr==0.1.6 + mako==1.3.9 + markdown==3.7 + markdown-it-py==3.0.0 + markupsafe==3.0.2 + matplotlib==3.10.1 + mdurl==0.1.2 + mlflow==2.21.2 + mlflow-skinny==2.21.2 + multidict==6.2.0 + multiprocess==0.70.16 + numpy==2.0.2 + openai==1.61.0 + opentelemetry-api==1.31.1 + opentelemetry-sdk==1.31.1 + opentelemetry-semantic-conventions==0.52b1 + optuna==4.2.1 + packaging==24.2 + pandas==2.2.3 + pillow==11.1.0 + propcache==0.3.1 + protobuf==5.29.4 + pyarrow==19.0.1 + pyasn1==0.6.1 + pyasn1-modules==0.4.1 + pydantic==2.11.0 + pydantic-core==2.33.0 + pygments==2.19.1 + pyparsing==3.2.3 + python-dateutil==2.9.0.post0 + python-dotenv==1.1.0 + pytz==2025.2 + pyyaml==6.0.2 + referencing==0.36.2 + regex==2024.11.6 + requests==2.32.3 + rich==13.9.4 + rpds-py==0.24.0 + rsa==4.9 + scikit-learn==1.6.1 + scipy==1.15.2 + six==1.17.0 + smmap==5.0.2 + sniffio==1.3.1 + sqlalchemy==2.0.40 + sqlparse==0.5.3 + starlette==0.46.1 + tenacity==9.0.0 + threadpoolctl==3.6.0 + tiktoken==0.9.0 + tokenizers==0.21.1 + tqdm==4.67.1 + typing-extensions==4.13.0 + typing-inspection==0.4.0 + tzdata==2025.2 + ujson==5.10.0 + urllib3==2.3.0 + uvicorn==0.34.0 + werkzeug==3.1.3 + wrapt==1.17.2 + xxhash==3.5.0 + yarl==1.18.3 + zipp==3.21.0 2025/03/27 23:21:18 INFO mlflow.utils.environment: === Running command '['bash', '-c', 'source /tmp/virtualenv_envs/mlflow-d0754190d85c75df3d4fb1a330acb20139572181/bin/activate && python -c ""']' 2025/03/27 23:21:18 INFO mlflow.utils.environment: === Running command '['bash', '-c', 'source /tmp/virtualenv_envs/mlflow-d0754190d85c75df3d4fb1a330acb20139572181/bin/activate && python /local_disk0/.ephemeral_nfs/envs/pythonEnv-855edb4f-2701-4b22-998f-f794b5693434/lib/python3.11/site-packages/mlflow/pyfunc/_mlflow_pyfunc_backend_predict.py --model-uri file:///local_disk0/user_tmp_data/spark-855edb4f-2701-4b22-998f-f7/tmp61usz_an/agent --content-type json --input-path /local_disk0/user_tmp_data/spark-855edb4f-2701-4b22-998f-f7/tmpa1jyeq4n/input.json']' {"messages": [{"role": "assistant", "content": "Hello! How can I assist you today?", "id": "c26d4b09eb474d86985b6de6a207980e"}]}

Register the model to Unity Catalog

Before you deploy the agent, you must register the agent to Unity Catalog.

  • TODO Update the catalog, schema, and model_name below to register the MLflow model to Unity Catalog.
13
mlflow.set_registry_uri("databricks-uc")

# TODO: define the catalog, schema, and model name for your UC model.
catalog = ""
schema = ""
model_name = ""
UC_MODEL_NAME = f"{catalog}.{schema}.{model_name}"

# register the model to UC
uc_registered_model_info = mlflow.register_model(model_uri=logged_agent_info.model_uri, name=UC_MODEL_NAME)
Registered model 'sidmurching_catalog.default.dspy_agent' already exists. Creating a new version of this model...
Created version '3' of model 'sidmurching_catalog.default.dspy_agent'.

Deploy the agent

15
from databricks import agents

agents.deploy(UC_MODEL_NAME, uc_registered_model_info.version, tags={"endpointSource": "docs"})
2025/03/27 23:21:38 INFO mlflow.store.artifact.artifact_repo: Failed to complete request, possibly due to credential expiration (Error: An error occurred (404) when calling the HeadObject operation: Not Found). Refreshing credentials and trying again...
2025/03/27 23:21:39 INFO mlflow.store.artifact.artifact_repo: Failed to complete request, possibly due to credential expiration (Error: An error occurred (404) when calling the HeadObject operation: Not Found). Refreshing credentials and trying again... Deployment of sidmurching_catalog.default.dspy_agent version 3 initiated. This can take up to 15 minutes and the Review App & Query Endpoint will not work until this deployment finishes. View status: https://db-ml-models-dev-us-west.cloud.databricks.com/ml/endpoints/agents_sidmurching_catalog-default-dspy_agent Review App: https://db-ml-models-dev-us-west.cloud.databricks.com/ml/review-v2/0b4fe73969fb43b79643ecc00b644d78/chat

Next steps

After your agent is deployed, you can chat with it in AI playground to perform additional checks, share it with SMEs in your organization for feedback, or embed it in a production application. See Databricks documentation (AWS | Azure).

;