如何运行自定义函数

先决条件

本指南假定您熟悉以下概念：

• LangChain 表达式语言 （LCEL）
• 链接 runnables

您可以将任意函数用作 Runnable。这对于格式化或需要其他 LangChain 组件未提供的功能非常有用，并且用作 Runnables 的自定义函数将调用RunnableLambdas.

请注意，这些函数的所有 inputs 都需要是 SINGLE 参数。如果你有一个接受多个参数的函数，你应该编写一个包装器，它接受单个 dict 输入并将其解压缩为多个参数。

本指南将涵盖：

• 如何使用RunnableLambdaconstructor 和便利性@chain装饰
• 在链中使用时将自定义函数强制转换为 runnables
• 如何在自定义函数中接受和使用运行元数据
• 如何通过让自定义函数返回生成器来使用自定义函数进行流式传输

使用构造函数

下面，我们使用RunnableLambda构造 函数：

%pip install -qU langchain langchain_openai

import os
from getpass import getpass

if "OPENAI_API_KEY" not in os.environ:
    os.environ["OPENAI_API_KEY"] = getpass()

from operator import itemgetter

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI


def length_function(text):
    return len(text)


def _multiple_length_function(text1, text2):
    return len(text1) * len(text2)


def multiple_length_function(_dict):
    return _multiple_length_function(_dict["text1"], _dict["text2"])


model = ChatOpenAI()

prompt = ChatPromptTemplate.from_template("what is {a} + {b}")

chain = (
    {
        "a": itemgetter("foo") | RunnableLambda(length_function),
        "b": {"text1": itemgetter("foo"), "text2": itemgetter("bar")}
        | RunnableLambda(multiple_length_function),
    }
    | prompt
    | model
)

chain.invoke({"foo": "bar", "bar": "gah"})

API 参考：ChatPromptTemplate | RunnableLambda | 聊天OpenAI

AIMessage(content='3 + 9 equals 12.', response_metadata={'token_usage': {'completion_tokens': 8, 'prompt_tokens': 14, 'total_tokens': 22}, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': 'fp_c2295e73ad', 'finish_reason': 'stop', 'logprobs': None}, id='run-73728de3-e483-49e3-ad54-51bd9570e71a-0')

便利性@chain装饰

您还可以通过添加@chain装饰。这在功能上等同于将函数包装在RunnableLambda构造函数，如上所示。下面是一个示例：

from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import chain

prompt1 = ChatPromptTemplate.from_template("Tell me a joke about {topic}")
prompt2 = ChatPromptTemplate.from_template("What is the subject of this joke: {joke}")


@chain
def custom_chain(text):
    prompt_val1 = prompt1.invoke({"topic": text})
    output1 = ChatOpenAI().invoke(prompt_val1)
    parsed_output1 = StrOutputParser().invoke(output1)
    chain2 = prompt2 | ChatOpenAI() | StrOutputParser()
    return chain2.invoke({"joke": parsed_output1})


custom_chain.invoke("bears")

API 参考：StrOutputParser | Chains

'The subject of the joke is the bear and his girlfriend.'

在上面，，@chaindecorator 用于将custom_chain导入到一个 runnable 中，我们使用.invoke()方法。

如果对 LangSmith 使用跟踪，则应看到custom_chain跟踪，下面嵌套了对 OpenAI 的调用。

Chains中的自动强制

当在链中使用自定义函数时，如果管道作符 （|），则可以省略RunnableLambda或@chain构造函数并依赖强制。下面是一个简单的示例，其中包含一个函数，该函数从模型获取输出并返回其前五个字母：

prompt = ChatPromptTemplate.from_template("tell me a story about {topic}")

model = ChatOpenAI()

chain_with_coerced_function = prompt | model | (lambda x: x.content[:5])

chain_with_coerced_function.invoke({"topic": "bears"})

'Once '

请注意，我们不需要包装自定义函数(lambda x: x.content[:5])在RunnableLambda构造函数，因为model在管道运算符的左侧已经是一个 Runnable。自定义函数被强制转换为可运行对象。有关更多信息，请参阅此部分。

传递运行元数据

可运行的 lambda 可以选择接受 RunnableConfig 参数，该参数可用于将回调、标签和其他配置信息传递给嵌套运行。

import json

from langchain_core.runnables import RunnableConfig


def parse_or_fix(text: str, config: RunnableConfig):
    fixing_chain = (
        ChatPromptTemplate.from_template(
            "Fix the following text:\n\n\`\`\`text\n{input}\n\`\`\`\nError: {error}"
            " Don't narrate, just respond with the fixed data."
        )
        | model
        | StrOutputParser()
    )
    for _ in range(3):
        try:
            return json.loads(text)
        except Exception as e:
            text = fixing_chain.invoke({"input": text, "error": e}, config)
    return "Failed to parse"


from langchain_community.callbacks import get_openai_callback

with get_openai_callback() as cb:
    output = RunnableLambda(parse_or_fix).invoke(
        "{foo: bar}", {"tags": ["my-tag"], "callbacks": [cb]}
    )
    print(output)
    print(cb)

API 参考：RunnableConfig | get_openai_callback

{'foo': 'bar'}
Tokens Used: 62
	Prompt Tokens: 56
	Completion Tokens: 6
Successful Requests: 1
Total Cost (USD): $9.6e-05

from langchain_community.callbacks import get_openai_callback

with get_openai_callback() as cb:
    output = RunnableLambda(parse_or_fix).invoke(
        "{foo: bar}", {"tags": ["my-tag"], "callbacks": [cb]}
    )
    print(output)
    print(cb)

API 参考：get_openai_callback

{'foo': 'bar'}
Tokens Used: 62
	Prompt Tokens: 56
	Completion Tokens: 6
Successful Requests: 1
Total Cost (USD): $9.6e-05

流

注意

RunnableLambda 最适合不需要支持流式处理的代码。如果您需要支持流式处理（即能够对 Importing 块进行作并产生输出块），请改用 RunnableGenerator，如下例所示。

您可以使用生成器函数（即使用yield关键字，并且其行为类似于迭代器）。

这些生成器的签名应该是Iterator[Input] -> Iterator[Output].或者对于异步生成器：AsyncIterator[Input] -> AsyncIterator[Output].

这些方法可用于：

• 实现自定义输出解析器
• 修改上一步的输出，同时保留流式处理功能

下面是逗号分隔列表的自定义输出解析器示例。首先，我们创建一个链，生成 text 这样的列表：

from typing import Iterator, List

prompt = ChatPromptTemplate.from_template(
    "Write a comma-separated list of 5 animals similar to: {animal}. Do not include numbers"
)

str_chain = prompt | model | StrOutputParser()

for chunk in str_chain.stream({"animal": "bear"}):
    print(chunk, end="", flush=True)

lion, tiger, wolf, gorilla, panda

接下来，我们定义一个自定义函数，该函数将聚合当前流式输出，并在模型生成列表中的下一个逗号时生成该输出：

# This is a custom parser that splits an iterator of llm tokens
# into a list of strings separated by commas
def split_into_list(input: Iterator[str]) -> Iterator[List[str]]:
    # hold partial input until we get a comma
    buffer = ""
    for chunk in input:
        # add current chunk to buffer
        buffer += chunk
        # while there are commas in the buffer
        while "," in buffer:
            # split buffer on comma
            comma_index = buffer.index(",")
            # yield everything before the comma
            yield [buffer[:comma_index].strip()]
            # save the rest for the next iteration
            buffer = buffer[comma_index + 1 :]
    # yield the last chunk
    yield [buffer.strip()]


list_chain = str_chain | split_into_list

for chunk in list_chain.stream({"animal": "bear"}):
    print(chunk, flush=True)

['lion']
['tiger']
['wolf']
['gorilla']
['raccoon']

调用 this 会得到一个完整的值数组：

list_chain.invoke({"animal": "bear"})

['lion', 'tiger', 'wolf', 'gorilla', 'raccoon']

异步版本

如果您在asyncenvironment 中，下面是一个async版本：

from typing import AsyncIterator


async def asplit_into_list(
    input: AsyncIterator[str],
) -> AsyncIterator[List[str]]:  # async def
    buffer = ""
    async for (
        chunk
    ) in input:  # `input` is a `async_generator` object, so use `async for`
        buffer += chunk
        while "," in buffer:
            comma_index = buffer.index(",")
            yield [buffer[:comma_index].strip()]
            buffer = buffer[comma_index + 1 :]
    yield [buffer.strip()]


list_chain = str_chain | asplit_into_list

async for chunk in list_chain.astream({"animal": "bear"}):
    print(chunk, flush=True)

['lion']
['tiger']
['wolf']
['gorilla']
['panda']

await list_chain.ainvoke({"animal": "bear"})

['lion', 'tiger', 'wolf', 'gorilla', 'panda']

后续步骤

现在，您已经了解了在链中使用自定义逻辑的几种不同方法，以及如何实现流式处理。

要了解更多信息，请参阅本节中有关 runnables 的其他操作指南。