Update app.py

app.py

@@ -1,101 +1,118 @@
 import os
 import json
 import operator
-import streamlit as st
 import tempfile
 from typing import TypedDict, Annotated, Sequence
+from dotenv import load_dotenv
+
 from langchain_openai import ChatOpenAI
+from langchain_core.messages import BaseMessage
 from langchain_core.tools import tool
 from langchain_core.utils.function_calling import convert_to_openai_tool
 from langgraph.graph import StateGraph, END
 
-# Environment Setup
-os.environ['OPENAI_API_KEY'] = os.getenv("OPENAI_API_KEY")
+import streamlit as st
+
+# ------------------- Streamlit UI Layout -------------------
+#st.set_page_config(page_title="Streamlit LLM Graph", layout="wide")
+st.title("Test App- Tool Calling and Conditional Graph")
+
+# ------------------- Environment Setup -------------------
 
-# Model Initialization
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+if not OPENAI_API_KEY:
+    st.error("OpenAI API Key not found! Please set it in your environment variables.")
+    st.stop()
+
+# ------------------- Model Initialization -------------------
 model = ChatOpenAI(temperature=0)
 
+# ------------------- Tool Definition -------------------
 @tool
 def multiply(first_number: int, second_number: int):
-    """Multiplies two numbers together and returns the result."""
+    """Multiplies two numbers together."""
     return first_number * second_number
 
 model_with_tools = model.bind(tools=[convert_to_openai_tool(multiply)])
 
-# State Setup
+# ------------------- State Graph Setup -------------------
 class AgentState(TypedDict):
-    messages: Annotated[Sequence, operator.add]
+    messages: Annotated[Sequence[BaseMessage], operator.add]
 
 graph = StateGraph(AgentState)
 
 def invoke_model(state):
-    question = state['messages'][-1]
+    messages = state['messages']
+    question = messages[-1]
     return {"messages": [model_with_tools.invoke(question)]}
 
 graph.add_node("agent", invoke_model)
 
 def invoke_tool(state):
     tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
+    multiply_call = None
+
     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."]}
+            multiply_call = tool_call
+
+    if multiply_call is None:
+        raise Exception("No multiply input found.")
+
+    res = multiply.invoke(
+        json.loads(multiply_call.get("function").get("arguments"))
+    )
+    return {"messages": [res]}
 
 graph.add_node("tool", invoke_tool)
 graph.add_edge("tool", END)
 graph.set_entry_point("agent")
 
 def router(state):
-    calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
-    return "multiply" if calls else "end"
-
-graph.add_conditional_edges("agent", router, {"multiply": "tool", "end": END})
-app_graph = graph.compile()
-
-# Save graph visualization as an image
-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
-st.title("Simple Tool Calling Demo")
-
-# Display the workflow graph
-st.image(graph_image_path, caption="Workflow Visualization")
-
-tab1, tab2 = st.tabs(["Try Multiplication", "Ask General Queries"])
-
-with tab1:
-    st.subheader("Try Multiplication")
-    col1, col2 = st.columns(2)
-
-    with col1:
-        first_number = st.number_input("First Number", value=0, step=1)
-    with col2:
-        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])
-
-with tab2:
-    st.subheader("General Query")
-    user_input = st.text_input("Enter your question here")
-
-    if st.button("Submit"):
-        if user_input:
-            try:
-                result = app_graph.invoke({"messages": [user_input]})
-                st.write("Response:")
-                st.success(result['messages'][-1])
-            except Exception as e:
-                st.error("Something went wrong. Try again!")
-        else:
-            st.warning("Please enter a valid input.")
-
-
-st.sidebar.title("References")
-st.sidebar.markdown("1. [LangGraph Tool Calling](https://github.com/aritrasen87/LLM_RAG_Model_Deployment/blob/main/LangGraph_02_ToolCalling.ipynb)")
+    tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
+    if len(tool_calls):
+        return "multiply"
+    else:
+        return "end"
+
+graph.add_conditional_edges("agent", router, {
+    "multiply": "tool",
+    "end": END,
+})
+
+app = graph.compile()
+
+# ------------------- Graph Visualization -------------------
+st.header("Graph Visualization")
+
+try:
+    # Generate graph visualization and save as image
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
+        graph_viz = app.get_graph(xray=True)
+        tmpfile.write(graph_viz.draw_mermaid_png())  # Write binary image data to file
+        graph_image_path = tmpfile.name
+
+    # Display the graph image
+    st.image(graph_image_path, caption="State Graph Visualization", use_column_width=True)
+except Exception as e:
+    st.error(f"Error generating graph visualization: {e}")
+
+# ------------------- Streamlit Interaction -------------------
+st.header("Ask the LLM")
+
+# User Input
+user_query = st.text_input("Enter your query:", placeholder="E.g., What is 123 * 456?")
+submit_button = st.button("Submit")
+
+# Process Query
+if submit_button and user_query:
+    try:
+        # Invoke the graph app
+        output = app.invoke({"messages": [user_query]})
+        final_response = output['messages'][-1]
+        
+        # Display Results
+        st.subheader("Response")
+        st.write(final_response)
+    except Exception as e:
+        st.error(f"An error occurred: {e}")