Update interim.py

interim.py

@@ -1,9 +1,7 @@
-#fix workflow
 import os
 import streamlit as st
 import pandas as pd
 import matplotlib.pyplot as plt
-import networkx as nx
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_openai import ChatOpenAI
 from langgraph.graph import MessagesState
@@ -11,6 +9,7 @@ from langgraph.graph import START, StateGraph
 from langgraph.prebuilt import tools_condition
 from langgraph.prebuilt import ToolNode
 from langchain_core.messages import HumanMessage, SystemMessage
+import tempfile
 
 # ------------------- Environment Variable Setup -------------------
 # Fetch API keys from environment variables
@@ -24,6 +23,7 @@ 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:
@@ -40,9 +40,10 @@ def divide(a: int, b: int) -> float:
         raise ValueError("Division by zero is not allowed.")
     return a / b
 
+# Combine tools
 tools = [add, multiply, divide, tavily_tool]
 
-# ------------------- LLM Setup -------------------
+# ------------------- 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.")
@@ -52,6 +53,7 @@ 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))
@@ -60,28 +62,22 @@ app_graph.add_conditional_edges("assistant", tools_condition)
 app_graph.add_edge("tools", "assistant")
 react_graph = app_graph.compile()
 
-# ------------------- Streamlit Interface -------------------
-st.title("ReAct Agent")
-
-# Display the workflow graph using NetworkX
-st.header("LangGraph Workflow Visualization")
+# Save graph visualization as an image
+with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
+    graph = react_graph.get_graph(xray=True)
+    tmpfile.write(graph.draw_mermaid_png())  # Write binary image data to file
+    graph_image_path = tmpfile.name
 
-G = nx.DiGraph()
-G.add_edge("START", "assistant")
-G.add_edge("assistant", "tools", label="tools_condition")
-G.add_edge("tools", "assistant", label="loop back")
+# ------------------- Streamlit Interface -------------------
+st.title("ReAct Agent for Arithmetic Ops & Web Search")
 
-plt.figure(figsize=(10, 6))
-pos = nx.spring_layout(G, seed=42)
-nx.draw(G, pos, with_labels=True, node_size=3000, node_color="lightblue", font_size=10, font_weight="bold")
-nx.draw_networkx_edge_labels(G, pos, edge_labels={
-    ("assistant", "tools"): "tools_condition",
-    ("tools", "assistant"): "loop back"
-}, font_color="red")
-st.pyplot(plt)
+# Display the workflow graph
+#st.header("LangGraph Workflow Visualization")
+st.image(graph_image_path, caption="Workflow Visualization")
 
-# User input
-user_question = st.text_area("Enter your question:", placeholder="Example: 'Add 3 and 4. Multiply the result by 2. Divide it by 5.'")
+# 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():
@@ -92,12 +88,14 @@ if st.button("Submit"):
     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 Questions
+# 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.")