创建自定义链式连接

要实现自己的自定义链式连接，您可以子类化 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!'

# 创建自定义链式连接

创建自定义链式连接