add main.py
This commit is contained in:
@@ -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())
|
||||
Reference in New Issue
Block a user