Update app.py

app.py

@@ -7,16 +7,16 @@ from typing import TypedDict, Annotated, Sequence
 from langchain_openai import ChatOpenAI
 from langchain_core.tools import tool
 from langchain_core.utils.function_calling import convert_to_openai_tool
-from langgraph.graph import StateGraph, END
-from langchain_core.messages import BaseMessage
+from langgraph.graph import StateGraph, END, START
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
 
-
-# Environment Setup
+# ------------------- Environment Setup -------------------
 os.environ['OPENAI_API_KEY'] = os.getenv("OPENAI_API_KEY")
 
-# Model Initialization
+# ------------------- Model Initialization -------------------
 model = ChatOpenAI(temperature=0)
 
+# ------------------- Define the Tool -------------------
 @tool
 def multiply(first_number: int, second_number: int):
     """Multiplies two numbers together and returns the result."""
@@ -24,48 +24,43 @@ def multiply(first_number: int, second_number: int):
 
 model_with_tools = model.bind(tools=[convert_to_openai_tool(multiply)])
 
-# State Setup
-class AgentState(TypedDict):
-    messages: Annotated[Sequence, operator.add]
-
-graph = StateGraph(AgentState)
-
-# Define MessagesState
+# ------------------- State Setup -------------------
 class MessagesState(TypedDict):
     messages: Annotated[Sequence[BaseMessage], operator.add]
 
-def invoke_model(state):
-    question = state['messages'][-1]
-    return {"messages": [model_with_tools.invoke(question)]}
+# ------------------- Assistant Node -------------------
+def assistant(state: MessagesState):
+    """Invoke the model to process messages."""
+    messages = state['messages']
+    response = model_with_tools.invoke(messages[-1])
+    return {"messages": messages + [response]}
 
-graph.add_node("agent", invoke_model)
-
-def invoke_tool(state):
+# ------------------- Tools Node -------------------
+def tools(state: MessagesState):
+    """Invoke tools based on tool calls."""
     tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
-    for tool_call in tool_calls:
-        if tool_call.get("function").get("name") == "multiply":
-            res = multiply.invoke(json.loads(tool_call.get("function").get("arguments")))
-            return {"messages": [f"Tool Result: {res}"]}
-    return {"messages": ["No tool input provided."]}
+    responses = []
 
-graph.add_node("tool", invoke_tool)
-graph.add_edge("tool", END)
-graph.set_entry_point("agent")
-
-# Define the router function
-def router(state):
+    for tool_call in tool_calls:
+        if tool_call["function"]["name"] == "multiply":
+            args = json.loads(tool_call["function"]["arguments"])
+            result = multiply.invoke(args)
+            responses.append(
+                AIMessage(content=f"Tool Result: {result}", name="multiply")
+            )
+    return {"messages": state["messages"] + responses}
+
+# ------------------- Router Logic -------------------
+def router(state: MessagesState):
     tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
-    if len(tool_calls):
-        return "tools"
-    else:
-        return "end"
+    return "tools" if tool_calls else "end"
 
-# Add graph nodes and edges
+# ------------------- Graph Definition -------------------
 app_graph = StateGraph(MessagesState)
 app_graph.add_node("assistant", assistant)
-app_graph.add_node("tools", ToolNode(tools))
+app_graph.add_node("tools", tools)
 
-# Add conditional routing logic
+# Define edges and conditional routing
 app_graph.add_edge(START, "assistant")
 app_graph.add_conditional_edges("assistant", router, {
     "tools": "tools",
@@ -76,18 +71,19 @@ app_graph.add_edge("tools", "assistant")
 # Compile the graph
 react_graph = app_graph.compile()
 
-# Save graph visualization as an image
+# ------------------- Save Graph Visualization -------------------
 with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
     graph_viz = app_graph.get_graph(xray=True)
     tmpfile.write(graph_viz.draw_mermaid_png())  # Write binary image data to file
     graph_image_path = tmpfile.name
 
-# Streamlit Interface
+# ------------------- Streamlit Interface -------------------
 st.title("Simple Tool Calling Demo")
 
 # Display the workflow graph
 st.image(graph_image_path, caption="Workflow Visualization")
 
+# ------------------- Tab 1: Multiplication -------------------
 tab1, tab2 = st.tabs(["Try Multiplication", "Ask General Queries"])
 
 with tab1:
@@ -100,25 +96,32 @@ with tab1:
         second_number = st.number_input("Second Number", value=0, step=1)
 
     if st.button("Multiply"):
-        question = f"What is {first_number} * {second_number}?"
-        output = app_graph.invoke({"messages": [question]})
-        st.success(output['messages'][-1])
-
+        question = HumanMessage(content=f"What is {first_number} * {second_number}?")
+        try:
+            output = react_graph.invoke({"messages": [question]})
+            st.success(output['messages'][-1].content)
+        except Exception as e:
+            st.error(f"Error: {e}")
+
+# ------------------- Tab 2: General Queries -------------------
 with tab2:
     st.subheader("General Query")
     user_input = st.text_input("Enter your question here")
 
     if st.button("Submit"):
         if user_input:
+            question = HumanMessage(content=user_input)
             try:
-                result = app_graph.invoke({"messages": [user_input]})
+                result = react_graph.invoke({"messages": [question]})
                 st.write("Response:")
-                st.success(result['messages'][-1])
+                st.success(result['messages'][-1].content)
             except Exception as e:
-                st.error("Something went wrong. Try again!")
+                st.error(f"Error: {e}")
         else:
             st.warning("Please enter a valid input.")
 
-
+# ------------------- Sidebar Reference -------------------
 st.sidebar.title("Reference")
-st.sidebar.markdown("1. [LangGraph Tool Calling](https://github.com/aritrasen87/LLM_RAG_Model_Deployment/blob/main/LangGraph_02_ToolCalling.ipynb)")
+st.sidebar.markdown(
+    "1. [LangGraph Tool Calling](https://github.com/aritrasen87/LLM_RAG_Model_Deployment/blob/main/LangGraph_02_ToolCalling.ipynb)"
+)