diff --git a/solution.py b/solution.py new file mode 100644 index 0000000..75534de --- /dev/null +++ b/solution.py @@ -0,0 +1,127 @@ +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() \ No newline at end of file