144 lines
5.5 KiB
Python
144 lines
5.5 KiB
Python
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() |