Add main.py
This commit is contained in:
@@ -0,0 +1,36 @@
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_openai import ChatOpenAI
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List
|
||||
|
||||
class TaskCard(BaseModel):
|
||||
title: str = Field(..., description="Task title")
|
||||
subject: str | None = Field(None, description="Subject or topic of the task")
|
||||
deadline_hint: str | None = Field(None, description="Short hint about deadline")
|
||||
deliverable_type: str | None = Field(None, description="Type of deliverable (e.g., code, report)")
|
||||
grading_hints: List[str] | None = Field(None, description="Hints for grading")
|
||||
|
||||
prompt_template = ChatPromptTemplate.from_messages([
|
||||
(
|
||||
"system",
|
||||
"You are an assistant that extracts structured task information from a raw text. Return JSON matching the TaskCard schema.",
|
||||
),
|
||||
("human", "{raw_text}"),
|
||||
])
|
||||
|
||||
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
|
||||
|
||||
async def parse_task(raw_text: str) -> TaskCard:
|
||||
chain = prompt_template | llm.with_structured_output(TaskCard)
|
||||
result = await chain.ainvoke({"raw_text": raw_text})
|
||||
return result
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys, json
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python main.py <path_to_raw_task_file>")
|
||||
sys.exit(1)
|
||||
with open(sys.argv[1], "r", encoding="utf-8") as f:
|
||||
raw = f.read()
|
||||
card = asyncio.run(parse_task(raw))
|
||||
print(json.dumps(card.model_dump(), indent=2))
|
||||
Reference in New Issue
Block a user