Pathway LLM xpack 实战：用 UDF 包装器接入 OpenAI、LiteLLM、HuggingFace 与 Cohere 聊天模型

本文以 Pathway Live Data Framework 的 pathway.xpacks.llm.llms 模块为主线，讲解如何用聊天包装器（Chat Wrapper）把 LLM 调用嵌入实时数据流：包括 OpenAIChat、LiteLLMChat、HFPipelineChat、CohereChat 四个包装器的构造参数、消息格式约定、按行覆盖参数的写法，以及异步并发、重试与缓存三大行为开关。读完本文，你可以直接复制示例代码，在表格列上批量驱动任意主流 LLM，并理解包装器底层的执行器与重试实现。

聊天包装器总览与 UDF 设计

Pathway 的 LLM xpack 开箱即用地提供文本生成（chat）与向量嵌入（embedding）两类包装器。对于文本生成，官方提供 OpenAI 聊天模型的原生包装器和可本地运行的 HuggingFace 包装器；而 Azure OpenAI、HuggingFace API、Gemini 等大量其他模型则可以经由 LiteLLMChat 包装器统一接入，支持的模型清单以 LiteLLM 官方文档为准。

当前仓库中，xpack 提供（或对应文档列出了）以下聊天包装器：

• OpenAIChat —— OpenAI 聊天 API 原生封装；
• LiteLLMChat —— 通过 LiteLLM 统一接入各家模型；
• HFPipelineChat —— 本地运行 HuggingFace transformers.pipeline；
• CohereChat —— Cohere 聊天服务，支持附带参考文档并返回引用。

对应实现集中在 llms.py，依赖则声明在 pyproject.toml 的 xpack-llm extras（含 openai、litellm、cohere 等）与 xpack-llm-local extras（含 transformers >= 4.50.2, < 5.0，用于本地推理）。按需安装即可，例如：

pip install "pathway[xpack-llm]"        # 接入云端 LLM API
pip install "pathway[xpack-llm-local]" # 本地跑 HuggingFace 模型

每个包装器本质上是一个 UDF

所有聊天包装器都继承自 BaseChat，而 BaseChat 又是 pw.UDF（用户自定义函数）的子类。UDF 的一般语义是：接受输入、处理、返回输出；在 Pathway 语境下，它让自定义逻辑（这里就是调用 LLM）能无缝作用在表格的列上。包装器把“向模型发请求、取回输出”封装成一个可在列表达式中调用的对象，从而可以在流式表上对每一行 prompt 独立发起模型调用。

使用任何包装器的通用模式是：先构造包装器实例，再把实例应用到包含 prompt 的列上。官方文档给出的示例统一使用下面这张查询表（后文示例均基于它）：

import pathway as pw
queries = pw.debug.table_from_markdown(
    """
questions | max_tokens
How many 'r' there are in 'strawberry'? | 400
""",
    split_on_whitespace=False,
)

OpenAIChat：原生 OpenAI 封装

基本用法

通过 OpenAIChat 类创建包装器，然后把 questions 列中的问题批量发给 OpenAI：

import os
from pathway.xpacks.llm import llms

model = llms.OpenAIChat(
    model="gpt-4o-mini",
    api_key=os.environ["OPENAI_API_KEY"],  # 从环境变量读取 OpenAI API key
)
# 把 queries 表的 question 列作为问题发给 OpenAI
responses = queries.select(result=model(llms.prompt_chat_single_qa(pw.this.questions)))
# 执行计算（含向 OpenAI 发请求）并打印结果表
pw.debug.compute_and_print(responses)

在 Pathway Templates（YAML 配置驱动的应用）中，同一个包装器可以直接以标签形式声明，API key 用环境变量占位：

chat: !pw.xpacks.llm.llms.OpenAIChat
  model: "gpt-4o-mini"
  api_key: $OPENAI_API_KEY

消息（Message）格式

OpenAIChat 要求消息符合 OpenAI Chat API 的格式——即一个字典列表，每个字典是对话中的单条消息（含 role 与 content）。对单轮问答，官方提供了 prompt_chat_single_qa 这个 UDF，它把问题字符串包装成 [{"role": "user", "content": ...}] 这样的 pw.Json：

# llms.py 中 prompt_chat_single_qa 的核心实现
@pw.udf
def prompt_chat_single_qa(question: str) -> pw.Json:
    return pw.Json([dict(role="user", content=question)])

如果需要更精细地控制对话（例如加 system 消息），可以自己构造消息列：

messages = pw.debug.table_from_rows(
    pw.schema_from_types(questions=list[dict]),
    rows=[
        (
            [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": "How many 'r' there are in 'strawberry'?"},
            ],
        )
    ],
)
responses = messages.select(result=model(pw.this.questions))
pw.debug.compute_and_print(responses)

从源码结构看，OpenAIChat.__wrapped__ 会先用 _prepare_messages 把 pw.Json 或原生 list[dict] 统一解码为普通字典，再合并构造期与调用期的 kwargs，最终通过 self.client.chat.completions.create 发出请求。构造函数内部创建的是 openai.AsyncOpenAI 客户端，并且显式设置 max_retries=0——因为重试被移到了 __wrapped__ 内的 retry_strategy.invoke 层，这样无论走框架执行器还是直接调用 __wrapped__，重试行为都一致。

模型参数：构造期默认值 + 调用期覆盖

OpenAI API 接受大量参数（model、api_key、max_tokens、temperature、top_p、frequency_penalty、response_format 等，完整清单见 llms.py 中 OpenAIChat 的 Args 文档）。OpenAIChat 的设计是：构造时设置默认值，调用时可用列表达式覆盖。

model = llms.OpenAIChat(
    model="gpt-4o-mini",
    api_key=os.environ["OPENAI_API_KEY"],
    max_tokens=200,  # 设置 max_tokens 默认值为 200
)
# 第一处：max_tokens 未在调用中指定，将使用默认值 200
responses = queries.select(result=model(llms.prompt_chat_single_qa(pw.this.questions)))
# 第二处：max_tokens 取自 max_tokens 列，逐行覆盖默认值
responses = queries.select(
    result=model(llms.prompt_chat_single_qa(pw.this.questions), max_tokens=pw.this.max_tokens)
)
pw.debug.compute_and_print(responses)

这种“默认值 + 逐行覆盖”的能力并非仅针对 max_tokens，而是通用机制：_extract_value_inside_dict 在每次调用时把 kwargs 中包装成列的 pw.Json 值还原为实际数据，从而实现按行取参。

更灵活的做法是连 model 本身都逐行指定：构造函数中把 model=None，调用时传 model=t.model（该用法在文档字符串与集成测试 test_llm_apply_openai 中都有体现）。

从源码结构看，框架在应用期还会通过 _check_model_accepts_arg 借助 litellm.get_supported_openai_params 校验当前模型是否支持某个参数（对应 OpenAIChat._accepts_call_arg），避免把 gpt-4o-mini 不支持的参数发给 API；单元测试 test_openai_call_args 即验证了 top_p、temperature、max_tokens 合法而 made_up_arg 非法。

LiteLLMChat：一个包装器接入 Gemini 等众多模型

LiteLLMChat 通过 LiteLLM 统一 API 接入各家模型。以 Gemini 为例，只需把 model 设为 "gemini/gemini-pro"：

from pathway.xpacks.llm import llms

model = llms.LiteLLMChat(
    model="gemini/gemini-pro",  # 选择目标模型
    api_key=os.environ["GEMINI_API_KEY"],  # 从环境变量读取 GEMINI API key
)
# 向 Gemini 提问（问题来自 prompt 列）
responses = queries.select(result=model(llms.prompt_chat_single_qa(pw.this.questions)))
pw.debug.compute_and_print(responses)

Templates 中同样支持 YAML 声明：

llm: !pw.xpacks.llm.llms.LiteLLMChat
  model: "gemini/gemini-pro"

与 OpenAIChat 不同，LiteLLMChat 没有默认模型——构造函数中 model 默认为 None，必须在构造时或每次调用时给出。从源码结构看，其 __wrapped__ 内部直接调用同步的 litellm.completion（代码注释说明：因 Ollama 的 json 模式行为差异，暂时关闭了 async 调用），并从 ret.choices[0]["message"]["content"] 取回答；同时它也会按 provider/model 前缀拆分模型名，复用 _check_model_accepts_arg 做参数校验，例如 "replicate/meta/meta-llama-3-8b" 会正确处理多层级命名。

HFPipelineChat：本地运行 HuggingFace 模型

对于希望本地运行的 HuggingFace 模型，使用 HFPipelineChat 包装器（若要调用 HuggingFace 的 API，则应走 LiteLLMChat）。构造实例时会同步初始化一个 transformers.pipeline，因此所有 pipeline 初始化参数（包括模型名）都必须在构造期给出；而 pipeline.__call__ 的参数可以像前文一样在构造期设默认、在调用期覆盖。

from pathway.xpacks.llm import llms

model = llms.HFPipelineChat(
    model="gpt2",  # 选择目标模型
)
responses = queries.select(result=model(pw.this.questions))
pw.debug.compute_and_print(responses)

llm: !pw.xpacks.llm.llms.HFPipelineChat
  model: "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

提问格式取决于模型

HuggingFace pipeline 期望的输入格式与模型相关：gpt2 这类基础模型期望 prompt 字符串；而对话模型则接受字典列表形式的消息，此时会使用模型自带的 prompt 模板渲染对话。例如对对话模型 TinyLlama/TinyLlama-1.1B-Chat-v1.0：

from pathway.xpacks.llm import llms

model = llms.HFPipelineChat(
    model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
)
responses = queries.select(result=model(llms.prompt_chat_single_qa(pw.this.questions)))
pw.debug.compute_and_print(responses)

HFPipelineChat 的构造参数在 llms.py 中定义：model（默认 "gpt2"）、call_kwargs（传给每次 pipeline 调用的参数，可在应用期覆盖）、device（默认 "cpu"）、batch_size（默认 32，即单批送入的最大行数，GPU 上加大批次通常更快）。它的 __wrapped__ 会做智能批处理：当所有行的调用参数一致且 tokenizer 具备 pad token 时整批调用 pipeline；否则退化为逐行推理。此外它还提供 crop_to_max_length 辅助方法，可用内置 tokenizer 把输入字符串按 token 数截断（默认 500 token），防止超长 prompt。

注意：模板（AI pipeline）场景下只接受对话模型，gpt2 这类非对话模型不可用。

CohereChat：附带文档并返回引用

CohereChat 封装 Cohere 的聊天服务，特色是允许把参考文档一并传入（RAG 场景），返回值是 (response, citations) 元组——回答文本与被引用的文档列表（无引用时为空列表）。模型名默认 "command"。

from pathway.xpacks.llm import llms

model = llms.CohereChat()
queries_with_docs = pw.debug.table_from_rows(
    schema=pw.schema_from_types(questions=str, docs=list[dict]),
    rows=[
        (
            "What is RAG?",
            [
                {"text": "Pathway is a high-throughput, low-latency data processing framework that handles live data & streaming for you."},
                {"text": "RAG stands for Retrieval Augmented Generation."},
            ],
        )
    ],
)

r = queries_with_docs.select(
    ret=model(llms.prompt_chat_single_qa(pw.this.questions), documents=pw.this.docs)
)
parsed_table = r.select(response=pw.this.ret[0], citations=pw.this.ret[1])
pw.debug.compute_and_print(parsed_table)

从源码结构看，其 __wrapped__ 会取消息列表的最后一条作为 message、其余作为 chat_history，再把 documents 交给 Cohere 客户端，最后把 ret.citations 序列化为字典列表返回——这正是示例中 pw.this.ret[0]/pw.this.ret[1] 能分别取出回答与引用的原因。

异步行为：capacity、retry_strategy 与 cache_strategy

OpenAIChat 与 LiteLLMChat（chat 与 embedding 包装器同理）都是异步的，Pathway 提供三个在构造期设置的参数控制其行为：

• capacity：允许的最大并发操作数；
• retry_strategy：失败重试策略；
• cache_strategy：调用结果的缓存机制。

model = llms.OpenAIChat(
    # 最大并发数为 10
    capacity=10,
    # 失败时重试 5 次，每次等待时间翻倍（初始 1000ms，backoff_factor=2）
    retry_strategy=pw.udfs.ExponentialBackoffRetryStrategy(
        max_retries=5, initial_delay=1000, backoff_factor=2
    ),
    # 若设置了 PATHWAY_PERSISTENT_STORAGE，则用它来缓存 LLM 调用
    cache_strategy=pw.udfs.DefaultCache(),
    # 选择模型
    model="gpt-4o-mini",
    # 从环境变量读取 OpenAI API key
    api_key=os.environ["OPENAI_API_KEY"],
)
responses = queries.select(result=model(llms.prompt_chat_single_qa(pw.this.questions)))
pw.debug.compute_and_print(responses)

YAML 等价写法：

chat: !pw.xpacks.llm.llms.OpenAIChat
  model: "gpt-4o-mini"
  capacity: 10
  retry_strategy: !pw.udfs.ExponentialBackoffRetryStrategy
    max_retries: 5
    initial_delay: 1000
    backoff_factor: 2

结合源码，这几个参数的落地方式可以进一步看清：

1. 执行器：OpenAIChat/LiteLLMChat/CohereChat 构造时经由 _prepare_executor 创建执行器，capacity 决定并发上限；执行器还有 async_mode 开关（默认 "batch_async"，也可设为 "fully_async"），分别对应 pw.udfs.async_executor 与 pw.udfs.fully_async_executor 两种执行模式。capacity=None 表示不显式限制并发。
2. 重试：各包装器 retry_strategy 默认是 ExponentialBackoffRetryStrategy()（指数退避），传 None 则降级为 NoRetryStrategy。单元测试 test_openai_chat_wrapped_retries_transient_errors 用 mock 验证了这一点：第一次调用抛运行时异常后，FixedDelayRetryStrategy(max_retries=4) 会再次发起请求并成功取回结果。
3. 缓存：cache_strategy 默认为 None（不缓存）；传入如 DefaultCache() 后，重复的调用可命中缓存，避免重复计费。单元测试 test_openai_chat_init 覆盖了 DiskCache()、None 等不同取值的构造行为。

小结与延伸阅读

本文沿 80.llm-chats.md 的脉络，完整覆盖了 Pathway LLM xpack 聊天包装器的四类实现：

包装器	适用场景	关键特性
OpenAIChat	OpenAI 原生 API	默认模型 gpt-3.5-turbo（见 构造函数），支持按行覆盖任意 OpenAI 参数
LiteLLMChat	Gemini、Azure OpenAI 等多模型	无默认模型，provider/model 命名约定
HFPipelineChat	本地推理	构造期初始化 pipeline，智能批处理，device/batch_size 可调
CohereChat	RAG 附文档问答	返回 (response, citations) 元组

三者共享“构造期设默认、调用期用列覆盖”的参数模型，且云端包装器均为异步执行，可通过 capacity / retry_strategy / cache_strategy 控制并发、重试与缓存。

进一步阅读可参考：

• 包装器完整实现：llms.py（含 prompt_chat_single_qa 与 BaseChat 基类）；
• 单元测试：test_llms.py（初始化、模型字段、重试、参数校验）；
• 集成测试：test_llms.py（真实 API 调用与按行覆盖 model 列的用法）；
• 依赖声明：pyproject.toml 中的 xpack-llm 与 xpack-llm-local extras；
• 同目录下的 LLM xpack 概览与 RAG 应用文档：10.overview.md、20.llm-app.md。