Agentic workflow with code¶
This notebook demonstrates a simple agentic workflow in DataRobot, showing how MLOps can be used to server, 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 the model-metadata.yaml for the agent deployment
- 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='5a61d7a0fbd723001a2f70d9' # Specify your prediction server here
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='5e8c889607389fe0f466c72d', # 3.9 drop-in
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 deployment ID of the calculator deployment created in step 1. In production you should use the DataRobot credential store to expose secrets in the deployment. DataRobot has omitted this step for brevity of this example.
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='5e8c889607389fe0f466c72d', # 3.9 drop-in,
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