Update app.py

app.py

@@ -139,6 +139,15 @@ from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
 from langchain_huggingface import HuggingFaceEndpoint
 from langgraph.graph import StateGraph
+from typing import TypedDict
+
+class InputState(TypedDict):
+    string_var :str
+    numeric_var :int 
+    
+def changeState(input: InputState):
+    print(f"Current value: {input}")
+    return input
 
 # Define the LLM models
 llm1 = HuggingFaceEndpoint(model='t5-small')
@@ -161,45 +170,35 @@ chain2 = LLMChain(llm=llm2, prompt=PromptTemplate(
     template="You are in state s2. {{query}}"
 ))
 
-# 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)
+graph = StateGraph(InputState)
 
 # Add states to the graph
-state1 = graph.add_state(name="s1")
-state2 = graph.add_state(name="s2")
+graph.add_node("s1",changeState)
+graph.add_node("s2",changeState)
 
 # Define transitions
-graph.add_edge(state1, state2, label="next")  # Transition from s1 to s2
-graph.add_edge(state2, state1, label="back")   # Transition from s2 to s1
+graph.add_edge(START, "s1")  # Transition from s1 to s2
+graph.add_edge("s1", "s2")   # Transition from s2 to s1
+graph.add_edge("s2", END)
 
 # Initialize the current state
-current_state = state1
+current_state = "s1"
 
 def handle_input(query):
     global current_state
     output = ''
     
     # Process user input based on current state
-    if current_state.name == "s1":
+    if current_state == "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.name == "s2":
+        current_state = "s2"  # Transition to state s2
+    elif current_state == "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
+        current_state = "s1"  # Transition back to state s1
 
     return response