80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
from dotenv import load_dotenv
|
||
import os
|
||
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain_core.prompts import PromptTemplate
|
||
from langchain_core.output_parsers import PydanticOutputParser
|
||
from pydantic import BaseModel, Field
|
||
|
||
# Load environment variables
|
||
load_dotenv()
|
||
api_key = os.getenv("OPENAI_API_KEY")
|
||
if not api_key:
|
||
raise RuntimeError("OPENAI_API_KEY is not set in .env file or environment.")
|
||
|
||
class TaskCard(BaseModel):
|
||
title: str = Field(..., description="Title of the task")
|
||
subject: str = Field(..., description="Subject of the task")
|
||
deadline_hint: str = Field(..., description="Deadline hint for the task")
|
||
deliverable_type: str = Field(..., description="Type of deliverable")
|
||
grading_hints: str = Field(..., description="Hints for grading")
|
||
|
||
# Prompt template instructing LLM to output JSON matching TaskCard schema
|
||
PROMPT_TEMPLATE = (
|
||
"You are a data extraction assistant. Extract the following fields from the text provided:\n"
|
||
"- title\n"
|
||
"- subject\n"
|
||
"- deadline_hint\n"
|
||
"- deliverable_type\n"
|
||
"- grading_hints\n\n"
|
||
"Return ONLY a JSON object with these keys and no additional explanation.\n\n"
|
||
"Text:\n{raw_text}\n"
|
||
)
|
||
|
||
prompt = PromptTemplate.from_template(PROMPT_TEMPLATE)
|
||
|
||
# LLM configuration – use the API key loaded from environment
|
||
llm = ChatOpenAI(model="gpt-4o", temperature=0, api_key=api_key)
|
||
|
||
# Parser to enforce TaskCard schema
|
||
parser = PydanticOutputParser(pydantic_object=TaskCard)
|
||
|
||
# Chain: prompt -> llm -> parser
|
||
from langchain_core.runnables import RunnablePassthrough
|
||
|
||
chain = (
|
||
RunnablePassthrough.assign(raw_text=lambda x: x["raw_text"]) | prompt
|
||
) | llm | parser
|
||
|
||
|
||
def parse_task(raw_text: str) -> TaskCard:
|
||
"""Parse raw text into a validated TaskCard using the LLM chain."""
|
||
try:
|
||
task_card = chain.invoke({"raw_text": raw_text})
|
||
return task_card
|
||
except Exception as e:
|
||
raise RuntimeError(f"Failed to parse task: {e}") from e
|
||
|
||
|
||
def summarize_task(task_card: TaskCard) -> str:
|
||
"""Return a human‑readable summary of the task."""
|
||
return (
|
||
f"Task '{task_card.title}' on subject '{task_card.subject}' has a deadline of '{task_card.deadline_hint}'. "
|
||
f"Deliverable type: '{task_card.deliverable_type}'. Grading hints: '{task_card.grading_hints}'."
|
||
)
|
||
|
||
|
||
def main():
|
||
raw_text = input("Enter the raw text to parse: ")
|
||
try:
|
||
task_card = parse_task(raw_text)
|
||
except RuntimeError as e:
|
||
print("Error:", e)
|
||
return
|
||
summary = summarize_task(task_card)
|
||
print("Parsed TaskCard:\n", task_card.json(indent=2))
|
||
print("\nSummary:\n", summary)
|
||
|
||
if __name__ == "__main__":
|
||
main()
|