Add agent.py

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, Field
-from pydantic import BaseModel
 
+# 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
+    title: str = Field(..., description="Title of the task")
-    subject: str
+    subject: str = Field(..., description="Subject of the task")
-    deadline_hint: str
+    deadline_hint: str = Field(..., description="Deadline hint for the task")
-    deliverable_type: str
+    deliverable_type: str = Field(..., description="Type of deliverable")
-    grading_hints: str
+    grading_hints: str = Field(..., description="Hints for grading")
 
-
+# Prompt template instructing LLM to output JSON matching TaskCard schema
-class State(TypedDict):
+PROMPT_TEMPLATE = (
-    raw_text: str
+    "You are a data extraction assistant. Extract the following fields from the text provided:\n"
-    task_card: Optional[TaskCard]
+    "- title\n"
-    summary: Optional[str]
+    "- subject\n"
-    error: Optional[str]
+    "- deadline_hint\n"
-
+    "- deliverable_type\n"
-
+    "- grading_hints\n\n"
-# Prompt that instructs the LLM to extract fields in JSON
+    "Return ONLY a JSON object with these keys and no additional explanation.\n\n"
-PROMPT_TEMPLATE = """
+    "Text:\n{raw_text}\n"
-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 = PromptTemplate.from_template(PROMPT_TEMPLATE)
 
-# LLM configuration
+# LLM configuration – use the API key loaded from environment
-llm = ChatOpenAI(model="gpt-4o", temperature=0)
+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
+# Chain: prompt -> llm -> parser
-from langchain_core.runnables import RunnableParallel, RunnablePassthrough
+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
+def parse_task(raw_text: str) -> TaskCard:
-async def process_task(state: State) -> State:
+    """Parse raw text into a validated TaskCard using the LLM chain."""
     try:
-        task_card = await chain.ainvoke({"raw_text": state["raw_text"]})
+        task_card = chain.invoke({"raw_text": raw_text})
-        summary = (
+        return task_card
-            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
     except Exception as e:
-        state["error"] = str(e)
+        raise RuntimeError(f"Failed to parse task: {e}") from e
-    return state
 
 
-# Node that checks for interrupt keyword
+def summarize_task(task_card: TaskCard) -> str:
-async def check_interrupt(state: State) -> State:
+    """Return a human‑readable summary of the task."""
-    if "interrupt" in state["raw_text"].lower():
+    return (
-        state["error"] = "Interrupt requested. Exiting."
+        f"Task '{task_card.title}' on subject '{task_card.subject}' has a deadline of '{task_card.deadline_hint}'. "
-    return state
+        f"Deliverable type: '{task_card.deliverable_type}'. Grading hints: '{task_card.grading_hints}'."
-
+    )
-
-# 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 main():
     raw_text = input("Enter the raw text to parse: ")
-    initial_state: State = {
+    try:
-        "raw_text": raw_text,
+        task_card = parse_task(raw_text)
-        "task_card": None,
+    except RuntimeError as e:
-        "summary": None,
+        print("Error:", e)
-        "error": None,
+        return
-    }
+    summary = summarize_task(task_card)
-    final_state = app.invoke(initial_state)
+    print("Parsed TaskCard:\n", task_card.json(indent=2))
-
+    print("\nSummary:\n", summary)
-    if final_state["error"]:
-        print("Error:", final_state["error"])
-    else:
-        print("Parsed TaskCard:", final_state["task_card"])
-        print("Summary:", final_state["summary"])
-
 
 if __name__ == "__main__":
     main()