Spaces:
Sleeping
Sleeping
File size: 2,208 Bytes
1ad4100 567735d 1ad4100 1e95e40 1ad4100 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
from typing import TypedDict, Any, Dict
from langgraph.graph import StateGraph, END
from agents.grammar_agent import GrammarAgent
from agents.style_agent import StyleAgent
from agents.clarity_agent import ClarityAgent
from agents.reviewer_agent import ReviewerAgent
# ---------------------------
# TextDoctor State Definition
# ---------------------------
class TextDoctorState(TypedDict, total=False):
input_text: str
grammar_result: Dict[str, Any]
style_result: Dict[str, Any]
clarity_result: Dict[str, Any]
review_result: Dict[str, Any]
# ---------------------------
# Initialize Agents Once
# ---------------------------
grammar_agent = GrammarAgent()
style_agent = StyleAgent()
clarity_agent = ClarityAgent()
reviewer_agent = ReviewerAgent()
# ---------------------------
# Node Definitions
# ---------------------------
def grammar_node(state: TextDoctorState) -> TextDoctorState:
state["grammar_result"] = grammar_agent.correct(state["input_text"])
return state
def style_node(state: TextDoctorState) -> TextDoctorState:
g = state["grammar_result"]
state["style_result"] = style_agent.stylize(g["corrected"])
return state
def clarity_node(state: TextDoctorState) -> TextDoctorState:
s = state["style_result"]
state["clarity_result"] = clarity_agent.clarify(s["styled"])
return state
def reviewer_node(state: TextDoctorState) -> TextDoctorState:
state["review_result"] = reviewer_agent.review(
original=state["input_text"],
grammar_result=state["grammar_result"],
style_result=state["style_result"],
clarity_result=state["clarity_result"],
)
return state
# ---------------------------
# Build LangGraph Pipeline
# ---------------------------
graph = StateGraph(TextDoctorState)
graph.add_node("grammar", grammar_node)
graph.add_node("style", style_node)
graph.add_node("clarity", clarity_node)
graph.add_node("reviewer", reviewer_node)
graph.set_entry_point("grammar")
graph.add_edge("grammar", "style")
graph.add_edge("style", "clarity")
graph.add_edge("clarity", "reviewer")
graph.add_edge("reviewer", END)
# The compiled pipeline ready for use in UI
textdoctor_app = graph.compile()
|