如何用 MLflow Agent Server 把 AI 代理托管为生产 REST API

如何用 MLflow Agent Server 把 AI 代理托管为生产 REST API

【免费下载链接】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

如果你已经用openai-agents-sdk或其他框架写好了一个 AI 代理，想把它暴露成一个可以被 HTTP 请求调用的服务，而不是只能在脚本里跑，MLflow 的 Agent Server 可以直接完成这件事。它是一个基于 FastAPI 的服务器，把代理函数注册在POST /invocations端点上，并自动处理请求/响应校验、MLflow tracing 记录，以及（可选的）流式响应。

适用前提：

• 需要安装openai-agents-sdk和mlflow>=3.0系列中的 3.6.0 及以上版本（文档给出的安装命令是'mlflow>=3.6.0'）；
• 代理底层调用 OpenAI 时需要设置OPENAI_API_KEY；
• 服务端代码必须导入你的 agent 模块，@invoke注册的函数才会被服务器发现。

安装依赖并准备环境

pip install -U openai-agents 'mlflow>=3.6.0' export OPENAI_API_KEY=sk-...

OPENAI_API_KEY替换为你自己的密钥。

定义带@invoke装饰器的代理函数

创建一个agent.py，在其中定义代理，并把处理请求的函数用@invoke()标注。函数签名要使用ResponsesAgentRequest/ResponsesAgentResponse这两个类型，服务器会据此自动做输入输出校验：

from agents import Agent, Runner from mlflow.genai.agent_server import invoke, stream from mlflow.types.responses import ResponsesAgentRequest, ResponsesAgentResponse agent = Agent( name="Math Tutor", instructions="You provide help with math problems. Explain your reasoning and include examples", ) @invoke() async def non_streaming(request: ResponsesAgentRequest) -> ResponsesAgentResponse: msgs = [i.model_dump() for i in request.input] result = await Runner.run(agent, msgs) return ResponsesAgentResponse(output=[item.to_input_item() for item in result.new_items]) # 可选：再注册一个 @stream 函数即可支持流式响应

两点限制需要注意：

• @invoke装饰器只能使用一次，重复注册会抛出ValueError；
• 只有注册了@stream函数后，请求体里的"stream": true才会生效，否则流式请求没有处理函数可用。

编写服务器入口start_server.py

# 必须导入 agent 模块，@invoke 注册的函数才能被服务器发现 import agent # noqa: F401 from mlflow.genai.agent_server import ( AgentServer, setup_mlflow_git_based_version_tracking, ) agent_server = AgentServer("ResponsesAgent") app = agent_server.app # 可选：开启基于 git 的版本跟踪， # 让每次 trace 对应到一个具体的 git commit setup_mlflow_git_based_version_tracking() def main(): # 支持多 worker 时，以 import string 的形式传入 app agent_server.run(app_import_string="start_server:app") if __name__ == "__main__": main()

• AgentServer("ResponsesAgent")中的类型参数开启自动的输入/输出校验和流式 tracing 聚合；目前只支持"ResponsesAgent"。
• setup_mlflow_git_based_version_tracking()是可选的，它会把活跃模型名与当前 git commit 短哈希关联（本地开发时 app 名为local，部署为 Databricks App 时取DATABRICKS_APP_NAME环境变量）。
• agent_server.run(app_import_string="start_server:app")以模块内app对象作为 import string 传给 uvicorn，这是多 worker 部署的方式。

启动服务器

python3 start_server.py --reload

--reload会在代码变更时自动重载服务器，适合开发阶段。生产环境常用的参数（来自同一启动入口的命令行参数解析）：

# 传入 worker 数以支持多个并发请求（默认 1） python3 start_server.py --workers 4 # 指定端口（默认 8000） python3 start_server.py --port 8000

服务器默认监听0.0.0.0，因此局域网内或其他容器可通过主机地址访问。

验证服务是否可用

启动成功后，向/invocations端点发送一个 JSON 请求测试：

curl -X POST http://localhost:8000/invocations \ -H "Content-Type: application/json" \ -d '{ "input": [{ "role": "user", "content": "What is the 14th Fibonacci number?"}]}'

请求成功时服务器返回代理的输出（ResponsesAgentResponse结构）。如果请求体不是合法 JSON 或不满足ResponsesAgent的校验，会收到 400；代理执行抛异常时会收到 500 及错误信息。

除了/invocations，服务器还暴露了以下端点，可用于健康检查和部署监控：

• GET /health：返回{"status": "healthy"}；
• GET /agent/info：返回 app 名、use_case: "agent"、mlflow 版本，ResponsesAgent 类型下还包含agent_api: "responses"；
• POST /responses：与/invocations同源的端点，用于兼容 OpenAI 客户端的client.responses.create(...)调用方式。

如果注册了@stream函数，可以在请求体中加入"stream": true发送流式请求，响应以data: <json>分块推送，并以data: [DONE]结束：

curl -X POST http://localhost:8000/invocations \ -H "Content-Type: application/json" \ -d '{ "input": [{ "role": "user", "content": "What is the 14th Fibonacci number?"}], "stream": true }'

查看自动记录的 Traces

Agent Server 会自动为每次调用创建 MLflow trace（请求作为输入、响应作为输出写入对应 span）。测试你的代理后，打开 MLflow UI，点击 "Traces" 标签页即可看到这些调用记录。若开启了setup_mlflow_git_based_version_tracking()，trace 会关联到对应的 git commit 版本，方便回溯是哪个代码版本产生了某次调用。

可选：用mlflow.genai.evaluate评估代理

测试通过后，可以用同一套@invoke注册的函数做离线评估，而不用真正启动服务器。创建一个eval_agent.py：

import asyncio import mlflow # 导入 agent 模块，@invoke 注册的函数才能被找到 from agent import agent # noqa: F401 from mlflow.genai.agent_server import get_invoke_function from mlflow.genai.scorers import RelevanceToQuery, Safety from mlflow.types.responses import ResponsesAgentRequest, ResponsesAgentResponse eval_dataset = [ { "inputs": { "request": {"input": [{"role": "user", "content": "What's the 15th Fibonacci number"}]} }, "expected_response": "The 15th Fibonacci number is 610.", } ] def sync_invoke_fn(request: dict) -> ResponsesAgentResponse: # 获取通过 @invoke 装饰器注册的函数 invoke_fn = get_invoke_function() return asyncio.run(invoke_fn(ResponsesAgentRequest(**request))) mlflow.genai.evaluate( data=eval_dataset, predict_fn=sync_invoke_fn, scorers=[RelevanceToQuery(), Safety()], )

运行python3 eval_agent.py，控制台会输出评估结果和 MLflow run 信息；在 MLflow UI 的实验页面可以找到对应的 run，点击 run 名称查看聚合的指标和元数据。

限制说明

• 服务器类型目前仅支持"ResponsesAgent"；不传agent_type时不做输入/输出校验，也不做流式 tracing 聚合。
• @invoke与@stream各只能注册一个函数。
• --reload依赖热重载，适合开发；生产部署应使用--workers控制并发。
• AgentServer还支持enable_chat_proxy=True参数，把未匹配的请求代理到CHAT_APP_PORT（默认 3000）端口上的 chat 应用，只放行/、/favicon.ico、/ping、/assets/*、/api/*、/chat/*等路径以防 SSRF，这是搭配前端聊天应用部署时的可选项，不是本场景必需步骤。

完整示例与参数说明可对照仓库中的 Agent Server 文档，端点与命令行参数的实现见 server.py。

【免费下载链接】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