Add agent.py
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user