How to deploy an agent into model registry using MLFlow

SandipCoder — Mon, 02 Sep 2024 08:39:47 GMT

Goal :

Here below I am attaching my code snippet where I am trying to build an agent on my custom functions .Things are working well when I am running agent.invoke(“Question”).
But now I want to deploy this agent into model registry using MLFlow but I am not able to . I have explored many databricks blogs but all of them ended up with deploying model and chain only.

SampleCode:

agent_name='PolicyStar_Agent'
agent_version = "1.0"

def multiply(a: int, b: int) -> int:
"""Multiply two numbers."""
return a * b

tools = [
StructuredTool.from_function(func=multiply),
StructuredTool.from_function(func=Fn_query_from_unstructured_data),
StructuredTool.from_function(func=Fn_query_from_structured_data),
StructuredTool.from_function(func=Fn_executesql_from_structured_data)
]

agent = initialize_agent(
agent_name=agent_name,
tools=tools,
llm=llm,
agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
verbose=True
)

response = agent.invoke("""
Use your tools to answer this question.

Step 1: Determine the data type. If the question is about user's personal information (like question pattern is 'My name is 'XXXX' or I am 'XXXX'), proceed to Step 2. If it is about general data, proceed to Step 4.
Step 2: Invoke the Fn_query_from_structured_data tool by passing the question mentioned in Step 6. Extract only sql and store the query in a variable named var_q.
Step 3: Invoke the Fn_executesql_from_structured_data tool, passing the SQL query stored in var_q as input. Execute the query and retrieve the result.
Step 4: If the question is about general data, invoke the Fn_query_from_unstructured_data tool to retrieve the answer.
Step 5: If any error occurs during the execution, stop immediately and report the error.
Step 6: Answer the question: My Name is Eric Ruiz. What are my yearly premium amount and sum insured amount?.
Please execute the steps sequentially and provide the final answer.
""")
print(response)

Re: How to deploy an agent into model registry using MLFlow

prar_shah — Mon, 02 Sep 2024 08:49:54 GMT

Had the same question when I was trying to deploy my custom llm. How we did it was wrapping it inside a custom mlflow model and registering it.

Re: How to deploy an agent into model registry using MLFlow

SandipCoder — Mon, 07 Apr 2025 11:05:00 GMT

Please go though this :

https://docs.databricks.com/aws/en/machine-learning/model-serving/custom-models

topic Re: How to deploy an agent into model registry using MLFlow in Generative AI

How to deploy an agent into model registry using MLFlow

Re: How to deploy an agent into model registry using MLFlow

Re: How to deploy an agent into model registry using MLFlow