task-6a1865008a94f887e50d471c/main.py

"""Структурированный вывод: PersonInfo / MeetingNotes через Pydantic + LangChain."""
from __future__ import annotations

import os
import sys
from typing import Literal

from dotenv import load_dotenv
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

load_dotenv()

BROJS_INFERENCE_URL = "https://platform.brojs.ru/jrnl-bh/api/inference/v1"
DEFAULT_MODEL = "openai/gpt-oss-20b:free"

PERSON_EXAMPLE = (
    "Анна, 28 лет, Python-разработчик. Навыки: FastAPI, Docker."
)
MEETING_EXAMPLE = (
    "Встреча 15.03.2026. Участники: Иван, Мария, Алексей. "
    "Обсудили релиз MVP и тестирование. Решили: завершить API до пятницы. "
    "Следующие шаги: написать интеграционные тесты, обновить документацию."
)


class PersonInfo(BaseModel):
    name: str = Field(description="Полное имя человека")
    age: int | None = Field(default=None, description="Возраст в годах, если указан")
    profession: str = Field(description="Профессия или роль")
    skills: list[str] = Field(description="Список навыков и технологий")


class MeetingNotes(BaseModel):
    date: str = Field(description="Дата встречи")
    participants: list[str] = Field(description="Список участников")
    topics: list[str] = Field(description="Обсуждённые темы")
    decisions: list[str] = Field(description="Принятые решения")
    next_steps: list[str] = Field(description="Следующие шаги после встречи")


SchemaType = Literal["person", "meeting"]


def _api_key() -> str:
    return (
        os.getenv("JOURNAL_MCP_PAT")
        or os.getenv("JOURNAL_TOKEN")
        or os.getenv("OPENAI_API_KEY")
        or ""
    )


def _base_url() -> str:
    if os.getenv("OPENAI_BASE_URL"):
        return os.environ["OPENAI_BASE_URL"]
    if os.getenv("OPENAI_API_KEY") and not os.getenv("JOURNAL_MCP_PAT"):
        return os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
    return BROJS_INFERENCE_URL


def build_llm() -> ChatOpenAI:
    return ChatOpenAI(
        model=os.getenv("OPENAI_MODEL", DEFAULT_MODEL),
        base_url=_base_url(),
        api_key=_api_key(),
        temperature=0.0,
    )


def detect_schema_type(text: str, llm: ChatOpenAI | None = None) -> SchemaType:
    """Эвристика + LLM: person vs meeting."""
    lower = text.lower()
    meeting_markers = ("встреч", "участник", "обсудил", "решил", "следующ", "meeting", "agenda")
    person_markers = ("лет", "навык", "разработчик", "професс", "skill", "age")

    meeting_score = sum(1 for m in meeting_markers if m in lower)
    person_score = sum(1 for m in person_markers if m in lower)

    if meeting_score > person_score:
        return "meeting"
    if person_score > meeting_score:
        return "person"

    llm = llm or build_llm()
    raw = llm.invoke(
        "Определи тип текста: person или meeting. "
        "Ответь одним словом.\n\nТекст:\n" + text
    ).content
    answer = str(raw).strip().lower()
    if "meeting" in answer:
        return "meeting"
    return "person"


def build_extraction_chain(schema: SchemaType, llm: ChatOpenAI | None = None):
    llm = llm or build_llm()
    model_cls = PersonInfo if schema == "person" else MeetingNotes
    parser = PydanticOutputParser(pydantic_object=model_cls)
    prompt = PromptTemplate(
        template=(
            "Извлеки структурированные данные из текста.\n\n"
            "{format_instructions}\n\n"
            "Текст:\n{text}\n"
        ),
        input_variables=["text"],
        partial_variables={"format_instructions": parser.get_format_instructions()},
    )
    return prompt | llm | parser, model_cls


def extract_structured(text: str, llm: ChatOpenAI | None = None) -> BaseModel:
    llm = llm or build_llm()
    schema = detect_schema_type(text, llm)
    chain, _ = build_extraction_chain(schema, llm)
    return chain.invoke({"text": text})


def print_result(label: str, obj: BaseModel) -> None:
    print(f"\n=== {label} ===")
    print(obj.model_dump())
    if isinstance(obj, PersonInfo):
        print(
            f"Сводка: {obj.name}, {obj.age or '?'} лет, {obj.profession}; "
            f"навыки: {', '.join(obj.skills)}"
        )
    elif isinstance(obj, MeetingNotes):
        print(
            f"Сводка: встреча {obj.date}; участники: {', '.join(obj.participants)}; "
            f"шагов: {len(obj.next_steps)}"
        )


def run_demo() -> None:
    llm = build_llm()
    for label, sample in (
        ("Person", PERSON_EXAMPLE),
        ("Meeting", MEETING_EXAMPLE),
    ):
        schema = detect_schema_type(sample, llm)
        chain, _ = build_extraction_chain(schema, llm)
        result = chain.invoke({"text": sample})
        print_result(label, result)


def main() -> int:
    if len(sys.argv) > 1:
        text = " ".join(sys.argv[1:]).strip()
        result = extract_structured(text)
        print_result("Custom input", result)
    else:
        run_demo()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())