Add agent.py

agent.py

@@ -0,0 +1,115 @@
+from typing import TypedDict, Optional
+
+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
+
+
+class TaskCard(BaseModel):
+    title: str
+    subject: str
+    deadline_hint: str
+    deliverable_type: str
+    grading_hints: str
+
+
+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 = PromptTemplate.from_template(PROMPT_TEMPLATE)
+
+# LLM configuration
+llm = ChatOpenAI(model="gpt-4o", temperature=0)
+
+# Parser to enforce TaskCard schema
+parser = PydanticOutputParser(pydantic_object=TaskCard)
+
+# Chain: prompt -> LLM -> parser
+from langchain_core.runnables import RunnableParallel, 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:
+    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
+    except Exception as e:
+        state["error"] = str(e)
+    return state
+
+
+# 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 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"])
+
+
+if __name__ == "__main__":
+    main()