Update app.py

app.py

@@ -1,6 +1,8 @@
 import os
+import io
 import chromadb
 import streamlit as st
+import matplotlib.pyplot as plt  # For matplotlib graph handling
 from langchain_openai import ChatOpenAI
 from langchain.agents import AgentExecutor, create_openai_tools_agent
 from langchain_core.messages import BaseMessage, HumanMessage
@@ -22,12 +24,10 @@ import operator
 from langchain_core.tools import tool
 from glob import glob
 
-
-# Clear ChromaDB cache to fix tenant issue
+# Clear ChromaDB cache
 chromadb.api.client.SharedSystemClient.clear_system_cache()
 
 # Load environment variables
-
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
 
@@ -35,7 +35,7 @@ if not OPENAI_API_KEY or not TAVILY_API_KEY:
     st.error("Please set OPENAI_API_KEY and TAVILY_API_KEY in your environment variables.")
     st.stop()
 
-# Initialize API keys and LLM
+# Initialize LLM
 llm = ChatOpenAI(model="gpt-4-1106-preview", openai_api_key=OPENAI_API_KEY)
 
 # Utility Functions
@@ -53,19 +53,15 @@ def agent_node(state, agent, name):
     result = agent.invoke(state)
     output_content = result["output"]
     
-    # Check if the output contains Python code that generates a graph
+    # Check if Python code generates a graph
     if "matplotlib" in output_content or "plt." in output_content:
         exec_locals = {}
         try:
-            exec(output_content, {}, exec_locals)  # Safely execute the code
-            fig = plt.gcf()  # Get the current matplotlib figure
-            
-            # Save the figure to a buffer
+            exec(output_content, {}, exec_locals)
+            fig = plt.gcf()
             buf = io.BytesIO()
             fig.savefig(buf, format="png")
             buf.seek(0)
-            
-            # Add image to session state for display
             st.session_state.graph_image = buf
         except Exception as e:
             output_content += f"\nError: {str(e)}"
@@ -88,28 +84,24 @@ def RAG(state):
     result = retrieval_chain.invoke(question)
     return result
 
-# Load Tools
+# Tools Setup
 tavily_tool = TavilySearchResults(max_results=5, tavily_api_key=TAVILY_API_KEY)
 python_repl_tool = PythonREPLTool()
 
 # Streamlit UI
-st.title("Multi-Agent w Supervisor")
+st.title("Multi-Agent Workflow with Supervisor")
 
-# Example questions for immediate testing
 example_questions = [
-    #"Code hello world and print it",
     "What is James McIlroy aiming for in sports?",
     "Fetch India's GDP over the past 5 years and draw a line graph.",
     "Fetch Japan's GDP over the past 4 years from RAG, then draw a line graph."
 ]
 
-# File Selection Section
 source_files = glob("sources/*.txt")
 selected_files = st.multiselect("Select files from the source directory:", source_files, default=source_files[:2])
-
 uploaded_files = st.file_uploader("Or upload your TXT files:", accept_multiple_files=True, type=['txt'])
 
-# Combine Files
+# Document Handling
 all_docs = []
 if selected_files:
     for file_path in selected_files:
@@ -122,18 +114,17 @@ if uploaded_files:
         all_docs.append(Document(page_content=content, metadata={"name": uploaded_file.name}))
 
 if not all_docs:
-    st.warning("Please select files from the source directory or upload TXT files.")
+    st.warning("Please select files or upload TXT files.")
     st.stop()
 
-# Process Documents
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10, length_function=len)
+# Document Splitting and Embedding
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10)
 split_docs = text_splitter.split_documents(all_docs)
-
 embeddings = HuggingFaceBgeEmbeddings(model_name="BAAI/bge-base-en-v1.5", model_kwargs={'device': 'cpu'}, encode_kwargs={'normalize_embeddings': True})
 db = Chroma.from_documents(split_docs, embeddings)
 retriever = db.as_retriever(search_kwargs={"k": 4})
 
-# Create Agents
+# Agents
 research_agent = create_agent(llm, [tavily_tool], "You are a web researcher.")
 code_agent = create_agent(llm, [python_repl_tool], "You may generate safe python code to analyze data and generate charts using matplotlib.")
 RAG_agent = create_agent(llm, [RAG], "Use this tool when questions are related to Japan or Sports category.")
@@ -143,10 +134,7 @@ code_node = functools.partial(agent_node, agent=code_agent, name="Coder")
 rag_node = functools.partial(agent_node, agent=RAG_agent, name="RAG")
 
 members = ["RAG", "Researcher", "Coder"]
-system_prompt = (
-    "You are a supervisor managing these workers: {members}. Respond with the next worker or FINISH. "
-    "Use RAG tool for Japan or Sports questions."
-)
+system_prompt = "You are a supervisor managing these workers: {members}. Respond with the next worker or FINISH."
 options = ["FINISH"] + members
 function_def = {
     "name": "route", "description": "Select the next role.",
@@ -157,10 +145,8 @@ prompt = ChatPromptTemplate.from_messages([
     MessagesPlaceholder(variable_name="messages"),
     ("system", "Given the conversation above, who should act next? Select one of: {options}"),
 ]).partial(options=str(options), members=", ".join(members))
-
 supervisor_chain = (prompt | llm.bind_functions(functions=[function_def], function_call="route") | JsonOutputFunctionsParser())
 
-# Workflow
 class AgentState(TypedDict):
     messages: Annotated[Sequence[BaseMessage], operator.add]
     next: str
@@ -188,6 +174,7 @@ user_input = st.text_area("Enter your task or question:", placeholder=example_qu
 def run_workflow(task):
     st.session_state.outputs.clear()
     st.session_state.outputs.append(f"User Input: {task}")
+    st.session_state.graph_image = None
     for state in graph.stream({"messages": [HumanMessage(content=task)]}):
         if "__end__" not in state:
             st.session_state.outputs.append(str(state))
@@ -199,10 +186,10 @@ if st.button("Run Workflow"):
     else:
         st.warning("Please enter a task or question.")
 
-st.subheader("Example Questions:")
-for example in example_questions:
-    st.text(f"- {example}")
-
 st.subheader("Workflow Output:")
 for output in st.session_state.outputs:
     st.text(output)
+
+if "graph_image" in st.session_state and st.session_state.graph_image:
+    st.subheader("Generated Graph:")
+    st.image(st.session_state.graph_image, caption="Generated Line Graph", use_column_width=True)