from typing import Literal from langgraph.graph import StateGraph, END from langchain_openai import ChatOpenAI from state import ConversationState from agents import ( RouterAgent, GeneralDesignAgent, BudgetAnalysisAgent, FloorplanAgent, FloorplanGeneratorAgent, RegulationAgent ) from detailed_budget_agent import DetailedBudgetAgent # Removed creative_specialists - they had NoneType formatting errors # Keeping only the core working agents def create_initial_state() -> ConversationState: """Create initial conversation state with enhanced multi-agent memory""" return { "messages": [], "current_topic": None, "user_requirements": { "budget": None, "location": "Montreal", # Default as specified in requirements "family_size": None, "lifestyle_preferences": [], "special_needs": [] }, "floorplan_requirements": { "num_floors": None, "total_sqft": None, "lot_shape": None, "lot_dimensions": None, "rooms": [] }, "detailed_floorplan": { "design_analysis": None, "detailed_rooms": [], "structural_elements": [], "circulation_plan": {}, "lot_utilization": {}, "architectural_features": [] }, "budget_breakdown": { "site_preparation": None, "foundation": None, "framing": None, "roofing": None, "exterior_finishes": None, "interior_finishes": None, "mechanical_systems": None, "electrical_systems": None, "plumbing_systems": None, "permits_fees": None, "professional_services": None, "contingency": None, "total_construction_cost": None, "cost_per_sqft": None, "budget_analysis": None }, "conversation_history": [], "agent_recommendations": [], "agent_memory": {}, # Shared memory between agents "next_agent": None, "floorplan_ready": False, "budget_ready": False } def should_generate_floorplan(state: ConversationState) -> Literal["generate_floorplan", "route_to_agent"]: """Decide whether to generate floorplan or continue conversation""" if state["floorplan_ready"]: return "generate_floorplan" return "route_to_agent" def should_generate_budget(state: ConversationState) -> Literal["generate_budget", "route_to_agent"]: """Decide whether to generate detailed budget or continue conversation""" # Check if we have enough info for detailed budget and floorplan is designed has_budget = state["user_requirements"]["budget"] is not None has_floorplan_details = state["detailed_floorplan"]["detailed_rooms"] != [] if has_budget and has_floorplan_details and not state["budget_ready"]: return "generate_budget" return "route_to_agent" def route_to_specialist(state: ConversationState) -> Literal[ "general_design", "budget_analysis", "floorplan", "detailed_budget", "regulation", "end" ]: """Route to appropriate specialist agent - core working agents plus regulation""" next_agent = state.get("next_agent") # Check if we should do detailed budget after floorplan if state["detailed_floorplan"]["detailed_rooms"] and not state["budget_ready"]: return "detailed_budget" # Route to specialist based on agent decision - core working agents routing_map = { "general": "general_design", "budget": "budget_analysis", "floorplan": "floorplan", "regulation": "regulation" } return routing_map.get(next_agent, "end") def create_architecture_assistant_graph(model: ChatOpenAI) -> StateGraph: """Create the LangGraph workflow for the architecture assistant with multi-agent collaboration""" # Initialize core working agents plus regulation router = RouterAgent(model) general_agent = GeneralDesignAgent(model) budget_agent = BudgetAnalysisAgent(model) floorplan_agent = FloorplanAgent(model) floorplan_generator = FloorplanGeneratorAgent(model) detailed_budget_agent = DetailedBudgetAgent(model) regulation_agent = RegulationAgent(model) # Create the graph workflow = StateGraph(ConversationState) # Add core working nodes plus regulation workflow.add_node("router", router.process) workflow.add_node("general_design", general_agent.process) workflow.add_node("budget_analysis", budget_agent.process) workflow.add_node("floorplan", floorplan_agent.process) workflow.add_node("generate_floorplan", floorplan_generator.process) workflow.add_node("detailed_budget", detailed_budget_agent.process) workflow.add_node("regulation", regulation_agent.process) # Set entry point workflow.set_entry_point("router") # Add conditional routing from router to core working specialists plus regulation workflow.add_conditional_edges( "router", route_to_specialist, { "general_design": "general_design", "budget_analysis": "budget_analysis", "floorplan": "floorplan", "detailed_budget": "detailed_budget", "regulation": "regulation", "end": END } ) # Add conditional routing from floorplan agent to check if ready for generation workflow.add_conditional_edges( "floorplan", should_generate_floorplan, { "generate_floorplan": "generate_floorplan", "route_to_agent": END } ) # After floorplan generation, automatically trigger detailed budget if budget available workflow.add_conditional_edges( "generate_floorplan", should_generate_budget, { "generate_budget": "detailed_budget", "route_to_agent": END } ) # All core agents end conversation after responding (user can continue with new input) workflow.add_edge("general_design", END) workflow.add_edge("budget_analysis", END) workflow.add_edge("detailed_budget", END) workflow.add_edge("regulation", END) return workflow.compile() class ArchitectureAssistant: """Main architecture assistant class with persistent state management""" def __init__(self, openai_api_key: str, user_id: str = None): self.model = ChatOpenAI( api_key=openai_api_key, model="gpt-4o-mini", temperature=0.7 ) self.graph = create_architecture_assistant_graph(self.model) self.state = create_initial_state() # Initialize user state management from user_state_manager import user_state_manager self.state_manager = user_state_manager self.user_id = user_id self.session_id = None # Start new session if user_id: self.session_id = self.state_manager.start_new_session(user_id) def chat(self, user_input: str, save_state: bool = True) -> str: """Process user input and return response with optional state saving""" # Add user message to state self.state["messages"].append({ "role": "user", "content": user_input }) # Process through the graph result = self.graph.invoke(self.state) # Update state with result self.state = result # Save state after each interaction if save_state and self.user_id: self.state_manager.save_user_state( self.state, self.user_id, self.session_id ) # Return the last assistant message assistant_messages = [msg for msg in self.state["messages"] if msg["role"] == "assistant"] if assistant_messages: return assistant_messages[-1]["content"] else: return "I'm here to help with your home design questions!" def get_conversation_summary(self) -> dict: """Get a summary of the current conversation state""" return { "user_requirements": self.state["user_requirements"], "floorplan_requirements": self.state["floorplan_requirements"], "current_topic": self.state["current_topic"], "total_messages": len(self.state["messages"]) } def reset_conversation(self, start_new_session: bool = True): """Reset the conversation state""" self.state = create_initial_state() # Start new session if requested if start_new_session and self.user_id: self.session_id = self.state_manager.start_new_session(self.user_id) def load_previous_state(self, session_id: str = None) -> bool: """Load a previous conversation state""" if not self.user_id: return False loaded_state = self.state_manager.load_user_state(self.user_id, session_id) if loaded_state: self.state = loaded_state if session_id: self.session_id = session_id return True return False def get_user_history(self) -> list: """Get conversation history for current user""" if not self.user_id: return [] return self.state_manager.get_user_history(self.user_id) def set_user_id(self, user_id: str): """Set or change the user ID""" self.user_id = user_id self.session_id = self.state_manager.start_new_session(user_id)