From 2f0dcc51a4bfc2bfc17e9b9a4bab995d58634629 Mon Sep 17 00:00:00 2001 From: balabanovan530 <175+balabanovan530@noreply.localhost> Date: Thu, 28 May 2026 05:39:33 +0000 Subject: [PATCH] Add agent.py --- agent.py | 115 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 115 insertions(+) create mode 100644 agent.py diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..993ae7a --- /dev/null +++ b/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()