From eb8b9b7ea138a5ada2ded16b7c389f0d40c0147a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=AD=D0=BC=D0=B8=D0=BB=D1=8C=20=D0=90=D0=BC=D0=B8=D1=80?= =?UTF-8?q?=D0=BE=D0=B2?= Date: Thu, 28 May 2026 16:28:57 +0000 Subject: [PATCH] add main.py with Pydantic models and structured output extraction --- main.py | 144 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 144 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..01eabdc --- /dev/null +++ b/main.py @@ -0,0 +1,144 @@ +import os +import sys +from typing import Optional, List +from pydantic import BaseModel, Field +from langchain_core.output_parsers import PydanticOutputParser +from langchain_core.prompts import PromptTemplate +from langchain_openai import ChatOpenAI +from langchain_ollama import ChatOllama + + +# Pydantic модели +class PersonInfo(BaseModel): + """Информация о человеке""" + name: str = Field(description="Имя человека") + age: Optional[int] = Field(default=None, description="Возраст человека (если указан)") + profession: str = Field(description="Профессия человека") + skills: List[str] = Field(default_factory=list, description="Список навыков человека") + + +class MeetingNotes(BaseModel): + """Заметки о встрече""" + date: str = Field(description="Дата встречи") + participants: List[str] = Field(default_factory=list, description="Список участников встрече") + topics: List[str] = Field(default_factory=list, description="Обсуждаемые темы") + decisions: List[str] = Field(default_factory=list, description="Принятые решения") + next_steps: List[str] = Field(default_factory=list, description="Следующие шаги") + + +def get_llm(): + """Получить LLM (OpenAI или Ollama)""" + if os.getenv("OPENAI_API_KEY"): + return ChatOpenAI(model="gpt-4o-mini", temperature=0) + return ChatOllama(model="llama3.2", temperature=0) + + +def create_person_chain(): + """Цепочка для извлечения информации о человеке""" + parser = PydanticOutputParser(pydantic_object=PersonInfo) + prompt = PromptTemplate( + template=""" +Извлеки информацию о человеке из текста. +{format_instructions} + +Текст: {text} +""", + input_variables=["text"], + partial_variables={"format_instructions": parser.get_format_instructions()} + ) + llm = get_llm() + return prompt | llm | parser + + +def create_meeting_chain(): + """Цепочка для извлечения заметок о встрече""" + parser = PydanticOutputParser(pydantic_object=MeetingNotes) + prompt = PromptTemplate( + template=""" +Извлеки информацию о встрече из текста. +{format_instructions} + +Текст: {text} +""", + input_variables=["text"], + partial_variables={"format_instructions": parser.get_format_instructions()} + ) + llm = get_llm() + return prompt | llm | parser + + +def classify_text(text: str) -> str: + """Определить тип текста: 'person' или 'meeting'""" + text_lower = text.lower() + + # Эвристика для определения типа текста + meeting_keywords = ["встреча", "meeting", "участники", "participants", "решения", + "decisions", "темы", "topics", "следующие шаги", "next steps", + "провели", "прошла", "дата", "date"] + + person_keywords = ["год", "лет", "years old", "возраст", "age", "профессия", + "profession", "навыки", "skills", "умеет", "знает"] + + meeting_score = sum(1 for kw in meeting_keywords if kw in text_lower) + person_score = sum(1 for kw in person_keywords if kw in text_lower) + + if meeting_score >= person_score and meeting_score > 0: + return "meeting" + return "person" + + +def extract_structured_data(text: str) -> BaseModel: + """Извлечь структурированные данные из текста""" + text_type = classify_text(text) + + if text_type == "meeting": + chain = create_meeting_chain() + else: + chain = create_person_chain() + + return chain.invoke({"text": text}) + + +def main(): + """CLI для демонстрации""" + # Встроенные примеры + person_example = "Анна, 28 лет, Python-разработчик. Навыки: FastAPI, Docker." + meeting_example = "Встреча 28.05.2026. Участники: Иван, Мария, Петр. Темы: план проекта, бюджет. Решения: утвердить бюджет. Следующие шаги: подготовить ТЗ." + + print("=== Структурированный вывод (Pydantic) ===\n") + + if len(sys.argv) > 1: + # Ввод с аргументом командной строки + text = " ".join(sys.argv[1:]) + else: + # Выбор примера + print("Выберите пример:") + print("1. Информация о человеке") + print("2. Заметки о встрече") + print("3. Ввести свой текст") + + choice = input("Ваш выбор (1-3): ").strip() + + if choice == "1": + text = person_example + elif choice == "2": + text = meeting_example + elif choice == "3": + text = input("Введите текст: ") + else: + print("Неверный выбор, используется пример 1") + text = person_example + + print(f"\nВход: {text}\n") + + try: + result = extract_structured_data(text) + print("Результат:") + print(result.model_dump_json(indent=2)) + print(f"\nТип: {type(result).__name__}") + except Exception as e: + print(f"Ошибка: {e}") + + +if __name__ == "__main__": + main() \ No newline at end of file