Update app.py

app.py

@@ -1,68 +1,101 @@
+import os
 import streamlit as st
 import pandas as pd
 import matplotlib.pyplot as plt
-from langchain_core.prompts import PromptTemplate
-from langchain_community.llms import OpenAI
-from langchain_core.chains import LLMChain
-
-
-# Fetch API Keys from environment variables
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-
-# Verify required environment variables are set
-if not OPENAI_API_KEY:
-    raise ValueError("Missing required environment variable: 'OPENAI_API_KEY'")
-
-
-# Title of the App
-st.title("ReAct Agent")
-
-# Step 1: User Input Collection
-user_input = st.text_area("Enter the input prompt for the LLM:", placeholder="Describe what you want...")
-num_intermediate_steps = st.number_input("Number of Intermediate Steps to Display:", min_value=1, value=3)
-
-# Check if inputs are provided
-if st.button("Run Processing"):
-    if not api_key or not user_input:
-        st.error("Please provide all required inputs!")
-    else:
-        st.success("Inputs received! Processing...")
-
-        # Step 2: LLM Initialization and Template
-        st.header("Step 2: LLM Processing - Intermediate Outputs")
-        try:
-            llm = OpenAI(api_key=api_key, temperature=0.7)
-            prompt = PromptTemplate(
-                input_variables=["text"],
-                template="Process the following text step-by-step: {text}"
-            )
-
-            chain = LLMChain(llm=llm, prompt=prompt)
-            intermediate_outputs = []
-
-            # Simulate multiple intermediate outputs
-            for step in range(num_intermediate_steps):
-                response = chain.run(user_input + f" Step {step+1}")
-                intermediate_outputs.append(response)
-                st.write(f"**Intermediate Step {step+1}:**")
-                st.info(response)
-
-            # Step 3: Final Processing
-            st.header("Step 3: Visualization")
-            # Convert responses to a DataFrame
-            df = pd.DataFrame({"Step": [f"Step {i+1}" for i in range(num_intermediate_steps)],
-                               "Output": intermediate_outputs})
-            st.dataframe(df)
-
-            # Plot the outputs' length as a graph
-            st.subheader("Output Length per Step")
-            df['Length'] = df['Output'].apply(len)
-            fig, ax = plt.subplots()
-            ax.bar(df['Step'], df['Length'])
-            ax.set_xlabel("Steps")
-            ax.set_ylabel("Output Length")
-            ax.set_title("Text Length at Each Step")
-            st.pyplot(fig)
-
-        except Exception as e:
-            st.error(f"An error occurred: {str(e)}")
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_openai import ChatOpenAI
+from langgraph.graph import MessagesState
+from langgraph.graph import START, StateGraph
+from langgraph.prebuilt import tools_condition
+from langgraph.prebuilt import ToolNode
+from langchain_core.messages import HumanMessage, SystemMessage
+from IPython.display import Image, display
+
+# ------------------- Environment Variable Setup -------------------
+# Fetch API keys from environment variables
+openai_api_key = os.getenv("OPENAI_API_KEY")
+tavily_api_key = os.getenv("TAVILY_API_KEY")
+
+# Verify if API keys are set
+if not openai_api_key:
+    raise ValueError("Missing required environment variable: OPENAI_API_KEY")
+if not tavily_api_key:
+    raise ValueError("Missing required environment variable: TAVILY_API_KEY")
+
+# ------------------- Tool Definitions -------------------
+# Tavily Search Tool
+tavily_tool = TavilySearchResults(max_results=5)
+
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers."""
+    return a * b
+
+def add(a: int, b: int) -> int:
+    """Add two numbers."""
+    return a + b
+
+def divide(a: int, b: int) -> float:
+    """Divide two numbers."""
+    if b == 0:
+        raise ValueError("Division by zero is not allowed.")
+    return a / b
+
+# Combine tools
+tools = [add, multiply, divide, tavily_tool]
+
+# ------------------- LLM and System Message Setup -------------------
+llm = ChatOpenAI(model="gpt-4o-mini")
+llm_with_tools = llm.bind_tools(tools, parallel_tool_calls=False)
+sys_msg = SystemMessage(content="You are a helpful assistant tasked with performing arithmetic and search on a set of inputs.")
+
+# ------------------- LangGraph Workflow -------------------
+def assistant(state: MessagesState):
+    """Assistant node to invoke LLM with tools."""
+    return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}
+
+# Define the graph
+app_graph = StateGraph(MessagesState)
+app_graph.add_node("assistant", assistant)
+app_graph.add_node("tools", ToolNode(tools))
+app_graph.add_edge(START, "assistant")
+app_graph.add_conditional_edges("assistant", tools_condition)
+app_graph.add_edge("tools", "assistant")
+react_graph = app_graph.compile()
+
+# Display graph visualization
+st.header("LangGraph Workflow Visualization")
+try:
+    image_path = react_graph.get_graph(xray=True).draw_mermaid_png()
+    display(Image(image_path))
+    st.image(image_path, caption="LangGraph Workflow Visualization")
+except Exception as e:
+    st.error(f"Failed to display graph visualization: {e}")
+
+# ------------------- Streamlit Interface -------------------
+#st.title("ReAct Agent with Arithmetic and Search")
+
+# Prompt user for inputs
+user_question = st.text_area("Enter your question:",
+                             placeholder="Example: 'Add 3 and 4. Multiply the result by 2. Divide it by 5.'")
+
+if st.button("Submit"):
+    if not user_question.strip():
+        st.error("Please enter a valid question.")
+        st.stop()
+    
+    st.info("Processing your question...")
+    messages = [HumanMessage(content=user_question)]
+    response = react_graph.invoke({"messages": messages})
+    
+    # Display results
+    st.subheader("Responses")
+    for m in response['messages']:
+        st.write(m.content)
+
+    st.success("Processing complete!")
+
+# Example Placeholder Suggestions
+st.sidebar.subheader("Example Questions")
+st.sidebar.write("- Add 3 and 4. Multiply the result by 2. Divide it by 5.")
+st.sidebar.write("- Tell me how many centuries Virat Kohli scored.")
+st.sidebar.write("- Search for the tallest building in the world.")