{ "cells": [ { "cell_type": "markdown", "id": "befb4f50-828c-487a-8922-7a57991fb9d6", "metadata": {}, "source": [ "# Build an agent using UC tools and PythonModel\n", "This notebook will guide you through how to utilize MLflow PythonModel with type hints and Unitycatalog functions as tools to build an agent.\n", "\n", "## Prerequisite\n", "\n", "Install required packages:\n", "```\n", "pip install mlflow==2.20.0 'unitycatalog-langchain[databricks]==0.1.1' langchain_openai==0.3.7\n", "```\n", "\n", "Follow the [instruction](https://docs.databricks.com/aws/en/dev-tools/cli/authentication#authentication-for-the-databricks-cli) to authenticate to your Databricks workspace. Alternatively, check the [UnityCatalog client guidance](https://docs.unitycatalog.io/ai/client/#using-the-client-for-agent-tool-calling) on how to use UnityCatalog Server." ] }, { "cell_type": "code", "execution_count": null, "id": "9062290c-f511-4094-ae8b-7a779e24a712", "metadata": {}, "outputs": [], "source": [ "import mlflow\n", "\n", "# start the mlflow server with `mlflow server` first, then set the tracking uri\n", "mlflow.set_tracking_uri(\"http://127.0.0.1:5000\")" ] }, { "cell_type": "code", "execution_count": null, "id": "f506b2e3-0de0-4663-a634-5042a26c4cca", "metadata": {}, "outputs": [], "source": [ "mlflow.langchain.autolog()" ] }, { "cell_type": "markdown", "id": "936eaccf-045f-4b83-9f3c-76e351d5e731", "metadata": {}, "source": [ "## Create a model" ] }, { "cell_type": "markdown", "id": "d2dc19af-33d9-431c-8b5f-243d7a908008", "metadata": {}, "source": [ "### Create tools" ] }, { "cell_type": "code", "execution_count": null, "id": "89ad5db3-b895-40a3-b382-fcc28c21a4c2", "metadata": {}, "outputs": [], "source": [ "from unitycatalog.ai.core.base import set_uc_function_client\n", "from unitycatalog.ai.core.databricks import DatabricksFunctionClient\n", "\n", "client = DatabricksFunctionClient()\n", "\n", "# sets the default uc function client\n", "set_uc_function_client(client)" ] }, { "cell_type": "code", "execution_count": null, "id": "2ec578f6-8c23-477d-93b0-d44abf85e4ff", "metadata": {}, "outputs": [], "source": [ "CATALOG = \"ml\"\n", "SCHEMA = \"serena_test\"" ] }, { "cell_type": "code", "execution_count": null, "id": "720a4c2e-5634-4e28-b64f-dfef2eabe9de", "metadata": {}, "outputs": [], "source": [ "# Define a python code execution function\n", "def execute_python_code(code: str) -> str:\n", " \"\"\"\n", " Executes the given python code and returns its stdout.\n", " Remember the code should print the final result to stdout.\n", "\n", " Args:\n", " code: Python code to execute. Remember to print the final result to stdout.\n", " \"\"\"\n", " # clint comment is used to disable lint check, you could delete them\n", " import sys # clint: disable=lazy-import\n", " from io import StringIO # clint: disable=lazy-import\n", "\n", " stdout = StringIO()\n", " sys.stdout = stdout\n", " exec(code) # noqa: S102\n", " return stdout.getvalue()\n", "\n", "\n", "function_info = client.create_python_function(\n", " func=execute_python_code, catalog=CATALOG, schema=SCHEMA, replace=True\n", ")\n", "python_execution_function_name = function_info.full_name" ] }, { "cell_type": "code", "execution_count": null, "id": "68fa8619-3f4c-410c-bc01-46f064f1c767", "metadata": { "scrolled": true }, "outputs": [ { "data": { "text/plain": [ "FunctionInfo(browse_only=None, catalog_name='ml', comment='translate the content to target language, currently only english <-> spanish translation is supported', created_at=1741264002930, created_by='serena.ruan@databricks.com', data_type=, external_language=None, external_name=None, full_data_type='STRING', full_name='ml.serena_test.translate', function_id='015ff53c-cdf9-453b-8e19-9a2332ca7b9f', input_params=FunctionParameterInfos(parameters=[FunctionParameterInfo(name='content', type_text='string', type_name=, position=0, comment='content to translate', parameter_default=None, parameter_mode=None, parameter_type=, type_interval_type=None, type_json='{\"name\":\"content\",\"type\":\"string\",\"nullable\":true,\"metadata\":{\"comment\":\"content to translate\"}}', type_precision=0, type_scale=0), FunctionParameterInfo(name='language', type_text='string', type_name=, position=1, comment='target language', parameter_default=None, parameter_mode=None, parameter_type=, type_interval_type=None, type_json='{\"name\":\"language\",\"type\":\"string\",\"nullable\":true,\"metadata\":{\"comment\":\"target language\"}}', type_precision=0, type_scale=0)]), is_deterministic=False, is_null_call=None, metastore_id='19a85dee-54bc-43a2-87ab-023d0ec16013', name='translate', owner='serena.ruan@databricks.com', parameter_style=, properties='{\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.taskWaitTimeInSeconds\":\"1000\",\"sqlConfig.spark.sql.ansi.enabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.decimal.dataType.enabled\":\"true\",\"referredTempFunctionsNames\":\"[]\",\"sqlConfig.spark.sql.streaming.statefulOperator.stateRebalancing.enabled\":\"false\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.dynamicPoolSizeEnabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.debugLogEnabled\":\"true\",\"catalogAndNamespace.part.0\":\"main\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.remoteHttpClient.maxConnections\":\"2048\",\"sqlConfig.spark.sql.legacy.createHiveTableByDefault\":\"false\",\"sqlConfig.spark.sql.shuffleDependency.skipMigration.enabled\":\"true\",\"sqlConfig.spark.sql.streaming.stopTimeout\":\"15s\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.remoteHttpClient.timeoutInSeconds\":\"360\",\"referredTempViewNames\":\"[]\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.enabled\":\"true\",\"sqlConfig.spark.sql.legacy.codingErrorAction\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.scaleUpThresholdCurrentQpsIncreaseRatio\":\"0.0\",\"sqlConfig.spark.databricks.sql.functions.builtInHttpReq.enabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.scaleUpThresholdTotalQpsIncreaseRatio\":\"0.0\",\"catalogAndNamespace.part.1\":\"default\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.useDedicatedHttpClient\":\"true\",\"sqlConfig.spark.sql.readSideCharPadding\":\"true\",\"sqlConfig.spark.sql.defaultStorage.zstdCompressionByDefault.enabled\":\"true\",\"sqlConfig.spark.sql.variable.substitute\":\"false\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.threadKeepAliveTimeInSeconds\":\"600\",\"sqlConfig.spark.databricks.sql.expression.aiFunctions.repartition\":\"0\",\"sqlConfig.spark.databricks.sql.functions.aiForecast.enabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.useDynamicTaskQueueExecutor\":\"false\",\"sqlConfig.spark.sql.sources.default\":\"delta\",\"sqlConfig.spark.sql.hive.convertCTAS\":\"true\",\"sqlConfig.spark.sql.legacy.jdbc.doNotUseCalendar\":\"false\",\"sqlConfig.spark.databricks.sql.functions.vectorSearch.enabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.safe.inference.enabled\":\"true\",\"sqlConfig.spark.sql.scripting.enabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.batch.aiQuery.embedding.request.size\":\"4\",\"referredTempVariableNames\":\"[]\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.batch.execution.size\":\"2048\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.clusterSizeBasedGlobalParallelism.scaleFactor\":\"32.0\",\"sqlConfig.spark.databricks.sql.functions.aiGen.endpointName\":\"databricks-meta-llama-3-3-70b-instruct\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.embeddingsEndpointName\":\"databricks-gte-large-en\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.batchSession.enabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.maxPoolSize\":\"2048\",\"sqlConfig.spark.sql.sources.commitProtocolClass\":\"com.databricks.sql.transaction.directory.DirectoryAtomicCommitProtocol\",\"sqlConfig.spark.sql.functions.remoteHttpClient.retryOn400TimeoutError\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.purposeBuiltFunctions.batch.execution.enabled\":\"true\",\"catalogAndNamespace.numParts\":\"2\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.scaleUpThresholdSuccessRatio\":\"0.95\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.modelEndpointTypeParsing.enabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.batchInferenceApi.enabled\":\"true\",\"sqlConfig.spark.sql.stableDerivedColumnAlias.enabled\":\"true\",\"sqlConfig.spark.sql.parquet.compression.codec\":\"snappy\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.model.parameters.enabled\":\"true\",\"sqlConfig.spark.sql.streaming.stateStore.providerClass\":\"com.databricks.sql.streaming.state.RocksDBStateStoreProvider\",\"sqlConfig.spark.sql.disableErrorOnCredentialMismatch\":\"true\"}', return_params=None, routine_body=, routine_definition='(SELECT ai_translate(content, language))', routine_dependencies=None, schema_name='serena_test', security_type=, specific_name='translate', sql_data_access=, sql_path=None, updated_at=1741264002930, updated_by='serena.ruan@databricks.com')" ] }, "execution_count": null, "metadata": {}, "output_type": "execute_result" } ], "source": [ "translate_function_name = f\"{CATALOG}.{SCHEMA}.translate\"\n", "sql_body = f\"\"\"CREATE OR REPLACE FUNCTION {translate_function_name}(content STRING COMMENT 'content to translate', language STRING COMMENT 'target language')\n", "RETURNS STRING\n", "COMMENT 'translate the content to target language, currently only english <-> spanish translation is supported'\n", "RETURN SELECT ai_translate(content, language)\n", "\"\"\"\n", "client.create_function(sql_function_body=sql_body)" ] }, { "cell_type": "code", "execution_count": null, "id": "648d5ecb-cfee-46f0-933d-29e914f973be", "metadata": {}, "outputs": [], "source": [ "from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit\n", "\n", "toolkit = UCFunctionToolkit(\n", " function_names=[\n", " python_execution_function_name,\n", " translate_function_name,\n", " ]\n", ")\n", "tools = toolkit.tools" ] }, { "cell_type": "code", "execution_count": null, "id": "6e38f029-d54f-4736-a32a-f8a4417bd64a", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[UnityCatalogTool(name='ml__serena_test__execute_python_code', description='Executes the given python code and returns its stdout. Remember the code should print the final result to stdout.', args_schema=, func=.func at 0x3152293f0>, uc_function_name='ml.serena_test.execute_python_code', client_config={'warehouse_id': None, 'profile': None}),\n", " UnityCatalogTool(name='ml__serena_test__translate', description='translate the content to target language, currently only english <-> spanish translation is supported', args_schema=, func=.func at 0x315f21a20>, uc_function_name='ml.serena_test.translate', client_config={'warehouse_id': None, 'profile': None})]" ] }, "execution_count": null, "metadata": {}, "output_type": "execute_result" } ], "source": [ "tools" ] }, { "cell_type": "markdown", "id": "20b796d3-b098-48c9-b62d-f725acfad29f", "metadata": {}, "source": [ "### Define langgraph model" ] }, { "cell_type": "code", "execution_count": null, "id": "df803a03-d0c0-46cc-af71-8b5330562c9e", "metadata": {}, "outputs": [], "source": [ "from typing import Literal\n", "\n", "from langchain_openai.chat_models import ChatOpenAI\n", "from langgraph.graph import END, START, MessagesState, StateGraph\n", "from langgraph.prebuilt import ToolNode\n", "\n", "tool_node = ToolNode(tools)\n", "model = ChatOpenAI(model=\"gpt-4o-mini\").bind_tools(tools)\n", "\n", "\n", "# Define the function that determines whether to continue or not\n", "def should_continue(state: MessagesState) -> Literal[\"tools\", END]:\n", " messages = state[\"messages\"]\n", " last_message = messages[-1]\n", " # If the LLM makes a tool call, then we route to the \"tools\" node\n", " if last_message.tool_calls:\n", " return \"tools\"\n", " # Otherwise, we stop (reply to the user)\n", " return END\n", "\n", "\n", "# Define the function that calls the model\n", "def call_model(state: MessagesState):\n", " messages = state[\"messages\"]\n", " response = model.invoke(messages)\n", " # We return a list, because this will get added to the existing list\n", " return {\"messages\": [response]}\n", "\n", "\n", "# Define a new graph\n", "workflow = StateGraph(MessagesState)\n", "\n", "# Define the two nodes we will cycle between\n", "workflow.add_node(\"agent\", call_model)\n", "workflow.add_node(\"tools\", tool_node)\n", "\n", "# Set the entrypoint as `agent`\n", "# This means that this node is the first one called\n", "workflow.add_edge(START, \"agent\")\n", "\n", "# We now add a conditional edge\n", "workflow.add_conditional_edges(\n", " # First, we define the start node. We use `agent`.\n", " # This means these are the edges taken after the `agent` node is called.\n", " \"agent\",\n", " # Next, we pass in the function that will determine which node is called next.\n", " should_continue,\n", ")\n", "\n", "# We now add a normal edge from `tools` to `agent`.\n", "# This means that after `tools` is called, `agent` node is called next.\n", "workflow.add_edge(\"tools\", \"agent\")\n", "\n", "app = workflow.compile()" ] }, { "cell_type": "markdown", "id": "538b0311-dd95-4c8d-bb82-d35ff41aeb4f", "metadata": {}, "source": [ "### Test the model first" ] }, { "cell_type": "code", "execution_count": null, "id": "99f3cc51-5aa0-4c52-b38d-67f7ef5951d8", "metadata": {}, "outputs": [], "source": [ "final_state = app.invoke(\n", " {\n", " \"messages\": [\n", " {\n", " \"role\": \"user\",\n", " \"content\": \"What is MLflow? Keep the response concise and reply in Chinese. Try using as many tools as possible\",\n", " }\n", " ]\n", " },\n", ")\n", "response = final_state[\"messages\"][-1].content" ] }, { "cell_type": "markdown", "id": "ad1ceb8d-8654-4a31-b3b6-72ca7358a939", "metadata": {}, "source": [ "## Save the model as python file --> agent.py" ] }, { "cell_type": "markdown", "id": "d3ee60f5-fe89-4371-9747-f2c7ab8697e1", "metadata": {}, "source": [ "## Log the model" ] }, { "cell_type": "code", "execution_count": null, "id": "37a9b50c-43a8-4799-9f9e-32ed93a4de98", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "2025/03/06 21:08:14 INFO mlflow.models.signature: Inferring model signature from type hints\n", "2025/03/06 21:08:14 INFO mlflow.models.signature: Running the predict function to generate output based on input example\n", "2025/03/06 21:08:52 INFO mlflow.models.model: Found the following environment variables used during model inference: [OPENAI_API_KEY]. Please check if you need to set them when deploying the model. To disable this message, set environment variable `MLFLOW_RECORD_ENV_VARS_IN_MODEL_LOGGING` to `false`.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "๐Ÿƒ View run angry-conch-727 at: http://127.0.0.1:5000/#/experiments/0/runs/30690b953587448597c77ca9b2cfa371\n", "๐Ÿงช View experiment at: http://127.0.0.1:5000/#/experiments/0\n" ] } ], "source": [ "input_example = [{\"messages\": [{\"role\": \"user\", \"content\": \"What is DSPy?\"}]}]\n", "\n", "with mlflow.start_run():\n", " model_info = mlflow.pyfunc.log_model(\n", " name=\"model\",\n", " # Pass the path to the saved model file\n", " python_model=\"agent.py\",\n", " input_example=input_example,\n", " )" ] }, { "cell_type": "markdown", "id": "23c1a000-bc49-4831-bb12-5e58c026eee0", "metadata": {}, "source": [ "## Load model as pyfunc and predict" ] }, { "cell_type": "code", "execution_count": null, "id": "c0db5686-6e88-41c7-9794-434afd0f8a9f", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "'DSPy is a framework for building and deploying decision-making systems in Python. It is used primarily in the context of machine learning and data science to create systems that can make predictions or decisions based on input data. DSPy facilitates the definition of decision rules, decision trees, and other logical constructs, allowing data scientists and developers to implement complex decision-making logic in a structured and manageable way.\\n\\nKey features of DSPy may include:\\n\\n1. **Declarative Syntax**: Provides a way to define rules and conditions in a readable and maintainable manner.\\n2. **Integration with Machine Learning**: Supports the incorporation of predictive models to enhance decision-making processes.\\n3. **Scalability**: Built to handle large datasets and complex decision logic without significant performance bottlenecks.\\n4. **Modularity**: Allows for reusable components and functions, making it easier to manage and update decision systems.\\n\\nOverall, DSPy aims to simplify the process of creating and managing decision-making systems, making it more accessible for data practitioners.'" ] }, "execution_count": null, "metadata": {}, "output_type": "execute_result" } ], "source": [ "pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)\n", "pyfunc_model.predict(input_example)" ] }, { "cell_type": "code", "execution_count": null, "id": "ec82ca5b-59ce-437c-bbe6-11aab8ebc6d2", "metadata": { "scrolled": true }, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "3a821cf96f9e4da7874645dbce12f0c5", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Downloading artifacts: 0%| | 0/9 [00:00