enterprise-kb-rag/2.2_LangChain基础知识/16_output_parser_custom.py

from pathlib import Path
from typing import Any, Dict
import sys

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

from _chapter2_config import get_chat_model, print_dependency_hint


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


class CustomKeyValueParser(BaseOutputParser[Dict[str, Any]]):
    """解析形如 'key: value' 的文本。"""

    def parse(self, text: str) -> Dict[str, Any]:
        result = {}
        lines = text.strip().split("\n")
        for line in lines:
            if ":" in line:
                key, value = line.split(":", 1)
                result[key.strip()] = value.strip()
        return result

    def get_format_instructions(self) -> str:
        return """请以'键: 值'的格式返回信息，每行一个键值对。
例如：
名称: 爱因斯坦
职业: 物理学家
贡献: 相对论"""


def main() -> None:
    try:
        model = get_chat_model(temperature=0)
    except ModuleNotFoundError as exc:
        print_dependency_hint(exc)
        raise SystemExit(1)

    custom_parser = CustomKeyValueParser()
    custom_prompt = ChatPromptTemplate.from_template(
        "提供关于{person}的基本信息。\n{format_instructions}"
    )

    custom_chain = custom_prompt | model | custom_parser
    result = custom_chain.invoke(
        {
            "person": "屠呦呦",
            "format_instructions": custom_parser.get_format_instructions(),
        }
    )
    print(result)


if __name__ == "__main__":
    main()