自然语言处理从入门到应用——LangChain：链（Chains）-[通用功能：自定义Chain和Chain的异步API]

分类目录：《自然语言处理从入门到应用》总目录

---

创建自定义Chain

要实现自己的自定义链式连接，我们可以子类化Chain并实现以下方法：

from __future__ import annotations
from typing import Any, Dict, List, Optional
from pydantic import Extra
from langchain.base_language import BaseLanguageModel
from langchain.callbacks.manager import (
    AsyncCallbackManagerForChainRun,
    CallbackManagerForChainRun,
)
from langchain.chains.base import Chain
from langchain.prompts.base import BasePromptTemplate


class MyCustomChain(Chain):
    """
    An example of a custom chain.
    """

    prompt: BasePromptTemplate
    """Prompt object to use."""
    llm: BaseLanguageModel
    output_key: str = "text"  #: :meta private:

    class Config:
        """Configuration for this pydantic object."""

        extra = Extra.forbid
        arbitrary_types_allowed = True

    @property
    def input_keys(self) -> List[str]:
        """Will be whatever keys the prompt expects.

        :meta private:
        """
        return self.prompt.input_variables

    @property
    def output_keys(self) -> List[str]:
        """Will always return text key.

        :meta private:
        """
        return [self.output_key]

    def _call(
        self,
        inputs: Dict[str, Any],
        run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> Dict[str, str]:
        # 在这里编写你的自定义链逻辑
        # 下面的示例仅模仿了 LLMChain
        prompt_value = self.prompt.format_prompt(**inputs)
        
        # 当调用语言模型或其他链时，应该将回调管理器传递给它。
        # 这样可以让内部运行受到外部运行注册的任何回调的跟踪。
        # 你可以通过调用 `run_manager.get_child()` 获取回调管理器，如下所示。
        response = self.llm.generate_prompt(
            [prompt_value],
            callbacks=run_manager.get_child() if run_manager else None
        )

        # 如果想要记录此次运行的某些信息，可以通过调用 `run_manager` 上的方法来实现。
        # 这将触发为该事件注册的任何回调。
        if run_manager:
            run_manager.on_text("记录此次运行的一些信息")
        
        return {self.output_key: response.generations[0][0].text}

    async def _acall(
        self,
        inputs: Dict[str, Any],
        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
    ) -> Dict[str, str]:
        # 在这里编写你的自定义链逻辑
        # 下面的示例仅模仿了 LLMChain
        prompt_value = self.prompt.format_prompt(**inputs)
        
        # 当调用语言模型或其他链时，应该将回调管理器传递给它。
        # 这样可以让内部运行受到外部运行注册的任何回调的跟踪。
        # 你可以通过调用 `run_manager.get_child()` 获取回调管理器，如下所示。
        response = await self.llm.agenerate_prompt(
            [prompt_value],
            callbacks=run_manager.get_child() if run_manager else None
        )

        # 如果想要记录此次运行的某些信息，可以通过调用 `run_manager` 上的方法来实现。
        # 这将触发为该事件注册的任何回调。
        if run_manager:
            await run_manager.on_text("记录此次运行的一些信息")
        
        return {self.output_key: response.generations[0][0].text}

    @property
    def _chain_type(self) -> str:
        return "my_custom_chain"

from langchain.callbacks.stdout import StdOutCallbackHandler
from langchain.chat_models.openai import ChatOpenAI
from langchain.prompts.prompt import PromptTemplate


chain = MyCustomChain(
    prompt=PromptTemplate.from_template('tell us a joke about {topic}'),
    llm=ChatOpenAI()
)

chain.run({'topic': 'callbacks'}, callbacks=[StdOutCallbackHandler()])

日志输出：

> Entering new MyCustomChain chain...
Log something about this run
> Finished chain.

输出：

Why did the callback function feel lonely? Because it was always waiting for someone to call it back!'

Chain 的异步 API

LangChain通过利用asyncio模块提供了对链式连接的异步支持。目前，LLMChain（通过 arun、apredict和acall方法）、LLMMathChain（通过arun和acall方法）、ChatVectorDBChain和问答链式连接支持异步方法。其他链式连接的异步支持正在计划中。

import asyncio
import time

from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain


def generate_serially():
    llm = OpenAI(temperature=0.9)
    prompt = PromptTemplate(
        input_variables=["product"],
        template="What is a good name for a company that makes {product}?",
    )
    chain = LLMChain(llm=llm, prompt=prompt)
    for _ in range(5):
        resp = chain.run(product="toothpaste")
        print(resp)


async def async_generate(chain):
    resp = await chain.arun(product="toothpaste")
    print(resp)


async def generate_concurrently():
    llm = OpenAI(temperature=0.9)
    prompt = PromptTemplate(
        input_variables=["product"],
        template="What is a good name for a company that makes {product}?",
    )
    chain = LLMChain(llm=llm, prompt=prompt)
    tasks = [async_generate(chain) for _ in range(5)]
    await asyncio.gather(*tasks)

s = time.perf_counter()
# If running this outside of Jupyter, use asyncio.run(generate_concurrently())
await generate_concurrently()
elapsed = time.perf_counter() - s
print('\033[1m' + f"Concurrent executed in {elapsed:0.2f} seconds." + '\033[0m')

s = time.perf_counter()
generate_serially()
elapsed = time.perf_counter() - s
print('\033[1m' + f"Serial executed in {elapsed:0.2f} seconds." + '\033[0m')

输出：

BrightSmile Toothpaste Company


BrightSmile Toothpaste Co.


BrightSmile Toothpaste


Gleaming Smile Inc.


SparkleSmile Toothpaste
Concurrent executed in 1.54 seconds.


BrightSmile Toothpaste Co.


MintyFresh Toothpaste Co.


SparkleSmile Toothpaste.


Pearly Whites Toothpaste Co.


BrightSmile Toothpaste.
Serial executed in 6.38 seconds.

参考文献：
[1] LangChain官方网站：https://www.langchain.com/
[2] LangChain ️ 中文网，跟着LangChain一起学LLM/GPT开发：https://www.langchain.com.cn/
[3] LangChain中文网 - LangChain 是一个用于开发由语言模型驱动的应用程序的框架：http://www.cnlangchain.com/