add main.py with Pydantic models and structured output extraction

This commit is contained in:
2026-05-28 16:28:57 +00:00
parent c1659e2104
commit eb8b9b7ea1
+144
View File
@@ -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()