企业官网建设流程全解析

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/mlflow

MLflow AI Gateway 是 MLflow 提供的一个统一模型服务网关，本指南以 examples/gateway/uc_functions/README.md 为骨架，完整演示如何在 MLflow AI Gateway 中集成 Databricks Unity Catalog（UC）函数：通过mlflow gateway start启动网关后，客户端即可用 OpenAI 兼容的tools协议声明调用 UC 函数，由网关在服务端自动将模型产生的函数调用翻译为 SQL 语句、经 Databricks SQL Warehouse 执行并把结果回传给模型。读完本文，你将掌握 UC 函数接入的完整环境配置、示例脚本运行方式，以及网关底层"函数元数据转 JSON Schema → 工具调用 → 参数化 SQL 执行"的完整调用链。

背景：为什么让 AI Gateway 调用 Unity Catalog 函数

MLflow AI Gateway 作为统一网关，把多个 LLM Provider（OpenAI、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 FUNCTION（SQL 与 Python）的说明。注意函数名将被示例脚本通过--uc-function-name参数引用，本文统一以my.uc_func.add为例。

3. 创建 SQL Warehouse

UC 函数的实际执行依赖 Databricks SQL Warehouse。请先在 Databricks 控制台创建一个 SQL Warehouse，并记下它的 HTTP Path——其末尾的 warehouse ID（形如/sql/1.0/warehouses/1234567890123456中最后的数字段）后续需要填入环境变量DATABRICKS_WAREHOUSE_ID。

环境变量与网关启动

按 README 的指引，先导出四组环境变量，再启动网关：

# 1) 认证 Databricks：HOST 形如 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_FUNCTIONS=true # 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 中被读取，若缺失会直接抛出AIGatewayException（HTTP 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/chat，Provider 为openai，模型gpt-4o-mini，并设置了每分钟 10 次调用的限流）。示例脚本请求中的model="chat"即对应此端点名。也可参照仓库 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.add

run.py 用 OpenAI SDK 指向网关地址http://localhost:7000/v1，依次演示两种场景：

场景一：单独调用 UC 函数

构造 OpenAI 兼容的tools参数，工具类型为uc_function：

uc_function = { "type": "uc_function", "uc_function": { "name": args.uc_function_name, }, } resp = client.chat.completions.create( model="chat", 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。用户侧只需声明函数全名，无需关心参数 schema，schema 由网关从 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( model="chat", messages=[msg], tools=[user_defined_function, uc_function], )

这里同时传入标准 OpenAI 函数multiply（在客户端本地执行）和 UC 函数add（在 Databricks 服务端执行）。模型会对同一个多问题消息做出多次工具调用；脚本随后按 OpenAI 工具调用的标准多轮协议，把multiply的调用结果以role: "tool"消息回填：

resp = client.chat.completions.create( model="chat", 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 类型。转换是有损的（不需要转回），映射表包括：

interval类型与未知类型会抛出TypeError。get_func_schema（uc_function_utils.py）进一步把参数名、注释（COMMENT）及默认值（parameter_default，形如(default: xxx)追加到描述中）组装进parameters.properties，并把无默认值的参数列入required——这正是"用户只传函数名、由网关自动补齐参数 schema"的实现基础。

2. 工具名截断与映射

OpenAI 对函数名有 64 字符上限，因此网关在_get_tool_name（uc_function_utils.py 对应 mlflow/gateway/uc_function_utils.py）中把 UC 函数全名拼接为catalog__schema__name并截断到 64 字符，同时维护uc_func_mapping把截断名映射回原始FunctionInfo，以便执行时还原真实函数调用。

3. 生成参数化 SQL 并执行

get_execute_function_sql_stmt（uc_function_utils.py）根据函数是否返回标量生成不同语句：

• 标量函数：SELECTcatalog.schema.func(:p1, :p2, ...)
• 表值函数：SELECT * FROMcatalog.schema.func(:p1, ...)

关键安全与正确性设计：

• _quote_identifier（uc_function_utils.py）对多段标识符逐段用反引号包裹，并拒绝包含内嵌反引号的非法标识符，从源头防止 SQL 注入；
• 所有参数一律使用StatementParameterListItem绑定（:占位符），不做字符串拼接；复杂类型（ARRAY/MAP/STRUCT）通过from_json(:param, 'type_text')还原，BINARY通过unbase64还原；
• 支持命名参数（name => :value）以应对"前面参数有默认值、后面参数被显式提供"的场景。

执行环节execute_function（uc_function_utils.py）通过ws.statement_execution.execute_statement提交到 SQL Warehouse，内置wait_timeout="30s"、row_limit=100、byte_limit=4096（源码注释标注这些限制暂不可配置），并根据结果状态区分成功与错误：标量函数把首行首列转成字符串（format: SCALAR），表值函数用 pandas 转成 CSV（format: CSV），统一封装为FunctionExecutionResult.to_json()供 LLM 消费。

4. 多轮工具调用编排

UC 函数与本地函数混用时，网关在_chat_uc_function（mlflow/gateway/providers/openai.py）中进入一个最多 20 轮的工具调用循环：

1. 把uc_function工具翻译成标准function工具并转发给上游；
2. 若返回中有tool_calls，逐条判断：属于 UC 函数的执行execute_function并记录为uc_func_calls；属于用户自定义函数的记录为user_tool_calls；
3. 若存在本地函数调用，则把 UC 调用结果（join_uc_functions生成的<uc_function_call>...</uc_function_call><uc_function_result>...</uc_function_result>文本块）拼接进回复内容并连同user_tool_calls返回给客户端，由客户端本地执行并回传；
4. 客户端回传role: "tool"消息后，网关通过parse_uc_functions正则（uc_function_utils.py）解析消息内容中的 UC 调用/结果块，与客户端工具调用合并后再次请求模型，直到得到无工具调用的最终回答。

整个过程的 token 用量由TokenUsageAccumulator（uc_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_FUNCTIONS=true；
• DATABRICKS_WAREHOUSE_ID、DATABRICKS_HOST、DATABRICKS_TOKEN三个环境变量缺一不可，否则网关直接报错；
• SQL 语句执行超时（30s）、行数（100）与字节（4096）上限当前为硬编码，暂不支持配置；
• interval数据类型、非 STRING 键的MAP类型不被支持，复杂类型参数通过from_json还原，需保证类型文本正确；
• 工具循环最多 20 轮，超过会返回Max iterations reached（HTTP 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