Spaces:
Sleeping
Sleeping
| """ | |
| LangGraph workflow for the Email Assistant multi-agent system. | |
| Orchestrates the flow between different agents using LangGraph. | |
| """ | |
| from typing import Dict, Any, Literal, Annotated | |
| from langgraph.graph import StateGraph, END | |
| from langgraph.graph.message import add_messages | |
| from models.state import EmailState, EmailInput | |
| from agents import ( | |
| InputParsingAgent, | |
| IntentDetectionAgent, | |
| DraftWriterAgent, | |
| ToneStylistAgent, | |
| ReviewAgent, | |
| RouterAgent | |
| ) | |
| class EmailAssistantWorkflow: | |
| """Orchestrates the multi-agent workflow for email generation.""" | |
| def __init__(self): | |
| # Initialize agents | |
| self.input_parser = InputParsingAgent() | |
| self.intent_detector = IntentDetectionAgent() | |
| self.draft_writer = DraftWriterAgent() | |
| self.tone_stylist = ToneStylistAgent() | |
| self.reviewer = ReviewAgent() | |
| self.router = RouterAgent() | |
| # Build the graph | |
| self.graph = self._build_graph() | |
| def _build_graph(self) -> StateGraph: | |
| """Build the LangGraph workflow.""" | |
| # Create the graph with the state schema | |
| workflow = StateGraph(dict) | |
| # Add nodes for each agent | |
| workflow.add_node("parse_input", self._parse_input_node) | |
| workflow.add_node("detect_intent", self._detect_intent_node) | |
| workflow.add_node("write_draft", self._write_draft_node) | |
| workflow.add_node("style_tone", self._style_tone_node) | |
| workflow.add_node("review_draft", self._review_draft_node) | |
| workflow.add_node("handle_retry", self._handle_retry_node) | |
| workflow.add_node("finalize", self._finalize_node) | |
| workflow.add_node("handle_error", self._handle_error_node) | |
| # Set the entry point | |
| workflow.set_entry_point("parse_input") | |
| # Add edges for the main flow | |
| workflow.add_edge("parse_input", "detect_intent") | |
| workflow.add_edge("detect_intent", "write_draft") | |
| workflow.add_edge("write_draft", "style_tone") | |
| workflow.add_edge("style_tone", "review_draft") | |
| # Add conditional edges after review | |
| workflow.add_conditional_edges( | |
| "review_draft", | |
| self._route_after_review, | |
| { | |
| "retry": "handle_retry", | |
| "complete": "finalize", | |
| "error": "handle_error" | |
| } | |
| ) | |
| # Retry goes back to draft writing | |
| workflow.add_edge("handle_retry", "write_draft") | |
| # Terminal nodes | |
| workflow.add_edge("finalize", END) | |
| workflow.add_edge("handle_error", END) | |
| return workflow.compile() | |
| def _parse_input_node(self, state: Dict[str, Any]) -> Dict[str, Any]: | |
| """Node for parsing user input.""" | |
| try: | |
| return self.input_parser.process(state) | |
| except Exception as e: | |
| state["error_message"] = f"Input parsing failed: {str(e)}" | |
| return state | |
| def _detect_intent_node(self, state: Dict[str, Any]) -> Dict[str, Any]: | |
| """Node for detecting email intent.""" | |
| try: | |
| return self.intent_detector.process(state) | |
| except Exception as e: | |
| state["error_message"] = f"Intent detection failed: {str(e)}" | |
| return state | |
| def _write_draft_node(self, state: Dict[str, Any]) -> Dict[str, Any]: | |
| """Node for writing the email draft.""" | |
| try: | |
| return self.draft_writer.process(state) | |
| except Exception as e: | |
| state["error_message"] = f"Draft writing failed: {str(e)}" | |
| return state | |
| def _style_tone_node(self, state: Dict[str, Any]) -> Dict[str, Any]: | |
| """Node for styling the email tone.""" | |
| try: | |
| return self.tone_stylist.process(state) | |
| except Exception as e: | |
| state["error_message"] = f"Tone styling failed: {str(e)}" | |
| return state | |
| def _review_draft_node(self, state: Dict[str, Any]) -> Dict[str, Any]: | |
| """Node for reviewing the email draft.""" | |
| try: | |
| return self.reviewer.process(state) | |
| except Exception as e: | |
| state["error_message"] = f"Review failed: {str(e)}" | |
| return state | |
| def _handle_retry_node(self, state: Dict[str, Any]) -> Dict[str, Any]: | |
| """Node for handling retry logic.""" | |
| return self.router.increment_retry(state) | |
| def _finalize_node(self, state: Dict[str, Any]) -> Dict[str, Any]: | |
| """Node for finalizing the workflow.""" | |
| return self.router.mark_complete(state) | |
| def _handle_error_node(self, state: Dict[str, Any]) -> Dict[str, Any]: | |
| """Node for handling errors.""" | |
| return self.router.handle_error(state) | |
| def _route_after_review(self, state: Dict[str, Any]) -> Literal["retry", "complete", "error"]: | |
| """Routing function after review.""" | |
| return self.router.route_after_review(state) | |
| def run(self, email_input: EmailInput) -> EmailState: | |
| """Run the workflow with the given email input.""" | |
| # Create initial state | |
| initial_state = EmailState(email_input=email_input) | |
| state_dict = initial_state.model_dump() | |
| # Run the graph | |
| final_state = self.graph.invoke(state_dict) | |
| # Return as EmailState | |
| return EmailState(**final_state) | |
| def run_dict(self, email_input: EmailInput) -> Dict[str, Any]: | |
| """Run the workflow and return the result as a dictionary.""" | |
| initial_state = EmailState(email_input=email_input) | |
| state_dict = initial_state.model_dump() | |
| return self.graph.invoke(state_dict) | |
| def create_workflow() -> EmailAssistantWorkflow: | |
| """Factory function to create a new workflow instance.""" | |
| return EmailAssistantWorkflow() | |