diff --git a/main.py b/main.py new file mode 100644 index 0000000..457e3ed --- /dev/null +++ b/main.py @@ -0,0 +1,159 @@ +"""Структурированный вывод: 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())