Skip to main content
Open on GitHub

MLflow AI Gateway for LLMs

LLM 的 MLflow AI 网关 是一个强大的工具,旨在简化组织内各种大型语言模型(LLM)提供商(如 OpenAI 和 Anthropic)的使用和管理。它提供了一个高级接口,通过提供一个统一的端点来处理特定的 LLM 相关请求,从而简化了与这些服务的交互。

安装与设置

安装带有 MLflow GenAI 依赖的 mlflow

pip install 'mlflow[genai]'

将 OpenAI API 密钥设置为环境变量:

export OPENAI_API_KEY=...

创建一个配置文件:

endpoints:
- name: completions
endpoint_type: llm/v1/completions
model:
provider: openai
name: text-davinci-003
config:
openai_api_key: $OPENAI_API_KEY

- name: embeddings
endpoint_type: llm/v1/embeddings
model:
provider: openai
name: text-embedding-ada-002
config:
openai_api_key: $OPENAI_API_KEY

启动网关服务器:

mlflow gateway start --config-path /path/to/config.yaml

示例由 MLflow

The mlflow.langchain 模块提供了用于记录和加载 LangChain 模型的 API。 此模块以 langchain 风格导出多元 LangChain 模型,并以 pyfunc 风格导出单变量 LangChain 模型。

有关更多信息,请参阅 API 文档和示例

补全示例

import mlflow
from langchain.chains import LLMChain, PromptTemplate
from langchain_community.llms import Mlflow

llm = Mlflow(
target_uri="http://127.0.0.1:5000",
endpoint="completions",
)

llm_chain = LLMChain(
llm=Mlflow,
prompt=PromptTemplate(
input_variables=["adjective"],
template="Tell me a {adjective} joke",
),
)
result = llm_chain.run(adjective="funny")
print(result)

with mlflow.start_run():
model_info = mlflow.langchain.log_model(chain, "model")

model = mlflow.pyfunc.load_model(model_info.model_uri)
print(model.predict([{"adjective": "funny"}]))
API 参考:LLMChain |MLflow

嵌入示例

from langchain_community.embeddings import MlflowEmbeddings

embeddings = MlflowEmbeddings(
target_uri="http://127.0.0.1:5000",
endpoint="embeddings",
)

print(embeddings.embed_query("hello"))
print(embeddings.embed_documents(["hello"]))
API 参考:Mlflow 嵌入

聊天示例

from langchain_community.chat_models import ChatMlflow
from langchain_core.messages import HumanMessage, SystemMessage

chat = ChatMlflow(
target_uri="http://127.0.0.1:5000",
endpoint="chat",
)

messages = [
SystemMessage(
content="You are a helpful assistant that translates English to French."
),
HumanMessage(
content="Translate this sentence from English to French: I love programming."
),
]
print(chat(messages))