206 lines
7.8 KiB
Python
206 lines
7.8 KiB
Python
"""
|
||
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()
|