diff --git a/main.py b/main.py new file mode 100644 index 0000000..685da28 --- /dev/null +++ b/main.py @@ -0,0 +1,100 @@ +"""CLI tool for extracting structured data from text using LangChain and Pydantic. + +The script supports two schemas: +- PersonInfo +- MeetingNotes + +It automatically detects the type of input text and runs the appropriate chain. +""" + +import argparse +import sys +from datetime import datetime +from typing import List, Optional + +from langchain_core.output_parsers import PydanticOutputParser +from langchain_core.prompts import PromptTemplate +from langchain_openai import ChatOpenAI +from pydantic import BaseModel, Field + +# --------------------------------------------------------------------------- +# Pydantic models +# --------------------------------------------------------------------------- +class PersonInfo(BaseModel): + name: str = Field(..., description="Имя человека") + age: Optional[int] = Field(None, description="Возраст, если известен") + profession: str = Field(..., description="Профессия человека") + skills: List[str] = Field(..., description="Список навыков") + +class MeetingNotes(BaseModel): + date: datetime = Field(..., description="Дата встречи") + participants: List[str] = Field(..., description="Участники встречи") + topics: List[str] = Field(..., description="Темы обсуждения") + decisions: List[str] = Field(..., description="Принятые решения") + next_steps: List[str] = Field(..., description="Следующие шаги") + +# --------------------------------------------------------------------------- +# Prompt templates +# --------------------------------------------------------------------------- +PERSON_PROMPT = PromptTemplate( + input_variables=["text"], + template=( + """ + Извлечь из следующего текста информацию о человеке в формате JSON, соответствующем схеме PersonInfo. + Текст: {text} + JSON: """ + ), +) + +MEETING_PROMPT = PromptTemplate( + input_variables=["text"], + template=( + """ + Извлечь из следующего текста информацию о встрече в формате JSON, соответствующем схеме MeetingNotes. + Текст: {text} + JSON: """ + ), +) + +# --------------------------------------------------------------------------- +# Helper functions +# --------------------------------------------------------------------------- + +def detect_schema(text: str) -> str: + """Простейшая эвристика: если в тексте встречается слово "встреча" или "meeting" – считаем это встречей. + Иначе – человек. + """ + lowered = text.lower() + if "встреча" in lowered or "meeting" in lowered: + return "meeting" + return "person" + + +def build_chain(schema: str): + if schema == "person": + parser = PydanticOutputParser(pydantic_object=PersonInfo) + chain = PERSON_PROMPT | ChatOpenAI(temperature=0) | parser + else: + parser = PydanticOutputParser(pydantic_object=MeetingNotes) + chain = MEETING_PROMPT | ChatOpenAI(temperature=0) | parser + return chain + +# --------------------------------------------------------------------------- +# CLI entry point +# --------------------------------------------------------------------------- + +def main(): + parser = argparse.ArgumentParser(description="Extract structured data from text.") + parser.add_argument("--text", required=True, help="Input text to parse") + args = parser.parse_args() + + schema = detect_schema(args.text) + chain = build_chain(schema) + result = chain.invoke({"text": args.text}) + + # result is a Pydantic model instance + print(result.model_dump(indent=2)) + + +if __name__ == "__main__": + main()