diff --git a/agent.py b/agent.py index 993ae7a..0e90b25 100644 --- a/agent.py +++ b/agent.py @@ -1,115 +1,79 @@ -from typing import TypedDict, Optional +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 langgraph.graph import StateGraph -from pydantic import BaseModel +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 - subject: str - deadline_hint: str - deliverable_type: str - grading_hints: str + 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") - -class State(TypedDict): - raw_text: str - task_card: Optional[TaskCard] - summary: Optional[str] - error: Optional[str] - - -# Prompt that instructs the LLM to extract fields in JSON -PROMPT_TEMPLATE = """ -You are a data extraction assistant. Extract the following fields from the text provided: -- title -- subject -- deadline_hint -- deliverable_type -- grading_hints - -Return ONLY a JSON object with these keys and no additional explanation. - -Text: -{raw_text} -""" +# 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 -llm = ChatOpenAI(model="gpt-4o", temperature=0) +# 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 RunnableParallel, RunnablePassthrough +# Chain: prompt -> llm -> parser +from langchain_core.runnables import RunnablePassthrough chain = ( RunnablePassthrough.assign(raw_text=lambda x: x["raw_text"]) | prompt ) | llm | parser -# Node that processes the raw text -async def process_task(state: State) -> State: +def parse_task(raw_text: str) -> TaskCard: + """Parse raw text into a validated TaskCard using the LLM chain.""" try: - task_card = await chain.ainvoke({"raw_text": state["raw_text"]}) - summary = ( - 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}'." - ) - state["task_card"] = task_card - state["summary"] = summary + task_card = chain.invoke({"raw_text": raw_text}) + return task_card except Exception as e: - state["error"] = str(e) - return state + raise RuntimeError(f"Failed to parse task: {e}") from e -# Node that checks for interrupt keyword -async def check_interrupt(state: State) -> State: - if "interrupt" in state["raw_text"].lower(): - state["error"] = "Interrupt requested. Exiting." - return state - - -# Build the StateGraph -graph = StateGraph(State) - -# Add nodes -graph.add_node("check_interrupt", check_interrupt) -graph.add_node("process_task", process_task) - -# Define edges: start at check_interrupt, then to process_task unless interrupted -graph.set_entry_point("check_interrupt") - -# Edge logic: after check_interrupt, if error set, go to END -graph.add_edge("check_interrupt", "process_task", condition=lambda s: s.get("error") is None) -graph.add_edge("check_interrupt", "END", condition=lambda s: s.get("error") is not None) -graph.add_edge("process_task", "END") - -# Compile graph -app = graph.compile() +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: ") - initial_state: State = { - "raw_text": raw_text, - "task_card": None, - "summary": None, - "error": None, - } - final_state = app.invoke(initial_state) - - if final_state["error"]: - print("Error:", final_state["error"]) - else: - print("Parsed TaskCard:", final_state["task_card"]) - print("Summary:", final_state["summary"]) - + 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()