enterprise-kb-rag/2.2_LangChain基础知识/11_chain_sequential.py

from pathlib import Path
import sys

sys.path.append(str(Path(__file__).resolve().parents[1]))

from _chapter2_config import as_text, get_chat_model, print_dependency_hint


try:
    from langchain_core.output_parsers import StrOutputParser
    from langchain_core.prompts import PromptTemplate
except ModuleNotFoundError as exc:
    print_dependency_hint(exc)
    raise SystemExit(1)


def main() -> None:
    first_prompt = PromptTemplate(
        input_variables=["lastname"],
        template="我的邻居姓{lastname}，他生了个儿子，给他儿子起个名字",
    )

    try:
        llm = get_chat_model(temperature=0)
    except ModuleNotFoundError as exc:
        print_dependency_hint(exc)
        raise SystemExit(1)

    first_chain = first_prompt | llm | StrOutputParser()

    second_prompt = PromptTemplate(
        input_variables=["child_name"],
        template="邻居的儿子名字叫{child_name}，给他起一个小名",
    )
    second_chain = second_prompt | llm

    overall_chain = first_chain | second_chain
    print(overall_chain)
    print("*" * 80)

    result = overall_chain.invoke({"lastname": "王"})
    print(as_text(result))


if __name__ == "__main__":
    main()