add main.py

2026-05-28 18:09:52 +00:00
parent d27dfb3e03
commit 35c2f70e26
1 changed files with 205 additions and 0 deletions
@@ -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()