Add main.py
This commit is contained in:
@@ -0,0 +1,127 @@
|
||||
"""Convert an informal assignment description into a validated flat card."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from dotenv import find_dotenv, 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
|
||||
|
||||
|
||||
DEFAULT_DESCRIPTION = (
|
||||
"Сдайте к пятнице мини-отчёт по LangChain: 2 страницы, упор на агентов. "
|
||||
"Оценка: за полноту и за пример кода."
|
||||
)
|
||||
|
||||
|
||||
class AssignmentCard(BaseModel):
|
||||
"""Flat data card extracted from one assignment description."""
|
||||
|
||||
title: str = Field(..., description="Short title of the assignment")
|
||||
subject: str = Field(..., description="Main topic or subject")
|
||||
deadline_hint: str | None = Field(
|
||||
None,
|
||||
description="Deadline as a free-form phrase from the source text",
|
||||
)
|
||||
deliverable_type: str | None = Field(
|
||||
None,
|
||||
description="Expected deliverable: report, code, presentation, etc.",
|
||||
)
|
||||
grading_hints: list[str] = Field(
|
||||
default_factory=list,
|
||||
description="Assessment hints mentioned in the description",
|
||||
)
|
||||
|
||||
|
||||
parser = PydanticOutputParser(pydantic_object=AssignmentCard)
|
||||
|
||||
prompt = PromptTemplate(
|
||||
template=(
|
||||
"You convert one informal study assignment description into a flat "
|
||||
"validated data card. Do not ask follow-up questions. Extract only "
|
||||
"information supported by the input. If a field is not mentioned, use "
|
||||
"null for optional strings or an empty list for grading_hints.\n\n"
|
||||
"{format_instructions}\n\n"
|
||||
"Assignment description:\n{description}"
|
||||
),
|
||||
input_variables=["description"],
|
||||
partial_variables={"format_instructions": parser.get_format_instructions()},
|
||||
)
|
||||
|
||||
|
||||
def build_llm() -> ChatOpenAI:
|
||||
"""Create an OpenAI-compatible chat model from environment variables."""
|
||||
|
||||
load_dotenv(find_dotenv(usecwd=True))
|
||||
base_url = os.getenv("OPENAI_BASE_URL") or os.getenv("GPT2GIGA_BASE_URL")
|
||||
api_key = os.getenv("OPENAI_API_KEY") or os.getenv("GPT2GIGA_OPENAI_API_KEY")
|
||||
model = (
|
||||
os.getenv("OPENAI_MODEL")
|
||||
or os.getenv("GPT2GIGA_MODEL")
|
||||
or "openai/gpt-oss-20b"
|
||||
)
|
||||
if not base_url:
|
||||
raise RuntimeError("Set OPENAI_BASE_URL or GPT2GIGA_BASE_URL")
|
||||
if not api_key:
|
||||
raise RuntimeError("Set OPENAI_API_KEY or GPT2GIGA_OPENAI_API_KEY")
|
||||
|
||||
return ChatOpenAI(
|
||||
model=model,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
|
||||
def parse_assignment(description: str) -> AssignmentCard:
|
||||
"""Run one prompt-model-parser chain and return a validated card."""
|
||||
|
||||
chain = prompt | build_llm() | parser
|
||||
return chain.invoke({"description": description})
|
||||
|
||||
|
||||
def format_summary(card: AssignmentCard) -> str:
|
||||
"""Create a short human-readable summary from the validated card."""
|
||||
|
||||
grading = ", ".join(card.grading_hints) if card.grading_hints else "не указано"
|
||||
return (
|
||||
f"{card.title}: {card.deliverable_type or 'тип сдачи не указан'} "
|
||||
f"по теме «{card.subject}». "
|
||||
f"Срок: {card.deadline_hint or 'не указан'}. "
|
||||
f"Критерии: {grading}."
|
||||
)
|
||||
|
||||
|
||||
def print_card(card: AssignmentCard) -> None:
|
||||
"""Print both machine-friendly and human-friendly output."""
|
||||
|
||||
print("Validated object:")
|
||||
for key, value in card.model_dump().items():
|
||||
print(f"{key}: {value}")
|
||||
print("\nSummary:")
|
||||
print(format_summary(card))
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> AssignmentCard:
|
||||
arg_parser = argparse.ArgumentParser(
|
||||
description="Parse one raw assignment description into a flat card."
|
||||
)
|
||||
arg_parser.add_argument(
|
||||
"description",
|
||||
nargs="*",
|
||||
help="One informal assignment description. Uses a sample if omitted.",
|
||||
)
|
||||
args = arg_parser.parse_args(argv)
|
||||
|
||||
description = " ".join(args.description).strip() or DEFAULT_DESCRIPTION
|
||||
card = parse_assignment(description)
|
||||
print_card(card)
|
||||
return card
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user