Экзамен: Структурированный вывод (Pydantic): solution.py

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()