enterprise-kb-rag/2.2_LangChain基础知识/15_output_parser_pydantic.py

from pathlib import Path
from typing import List
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 PydanticOutputParser
    from langchain_core.prompts import ChatPromptTemplate
    from pydantic import BaseModel, Field
except ModuleNotFoundError as exc:
    print_dependency_hint(exc)
    raise SystemExit(1)


class Movie(BaseModel):
    title: str = Field(description="电影标题")
    director: str = Field(description="导演姓名")
    year: int = Field(description="上映年份")
    genre: List[str] = Field(description="电影类型")
    rating: float = Field(description="评分（1-10）")


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

    pydantic_parser = PydanticOutputParser(pydantic_object=Movie)
    format_instructions = pydantic_parser.get_format_instructions()
    pydantic_prompt = ChatPromptTemplate.from_template(
        "生成一部{genre}电影的信息。\n{format_instructions}"
    )

    pydantic_chain = pydantic_prompt | model | pydantic_parser
    movie_data = pydantic_chain.invoke(
        {
            "genre": "科幻",
            "format_instructions": format_instructions,
        }
    )
    print(movie_data)


if __name__ == "__main__":
    main()