Agentic workflow with code¶
This notebook demonstrates a simple agentic workflow in DataRobot, showing how MLOps can be used to serve, monitor, and govern the workflow.
After running the notebook, you will have DataRobot MLOps deployments for a simple agent that can also reliably perform basic arithmetic. This example also produces a separate deployment for the calculator tool used by the agent.
Code asset overview¶
The following code-based assets are used in this workflow, all of which you can download here.
calculator/custom.py: Custom deployment logic for the "calculator" tool that will allow the LLM to reliably perform simple arithmetic operations.agent/custom.py: Custom deployment logic for the agent that will either respond directly to user prompts or alternatively first delegate to the calculator tool.agent/requirements.txt: Python dependencies for the agent custom deployment.agent/model-metadata.yaml: A configuration file for the agent deployment that specifies Azure OpenAI credentials and the identifier of the calculator deployment.create_deployments.ipynb: This notebook file; includes code for creating and testing the deployments.
Workflow outline¶
- Create the calculator deployment
- Update model-metadata.yaml and requirements.txt
- Create the agent deployment
- Test and make predictions with the agent deployment
1. Create the calculator deployment¶
The following cell deploys the files in the calculator directory (calculator/custom.py). Create a custom model deployment by importing the DataRobot package and using DataRobot MLOps' deployment creation methods. This model functions as a calculator. You can bring two numbers and a mathematic operation to the deployment, and the model will return the answer.
import datarobot as dr
default_prediction_server_id = '<YOUR_PREDICTION_SERVER_ID>' # Specify your prediction server here
execution_environment_id = "5e8c889607389fe0f466c72d"
cm_calc = dr.CustomInferenceModel.create(name='Calculator',
target_name='result',
target_type='TextGeneration')
cmv_calc = dr.CustomModelVersion.create_clean(cm_calc.id,
base_environment_id=execution_environment_id,
folder_path='./calculator')
rmv_calc = dr.RegisteredModelVersion.create_for_custom_model_version(cmv_calc.id)
d_calc = dr.Deployment.create_from_registered_model_version(rmv_calc.id,
'Calculator',
default_prediction_server_id=default_prediction_server_id)
2. Provide credentials¶
In your text editor of choice, update agent/model-metadata.yaml with your Azure Open AI credentials and the calculator deployment ID from step 1 (d_calc.id). In production you should use the DataRobot credential store to expose secrets in the deployment.
Update package versions in agent/requirements.txt to the following:
openai==1.55.3
pydantic==2.5.2
datarobot-predict==1.13.5
datarobot==3.4.0
3. Create the agent deployment¶
Use the code below to create a deployment for the agent.
cm_agent = dr.CustomInferenceModel.create(name='Agent',
target_name='completion',
target_type='TextGeneration')
cmv_agent = dr.CustomModelVersion.create_clean(cm_agent.id,
base_environment_id=execution_environment_id,
folder_path='./agent')
dr.CustomModelVersionDependencyBuild.start_build(cm_agent.id, cmv_agent.id)
rmv_agent = dr.RegisteredModelVersion.create_for_custom_model_version(cmv_agent.id)
d_agent = dr.Deployment.create_from_registered_model_version(rmv_agent.id,
'Agent',
default_prediction_server_id=default_prediction_server_id)
4. Test the deployment¶
The cells below communicate with the deployments by asking the calculator model a math problem (what is 4 x 752) and receiving the answer retrieved by the agent.
from datarobot_predict.deployment import predict
import pandas as pd
import json
messages = [
{'role': 'user',
'content': 'what is 4*752',}
]
df, _ = predict(d_agent, pd.DataFrame([{'messages': json.dumps(messages)}]))
df
messages = [
{'role': 'user',
'content': 'hello',}
]
df, _ = predict(d_agent, pd.DataFrame([{'messages': json.dumps(messages)}]))
df