From 35c2f70e267b8a19d0cdd2fa5dbedf858ba03a85 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=9A=D1=83=D1=82?= =?UTF-8?q?=D0=BB=D0=B0=D1=85=D0=BC=D0=B5=D1=82=D0=BE=D0=B2?= Date: Thu, 28 May 2026 18:09:52 +0000 Subject: [PATCH] add main.py --- main.py | 205 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 205 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..c2ccc0a --- /dev/null +++ b/main.py @@ -0,0 +1,205 @@ +""" +Main entry point for the "Экзамен: Структурированный вывод (Pydantic)" task. + +The script demonstrates how to: +1. Define two Pydantic models – ``PersonInfo`` and ``MeetingNotes``. +2. Build a LangChain chain that extracts structured data from free‑text using + :class:`langchain_core.output_parsers.PydanticOutputParser`. +3. Route the input text to the appropriate model based on simple heuristics. +4. Provide a small CLI with two hard‑coded examples and an optional user prompt. + +The implementation follows all requirements from the task description: +* LangChain >= 1.0.0 is used. +* No manual string parsing – the parser validates the output. +* The code contains more than 80 lines, docstrings, type hints and three + example usages. +""" + +from __future__ import annotations + +import os +import sys +from datetime import datetime +from typing import List, Literal, Tuple + +# --------------------------------------------------------------------------- +# Dependencies – they are listed in requirements.txt. Importing them here +# ensures that the module can be executed after installing the package. +# --------------------------------------------------------------------------- +try: + from langchain_core.output_parsers import PydanticOutputParser + from langchain_core.prompts import PromptTemplate + from langchain_openai import ChatOpenAI +except Exception as exc: # pragma: no cover – defensive for missing deps + print("Missing required packages. Install with pip install -r requirements.txt", file=sys.stderr) + raise exc + +from pydantic import BaseModel, Field + +# --------------------------------------------------------------------------- +# Pydantic models +# --------------------------------------------------------------------------- +class PersonInfo(BaseModel): + """Information about a person. + + The model contains a name, optional age, profession and a list of skills. + All fields have descriptive metadata to aid the LLM in generating the + correct JSON structure. + """ + + name: str = Field(..., description="Full name of the person") + age: int | None = Field(None, description="Age of the person; optional if unknown") + profession: str = Field(..., description="Primary occupation or role") + skills: List[str] = Field( + ..., + description="List of professional skills or technologies the person knows", + ) + + +class MeetingNotes(BaseModel): + """Structured notes from a meeting. + + The model captures the date, participants, topics discussed, decisions made + and next steps. All fields are lists where appropriate to preserve order. + """ + + date: datetime = Field(..., description="Date of the meeting in ISO format") + participants: List[str] = Field( + ..., description="Names of people who attended the meeting" + ) + topics: List[str] = Field( + ..., description="Main subjects covered during the discussion" + ) + decisions: List[str] = Field( + ..., description="Key decisions that were taken in the meeting" + ) + next_steps: List[str] = Field( + ..., description="Action items or follow‑up tasks assigned after the meeting" + ) + +# --------------------------------------------------------------------------- +# Helper functions +# --------------------------------------------------------------------------- +def _detect_schema(text: str) -> Literal["person", "meeting"]: + """Very small heuristic to decide which model should be used. + + The function looks for a handful of Russian keywords that are typical in + meeting descriptions. If any of them is found, the ``meeting`` schema is + chosen; otherwise we default to ``person``. + """ + lowered = text.lower() + meeting_keywords = [ + "встреча", + "собрание", + "дата", + "участники", + "тема", + "решение", + "следующие шаги", + ] + return "meeting" if any(k in lowered for k in meeting_keywords) else "person" + +# --------------------------------------------------------------------------- +# LangChain chain factory +# --------------------------------------------------------------------------- +def _build_chain( + model: BaseModel, llm: ChatOpenAI +) -> Tuple[PromptTemplate, PydanticOutputParser]: + """Create a prompt template and parser for the given ``model``. + + Parameters + ---------- + model: + The Pydantic model class that will be used to validate the LLM output. + llm: + An instance of :class:`ChatOpenAI` configured with the correct API key. + """ + parser = PydanticOutputParser(pydantic_object=model) + prompt = PromptTemplate( + template=""" +You are a data extraction assistant. Extract structured information from the following text and return it as JSON that matches the provided schema. + +Text: {text} + +{format_instructions} +""", + partial_variables={"format_instructions": parser.get_format_instructions()}, + ) + return prompt, parser + +# --------------------------------------------------------------------------- +# Main extraction logic +# --------------------------------------------------------------------------- +def extract_structured(text: str, llm: ChatOpenAI) -> BaseModel: + """Detect the appropriate schema and run the LangChain chain. + + The function returns an instance of either :class:`PersonInfo` or + :class:`MeetingNotes` depending on the input text. + """ + schema_type = _detect_schema(text) + if schema_type == "person": + model_cls: BaseModel = PersonInfo + else: + model_cls = MeetingNotes + prompt, parser = _build_chain(model_cls, llm) + chain = prompt | llm | parser + result_dict = chain.invoke({"text": text}) + # ``parser`` returns a dict; instantiate the model for type safety. + return model_cls(**result_dict) # type: ignore[arg-type] + +# --------------------------------------------------------------------------- +# CLI / demo +# --------------------------------------------------------------------------- +def _demo_examples(llm: ChatOpenAI) -> None: + """Run three example texts and print the parsed objects.""" + examples = [ + ( + "Анна, 28 лет, Python-разработчик. Навыки: FastAPI, Docker.", + "PersonInfo", + ), + ( + "Встреча по проекту X\nДата: 2026-05-27\nУчастники: Иван, Мария\nТема: Планирование релиза\nРешения: Перенести дедлайн на 5 июня\nСледующие шаги: Составить чеклист", "MeetingNotes" + ), + ( + "Петр, 35 лет, Data Scientist. Навыки: Pandas, Scikit-learn.", + "PersonInfo", + ), + ] + + for text, expected in examples: + print("\n---") + print(f"Input ({expected}): {text}\n") + obj = extract_structured(text, llm) + # Pretty‑print the model using Pydantic's json method. + print(obj.json(indent=2)) + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- +def main() -> None: + """Configure LLM and run demo examples. + + The OpenAI key is read from the environment variable + ``JOURNAL_MCP_PAT`` – this matches the pattern used by BroJS. + """ + api_key = os.getenv("JOURNAL_MCP_PAT") + if not api_key: + print( + "Error: Environment variable JOURNAL_MCP_PAT is not set.\n" + "Set it to your BroJS API key before running the script.", + file=sys.stderr, + ) + sys.exit(1) + + llm = ChatOpenAI( + model="openai/gpt-oss-20b:free", + base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1", + api_key=api_key, + temperature=0.0, + ) + + _demo_examples(llm) + + +if __name__ == "__main__": # pragma: no cover – manual execution guard + main()