add main.py with Pydantic models and structured output extraction
This commit is contained in:
@@ -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()
|
||||||
Reference in New Issue
Block a user