116 lines
3.1 KiB
Python
116 lines
3.1 KiB
Python
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()
|