127 lines
5.3 KiB
Python
127 lines
5.3 KiB
Python
import os
|
|
import argparse
|
|
from datetime import datetime
|
|
from typing import List, Union
|
|
|
|
from langchain_core.prompts import PromptTemplate
|
|
from langchain_core.output_parsers import PydanticOutputParser
|
|
from langchain_core.runnables import Runnable
|
|
from langchain_openai import ChatOpenAI
|
|
from pydantic import BaseModel, Field
|
|
|
|
# ------------------------------------------------------------------
|
|
# 1. Pydantic модели
|
|
# ------------------------------------------------------------------
|
|
class PersonInfo(BaseModel):
|
|
name: str = Field(..., description="Имя человека")
|
|
age: int | None = Field(None, description="Возраст (необязательно)")
|
|
profession: str = Field(..., description="Профессия")
|
|
skills: List[str] = Field(..., description="Список навыков")
|
|
|
|
class MeetingNotes(BaseModel):
|
|
date: datetime = Field(..., description="Дата встречи в формате YYYY-MM-DD")
|
|
participants: List[str] = Field(..., description="Список участников")
|
|
topics: List[str] = Field(..., description="Темы обсуждения")
|
|
decisions: List[str] = Field(..., description="Принятые решения")
|
|
next_steps: List[str] = Field(..., description="Следующие шаги")
|
|
|
|
# ------------------------------------------------------------------
|
|
# 2. Шаблоны и парсеры
|
|
# ------------------------------------------------------------------
|
|
person_prompt = PromptTemplate(
|
|
input_variables=["text"],
|
|
template=(
|
|
"Найди в тексте информацию о человеке и выведи её в формате JSON, "
|
|
"соответствующем схеме PersonInfo.\n"
|
|
"Текст: {text}\n"
|
|
"Формат: {format_instructions}"
|
|
),
|
|
)
|
|
|
|
meeting_prompt = PromptTemplate(
|
|
input_variables=["text"],
|
|
template=(
|
|
"Найди в тексте информацию о встрече и выведи её в формате JSON, "
|
|
"соответствующем схеме MeetingNotes.\n"
|
|
"Текст: {text}\n"
|
|
"Формат: {format_instructions}"
|
|
),
|
|
)
|
|
|
|
person_parser = PydanticOutputParser(pydantic_object=PersonInfo)
|
|
meeting_parser = PydanticOutputParser(pydantic_object=MeetingNotes)
|
|
|
|
person_chain = (
|
|
person_prompt | ChatOpenAI(temperature=0) | person_parser
|
|
)
|
|
meeting_chain = (
|
|
meeting_prompt | ChatOpenAI(temperature=0) | meeting_parser
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# 3. Выбор схемы
|
|
# ------------------------------------------------------------------
|
|
def choose_chain(text: str) -> Runnable:
|
|
"""Определяем, какой тип данных в тексте."""
|
|
lowered = text.lower()
|
|
# простая эвристика: наличие ключевых слов
|
|
if any(word in lowered for word in ("встреча", "meeting", "дата", "participants", "topics")):
|
|
return meeting_chain
|
|
return person_chain
|
|
|
|
# ------------------------------------------------------------------
|
|
# 4. CLI
|
|
# ------------------------------------------------------------------
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Structured output extractor")
|
|
parser.add_argument(
|
|
"--text",
|
|
type=str,
|
|
help="Текст для обработки. Если не задан, используется пример.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if args.text:
|
|
input_text = args.text
|
|
else:
|
|
# пример для PersonInfo
|
|
input_text = (
|
|
"Анна, 28 лет, Python-разработчик. Навыки: FastAPI, Docker."
|
|
)
|
|
# пример для MeetingNotes
|
|
# input_text = (
|
|
# "Дата: 2024-05-01. Участники: Анна, Боб. Темы: проект X, бюджет. "
|
|
# "Решения: утвердить бюджет. Next steps: подготовить презентацию."
|
|
# )
|
|
|
|
chain = choose_chain(input_text)
|
|
result = chain.invoke({"text": input_text})
|
|
|
|
# result уже валидированный объект Pydantic
|
|
print("\nВыполнено. Валидированный объект:")
|
|
print(result.model_dump(indent=2))
|
|
|
|
# Краткая сводка
|
|
if isinstance(result, PersonInfo):
|
|
summary = (
|
|
f"Человек: {result.name}, "
|
|
f"возраст: {result.age if result.age is not None else 'не указан'}, "
|
|
f"профессия: {result.profession}, "
|
|
f"навыки: {', '.join(result.skills)}"
|
|
)
|
|
else:
|
|
summary = (
|
|
f"Встреча: {result.date.date()}, "
|
|
f"участники: {', '.join(result.participants)}, "
|
|
f"темы: {', '.join(result.topics)}, "
|
|
f"решения: {', '.join(result.decisions)}, "
|
|
f"следующие шаги: {', '.join(result.next_steps)}"
|
|
)
|
|
print("\nСводка:")
|
|
print(summary)
|
|
|
|
if __name__ == "__main__":
|
|
# Для работы с OpenAI требуется переменная окружения OPENAI_API_KEY
|
|
if "OPENAI_API_KEY" not in os.environ:
|
|
raise RuntimeError("Необходимо задать переменную окружения OPENAI_API_KEY")
|
|
main() |