LangChain 中的 Runnable 接口

379 阅读4分钟

大家好,我是雨飞。今天为大家带来 LangChain 中的 Runnable 接口的讲解。

Runnable 接口是 LangChain LCEL 中一个非常重要的概念,很多 LangChain 的组件都实现了这个接口,比如 chat model、LLMs、output parser、retrievers、Prompt 等。

Runnable 接口

runnable 接口定义了很多同步和异步的方法去保证可以方便的进行 chain 的生成。核心的方法包括下面几个:

  • stream:流式的输出模式
  • invoke:使用单个输入去调用链
  • batch:使用一个数组的输入去调用链

异步的方法,需要使用 await 语法实现异步等待:

  • astream:异步的流式输出模式
  • ainvoke:异步的使用单个输入去调用链
  • abatch:异步的使用一个数组的输入去调用链
  • astream_log:中间步骤发生时进行回调

输入输出类型

组件输入类型输出类型
Promptpython format 格式的字符串PromptValue
ChatModelString、ChatMessage、PromptValueChatMessage
LLMString、ChatMessage、PromptValueString
OutputParserString、ChatMessage依赖 parser 的解析规则
RetrieverStringDocuments
ToolString、dictionary依赖 tool 的规则

我们以一个简单的 Prompt+ChatModel 的链作为示例,去讲一些各个方法的使用。下面是示例代码:

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from model_factory import url_base,api_key2,model_name
model = ChatOpenAI(openai_api_base=url_base, openai_api_key=api_key2, model=model_name)
prompt = ChatPromptTemplate.from_template("请帮我讲一个关于 {topic} 的笑话")
chain = prompt | model

Input Schema 输入模式

你可以使用 .schema() 方法去获取一个 JSON 格式的表示,这个表示就是对 Runnable 接口的输入做了详细的描述。我们上面提到提示词、模型、链都继承了 Runnable 接口,因此这些对象都有.schema() 方法可以调用。

print(chain.input_schema.schema())

print(prompt.input_schema.schema())

print(model.input_schema.schema())

Output Schema 输出模式

同样的,我们也可以用.schema() 方法去获取 Runnable 接口对应的输出。

 print(chain.output_schema.schema())

Stream

for s in chain.stream({"topic": "小兔子"}):
    print(s.content, end="", flush=True)

Invoke

print(chain.invoke({"topic": "小兔子"}))

Batch

print(chain.batch([{"topic": "小兔子"}, {"topic": "大灰狼"}]))

值得注意的是,你可以使用 max_concurrency 这个参数,去控制并发请求的数量。

print(chain.batch([{"topic": "小兔子"}, {"topic": "大灰狼"}], config={"max_concurrency": 5}))

Async Stream

async for s in chain.stream({"topic": "小兔子"}):
    print(s.content, end="", flush=True)

Async Invoke

async def output():
    result = await chain.ainvoke({"topic": "小兔子"})
    print(result)
import asyncio

asyncio.run(output())

Async Batch

await chain.abatch([{"topic": "小兔子"}])

Async Stream Events (beta)

事件流是一个测试版本的 API,目前在 LangChain 0.2.0 版本才被引进。

使用这个 API 时,需要注意下面几点:

  • 在代码中使用异步的方法、工具
  • 需要在自己定义的函数里增加回调
  • 如果在非 LECL 的情况下使用,需要使用 .astream() 方法去让 LLM 进行流式输出

Event Reference

值得注意的是,runnable 接口在进行流式处理时,输入的内容将在输入流全部耗尽后才可用,这意味着输入将在相应的结束钩子而不是开始事件时可用。

下面是一个示例。

from typing import List
from langchain_core.documents import Document
from langchain_core.tools import tool
from langchain_community.vectorstores import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from model_factory import huggingface_bge_embedding
template = """Answer the question based only on the following context:
{context}

Question: {question}
"""
prompt = ChatPromptTemplate.from_template(template)

vectorstore = Chroma.from_texts(
    ["harrison worked at kensho"], embedding=huggingface_bge_embedding
)
retriever = vectorstore.as_retriever()

retrieval_chain = (
    {
        "context": retriever.with_config(run_name="Docs"),
        "question": RunnablePassthrough(),
    }
    | prompt
    | model.with_config(run_name="my_llm")
    | StrOutputParser()
)

async  def stream_events():
    async for event in retrieval_chain.astream_events(
        "where did harrison work?", version="v1", include_names=["Docs", "my_llm"]
    ):
        kind = event["event"]
        # 
        if kind == "on_chat_model_stream":
            print(event["data"]["chunk"].content, end="|")
        elif kind in {"on_chat_model_start"}:
            print()
            print("Streaming LLM:")
        elif kind in {"on_chat_model_end"}:
            print()
            print("Done streaming LLM.")
        elif kind == "on_retriever_end":
            print("--")
            print("Retrieved the following documents:")
            print(event["data"]["output"]["documents"])
        elif kind == "on_tool_end":
            print(f"Ended tool: {event['name']}")
        else:
            pass
import asyncio
asyncio.run(stream_events())

Async Stream Intermediate Steps

所有的 runnable 对象都有一个.astream_log() 的方法,这个方法允许我们输出整个链路上的所有中间结果。

这对于我们 debug 或者展示整体的进度是十分有帮助的,你可以使用 name、tags、metadata 去获取整个链路的数据。

async def stream_log():
    async for chunk in retrieval_chain.astream_log(
        "where did harrison work?", include_names=["Docs"]
    ):
        print("-" * 40)
        print(chunk)
解释

Parallelism

LECL 支持并行的请求,比如我们可以使用 RunnableParallel 去并行的输出每一个元素。这个方法不仅可以用在 invoke 上,也可以在 batch 中进行使用。

from langchain_core.runnables import RunnableParallel

import time

def timeit(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        elapsed_time = end_time - start_time
        print(f"{func.__name__} executed in {elapsed_time:.6f} seconds")
        return result
    return wrapper

chain1 = ChatPromptTemplate.from_template("tell me a joke about {topic}") | model
chain2 = (
    ChatPromptTemplate.from_template("write a short (2 line) poem about {topic}")
    | model
)
combined = RunnableParallel(joke=chain1, poem=chain2)

可以测试不同方法的调用时间。

@timeit
def chain_output():
    result = chain1.invoke({"topic":"小兔子"})
    print(result)
@timeit
def chain2_output():
    result = chain2.invoke({"topic":"小兔子"})
    print(result)

@timeit
def chain3_output():
    result = combined.invoke({"topic":"小兔子"})
    print(result)

chain_output()
chain2_output()
chain3_output()

好了,我写完了,感谢阅读到这里,我们下期再见。