databricks_agents_langchain_deploy_unity_catalog_mlflow.py

1import mlflow
2from databricks import agents
3from langchain.agents import initialize_agent, Tool
4from langchain.chat_models import ChatDatabricks
5
6# 1. Define a tool for the agent to use
7def get_weather(location):
8    """Returns the weather for a given location."""
9    return f"The weather in {location} is sunny."
10
11tools = [
12    Tool(
13        name="GetWeather",
14        func=get_weather,
15        description="useful for when you need to answer questions about weather"
16    )
17]
18
19# 2. Initialize the LLM and Agent
20llm = ChatDatabricks(endpoint="databricks-meta-llama-3-1-70b-instruct")
21agent = initialize_agent(tools, llm, agent="zero-shot-react-description", verbose=True)
22
23# 3. Log the agent to MLflow using the Agent Framework
24# Note: In a notebook, this logs the model to the current experiment
25with mlflow.start_run():
26    logged_agent = mlflow.langchain.log_model(
27        lc_model=agent,
28        artifact_path="agent",
29        input_example={"input": "What is the weather in San Francisco?"}
30    )
31
32# 4. Deploy the agent to a Unity Catalog model version
33# Replace with your desired catalog, schema, and model name
34model_name = "main.default.weather_agent"
35model_info = mlflow.register_model(model_uri=logged_agent.model_uri, name=model_name)
36
37# 5. Use the Databricks Agents SDK to deploy the registered model
38agents.deploy(model_name, model_info.version)
39
40print(f"Agent successfully deployed to: {model_name}")