Update app.py

app.py

@@ -161,16 +161,28 @@ chain2 = LLMChain(llm=llm2, prompt=PromptTemplate(
     template="You are in state s2. {{query}}"
 ))
 
-# Create a state graph for managing the chatbot's states
-graph = StateGraph()
-
-# Create states and add them to the graph
-state1 = graph.add_state("s1")  # State for the first agent
-state2 = graph.add_state("s2")  # State for the second agent
+# Define the state schema
+state_schema = {
+    "s1": {
+        "inputs": ["query"],
+        "outputs": ["response"]
+    },
+    "s2": {
+        "inputs": ["query"],
+        "outputs": ["response"]
+    }
+}
+
+# Create a state graph with required schemas for inputs and outputs
+graph = StateGraph(state_schema=state_schema)
+
+# Add states to the graph
+state1 = graph.add_state(name="s1")
+state2 = graph.add_state(name="s2")
 
 # Define transitions
-graph.add_edge(state1, state2, "next")  # Transition from s1 to s2
-graph.add_edge(state2, state1, "back")   # Transition from s2 to s1
+graph.add_edge(state1, state2, label="next")  # Transition from s1 to s2
+graph.add_edge(state2, state1, label="back")   # Transition from s2 to s1
 
 # Initialize the current state
 current_state = state1
@@ -180,11 +192,11 @@ def handle_input(query):
     output = ''
     
     # Process user input based on current state
-    if current_state == state1:
+    if current_state.name == "s1":
         output = chain1.invoke(input=query)  # Invoke chain1 with user input
         response = agent1(output)  # Process output through Agent 1
         current_state = state2  # Transition to state s2
-    elif current_state == state2:
+    elif current_state.name == "s2":
         output = chain2.invoke(input=query)  # Invoke chain2 with user input
         response = agent2(output)  # Process output through Agent 2
         current_state = state1  # Transition back to state s1
@@ -207,3 +219,4 @@ with gr.Blocks() as demo:
 
 # Launch the Gradio application
 demo.launch()
+