Update app.py

app.py

@@ -133,11 +133,12 @@ with gr.Blocks() as demo:
 # Launch the Gradio application
 demo.launch()
 '''
+
 import gradio as gr
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
 from langchain_huggingface import HuggingFaceEndpoint
-from langgraph.graph import StateGraph, END, START
+from langgraph.graph import StateGraph
 
 # Define the LLM models
 llm1 = HuggingFaceEndpoint(model='t5-small')
@@ -151,34 +152,42 @@ def agent2(response):
     return f"Agent 2: {response}"
 
 # Define the prompts and LLM chains
-chain1 = LLMChain(llm=llm1, prompt=PromptTemplate(input_variables=["query"], template="You are in state s1. {{query}}"))
-chain2 = LLMChain(llm=llm2, prompt=PromptTemplate(input_variables=["query"], template="You are in state s2. {{query}}"))
-
-# State definitions
-s1 = StateGraph("s1")
-s2 = StateGraph("s2")
-
-# Create transitions in the states
-s1.add_edge(s2, 'next')  # From state s1 to s2
-s2.add_edge(s1, 'back')   # From state s2 to s1
+chain1 = LLMChain(llm=llm1, prompt=PromptTemplate(
+    input_variables=["query"], 
+    template="You are in state s1. {{query}}"
+))
+chain2 = LLMChain(llm=llm2, prompt=PromptTemplate(
+    input_variables=["query"], 
+    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 transitions
+graph.add_edge(state1, state2, "next")  # Transition from s1 to s2
+graph.add_edge(state2, state1, "back")   # Transition from s2 to s1
 
 # Initialize the current state
-current_state = s1
+current_state = state1
 
 def handle_input(query):
     global current_state
     output = ''
     
-    if current_state == s1:
-        # Use LLM Chain in s1
-        output = chain1.invoke(input=query)  # Invoke chain1 with the user input
+    # Process user input based on current state
+    if current_state == state1:
+        output = chain1.invoke(input=query)  # Invoke chain1 with user input
         response = agent1(output)  # Process output through Agent 1
-        current_state = s2  # Transition to state s2
-    elif current_state == s2:
-        # Use LLM Chain in s2
-        output = chain2.invoke(input=query)  # Invoke chain2 with the user input
+        current_state = state2  # Transition to state s2
+    elif current_state == state2:
+        output = chain2.invoke(input=query)  # Invoke chain2 with user input
         response = agent2(output)  # Process output through Agent 2
-        current_state = s1  # Transition back to state s1
+        current_state = state1  # Transition back to state s1
 
     return response