MLflow AI Gateway 接入 Unity Catalog 函数:从配置到源码级原理 📅 发布时间:2026/9/12 1:36:30 👁 浏览次数: MLflow AI Gateway 接入 Unity Catalog 函数从配置到源码级原理【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflowMLflow AI Gateway 是 MLflow 提供的一个统一模型服务网关本指南以 examples/gateway/uc_functions/README.md 为骨架完整演示如何在 MLflow AI Gateway 中集成 Databricks Unity CatalogUC函数通过mlflow gateway start启动网关后客户端即可用 OpenAI 兼容的tools协议声明调用 UC 函数由网关在服务端自动将模型产生的函数调用翻译为 SQL 语句、经 Databricks SQL Warehouse 执行并把结果回传给模型。读完本文你将掌握 UC 函数接入的完整环境配置、示例脚本运行方式以及网关底层函数元数据转 JSON Schema → 工具调用 → 参数化 SQL 执行的完整调用链。背景为什么让 AI Gateway 调用 Unity Catalog 函数MLflow AI Gateway 作为统一网关把多个 LLM ProviderOpenAI、Anthropic 等收敛到 OpenAI 兼容的 REST 接口之下。而 Unity Catalog 是 Databricks 上的统一元数据与数据治理体系用户可以用 SQL/Python 注册业务函数如计算、查询、特征提取。将两者打通后LLM 在对话过程中可以通过 Function Calling 机制动态调用 UC 中的既有函数复用企业级的数据与计算逻辑而无需把敏感凭证暴露给客户端——UC 函数的执行完全发生在网关服务端。从仓库结构看UC 函数集成是 AI Gateway 中 OpenAI Provider 的一个专属能力分支源码位于 mlflow/gateway/providers/openai.py并由独立的工具模块 mlflow/gateway/uc_function_utils.py 支撑。示例代码位于 examples/gateway/uc_functions/包含 README.md 与 run.py 两个文件。前置条件1. 安装依赖包pip install mlflow openai databricks-sdk其中databricks-sdk是执行 UC 函数的关键——网关正是通过它调用 Databricks Workspace 的functions与statement_executionAPI。openai则用于示例客户端脚本。2. 在 Databricks Notebook 中创建 UC 函数在 Databricks Notebook 中执行如下 SQL创建一个最简单的加法函数%sql CREATE OR REPLACE FUNCTION my.uc_func.add ( x INTEGER COMMENT The first number to add., y INTEGER COMMENT The second number to add. ) RETURNS INTEGER LANGUAGE SQL RETURN x y该函数位于my.uc_func模式下接收两个整数并返回它们的和。若需自定义函数如 Python 函数、表值函数可参考 Databricks 官方 SQL 语言手册中关于CREATE FUNCTIONSQL 与 Python的说明。注意函数名将被示例脚本通过--uc-function-name参数引用本文统一以my.uc_func.add为例。3. 创建 SQL WarehouseUC 函数的实际执行依赖 Databricks SQL Warehouse。请先在 Databricks 控制台创建一个 SQL Warehouse并记下它的 HTTP Path——其末尾的 warehouse ID形如/sql/1.0/warehouses/1234567890123456中最后的数字段后续需要填入环境变量DATABRICKS_WAREHOUSE_ID。环境变量与网关启动按 README 的指引先导出四组环境变量再启动网关# 1) 认证 DatabricksHOST 形如 https://my.databricks.com export DATABRICKS_HOST... export DATABRICKS_TOKEN... # 2) 执行 UC 函数所需的 SQL Warehouse ID export DATABRICKS_WAREHOUSE_ID... # 3) 开启 Unity Catalog 集成开关 export MLFLOW_ENABLE_UC_FUNCTIONStrue # 4) 启动 AI Gateway使用 OpenAI 示例配置端口 7000 mlflow gateway start --config-path examples/gateway/openai/config.yaml --port 7000关于这些变量在源码中的位置MLFLOW_ENABLE_UC_FUNCTIONS是一个布尔型开关默认值为False定义于 mlflow/environment_variables.py。只有显式设为true时AI Gateway 才会在处理请求时进入 UC 函数分支。DATABRICKS_WAREHOUSE_ID在 mlflow/gateway/providers/openai.py 中被读取若缺失会直接抛出AIGatewayExceptionHTTP 400提示DATABRICKS_WAREHOUSE_ID environment variable is not set。DATABRICKS_HOST/DATABRICKS_TOKEN由databricks-sdk的WorkspaceClient在 openai.py 中通过_get_workspace_client()隐式读取用于创建工作区客户端。启动时使用的 examples/gateway/openai/config.yaml 定义了名为chat的聊天端点endpoint_type: llm/v1/chatProvider 为openai模型gpt-4o-mini并设置了每分钟 10 次调用的限流。示例脚本请求中的modelchat即对应此端点名。也可参照仓库 examples/gateway/ 下的其他 Provider 配置自行替换。运行示例脚本网关启动后另开终端执行# Replace my.uc_func.add if your UC function has a different name python examples/gateway/uc_functions/run.py --uc-function-name my.uc_func.addrun.py 用 OpenAI SDK 指向网关地址http://localhost:7000/v1依次演示两种场景场景一单独调用 UC 函数构造 OpenAI 兼容的tools参数工具类型为uc_functionuc_function { type: uc_function, uc_function: { name: args.uc_function_name, }, } resp client.chat.completions.create( modelchat, messages[{role: user, content: What is the result of 1 2?}], tools[uc_function], ) print(resp.choices[0].message.content)注意type: uc_function不是标准的 OpenAI 工具类型而是 MLflow AI Gateway 的扩展类型——网关在服务端把它翻译成标准的function工具后再转发给上游 LLM。用户侧只需声明函数全名无需关心参数 schemaschema 由网关从 UC 元数据自动生成详见下文源码解析。场景二UC 函数与用户自定义函数混合使用user_defined_function { type: function, function: { description: Multiply numbers, name: multiply, parameters: { type: object, properties: { x: {type: integer, description: First number}, y: {type: integer, description: Second number}, }, required: [x, y], }, }, } def multiply(x: int, y: int) - int: return x * y msg { role: user, content: What is the result of 1 2? What is the result of 3 4? What is the result of 5 * 6?, } resp client.chat.completions.create( modelchat, messages[msg], tools[user_defined_function, uc_function], )这里同时传入标准 OpenAI 函数multiply在客户端本地执行和 UC 函数add在 Databricks 服务端执行。模型会对同一个多问题消息做出多次工具调用脚本随后按 OpenAI 工具调用的标准多轮协议把multiply的调用结果以role: tool消息回填resp client.chat.completions.create( modelchat, messages[ msg, {role: assistant, content: resp.choices[0].message.content}, {role: assistant, content: , tool_calls: resp.choices[0].message.tool_calls}, { role: tool, tool_call_id: resp.choices[0].message.tool_calls[0].id, content: str(multiply(**json.loads(multiply_call.arguments))), }, ], )脚本最后断言第一个工具调用确实来自multiply验证了UC 函数调用 本地函数调用混排时网关能正确分流。源码级原理网关如何执行 UC 函数1. 从 UC 元数据生成 JSON Schema类型映射当请求包含type uc_function的工具时网关用workspace_client.functions.get(function_name)拉取函数元数据FunctionInfo再调用 mlflow/gateway/uc_function_utils.py 中的uc_type_to_json_schema_type把 UC 数据类型转换为 JSON Schema 类型。转换是有损的不需要转回映射表包括UC 类型JSON Schema 类型long/integer/short/byteintegerdouble/float/decimal*numberstring/binarystringbooleanbooleandatestringformat: datetimestamp/timestamp_ntzstringformat: date-timevoidnullarrayarrayitems递归map仅支持 STRING 键objectadditionalPropertiesstructobjectproperties逐字段递归interval类型与未知类型会抛出TypeError。get_func_schemauc_function_utils.py进一步把参数名、注释COMMENT及默认值parameter_default形如(default: xxx)追加到描述中组装进parameters.properties并把无默认值的参数列入required——这正是用户只传函数名、由网关自动补齐参数 schema的实现基础。2. 工具名截断与映射OpenAI 对函数名有 64 字符上限因此网关在_get_tool_nameuc_function_utils.py 对应 mlflow/gateway/uc_function_utils.py中把 UC 函数全名拼接为catalog__schema__name并截断到 64 字符同时维护uc_func_mapping把截断名映射回原始FunctionInfo以便执行时还原真实函数调用。3. 生成参数化 SQL 并执行get_execute_function_sql_stmtuc_function_utils.py根据函数是否返回标量生成不同语句标量函数SELECTcatalog.schema.func(:p1, :p2, ...)表值函数SELECT * FROMcatalog.schema.func(:p1, ...)关键安全与正确性设计_quote_identifieruc_function_utils.py对多段标识符逐段用反引号包裹并拒绝包含内嵌反引号的非法标识符从源头防止 SQL 注入所有参数一律使用StatementParameterListItem绑定:占位符不做字符串拼接复杂类型ARRAY/MAP/STRUCT通过from_json(:param, type_text)还原BINARY通过unbase64还原支持命名参数name :value以应对前面参数有默认值、后面参数被显式提供的场景。执行环节execute_functionuc_function_utils.py通过ws.statement_execution.execute_statement提交到 SQL Warehouse内置wait_timeout30s、row_limit100、byte_limit4096源码注释标注这些限制暂不可配置并根据结果状态区分成功与错误标量函数把首行首列转成字符串format: SCALAR表值函数用 pandas 转成 CSVformat: CSV统一封装为FunctionExecutionResult.to_json()供 LLM 消费。4. 多轮工具调用编排UC 函数与本地函数混用时网关在_chat_uc_functionmlflow/gateway/providers/openai.py中进入一个最多 20 轮的工具调用循环把uc_function工具翻译成标准function工具并转发给上游若返回中有tool_calls逐条判断属于 UC 函数的执行execute_function并记录为uc_func_calls属于用户自定义函数的记录为user_tool_calls若存在本地函数调用则把 UC 调用结果join_uc_functions生成的uc_function_call.../uc_function_calluc_function_result.../uc_function_result文本块拼接进回复内容并连同user_tool_calls返回给客户端由客户端本地执行并回传客户端回传role: tool消息后网关通过parse_uc_functions正则uc_function_utils.py解析消息内容中的 UC 调用/结果块与客户端工具调用合并后再次请求模型直到得到无工具调用的最终回答。整个过程的 token 用量由TokenUsageAccumulatoruc_function_utils.py累计并写回响应usage字段。5. 测试验证该功能有独立的测试套件 tests/gateway/providers/test_openai_uc_functions.py覆盖了工具 schema 生成、参数化 SQL 语句构建、UC 函数与本地函数混排等核心路径是理解预期行为的另一份权威参考。注意事项与已知限制从源码中可以确认以下限制使用前需知悉UC 函数集成仅在 OpenAI 兼容的 Chat 端点llm/v1/chat生效且需显式设置MLFLOW_ENABLE_UC_FUNCTIONStrueDATABRICKS_WAREHOUSE_ID、DATABRICKS_HOST、DATABRICKS_TOKEN三个环境变量缺一不可否则网关直接报错SQL 语句执行超时30s、行数100与字节4096上限当前为硬编码暂不支持配置interval数据类型、非 STRING 键的MAP类型不被支持复杂类型参数通过from_json还原需保证类型文本正确工具循环最多 20 轮超过会返回Max iterations reachedHTTP 500函数名因 OpenAI 64 字符限制会被截断保留末尾 64 字符极端情况下可能丢失前缀信息建议 UC 函数命名保持精简示例默认使用gpt-4o-mini与 OpenAI 官方 API若改用其他 Provider请同步调整 examples/gateway/openai/config.yaml 中的模型配置与 API Key 环境变量。小结本文以 examples/gateway/uc_functions/README.md 为主线走通了安装依赖 → 创建 UC 函数与 SQL Warehouse → 配置环境变量 → 启动 AI Gateway → 运行 run.py的完整流程并结合 mlflow/gateway/uc_function_utils.py 与 mlflow/gateway/providers/openai.py 的源码剖析了类型映射、参数化 SQL、SQL 注入防护、多轮工具编排等底层机制。这套链路让 MLflow AI Gateway 成为连接 LLM 与企业级 Unity Catalog 数据资产的桥梁是构建模型对话 数据函数即服务应用的实用范本。【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考