Update interim.py

interim.py

@@ -1,5 +1,4 @@
 import os
-import chromadb
 import streamlit as st
 from dotenv import load_dotenv
 from langchain_openai import ChatOpenAI
@@ -7,7 +6,7 @@ from langchain.agents import AgentExecutor, create_openai_tools_agent
 from langchain_core.messages import BaseMessage, HumanMessage
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_experimental.tools import PythonREPLTool
-from langchain_community.document_loaders import DirectoryLoader
+from langchain_community.document_loaders import DirectoryLoader, TextLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain.embeddings import HuggingFaceBgeEmbeddings
@@ -21,12 +20,10 @@ from typing import Annotated, Sequence, TypedDict
 import functools
 import operator
 from langchain_core.tools import tool
-
-
-# Clear ChromaDB cache to fix tenant issue
-chromadb.api.client.SharedSystemClient.clear_system_cache()
+from glob import glob
 
 # Load environment variables
+
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
 
@@ -67,29 +64,51 @@ def RAG(state):
     result = retrieval_chain.invoke(question)
     return result
 
-# Load Tools and Retriever
+# Load Tools
 tavily_tool = TavilySearchResults(max_results=5, tavily_api_key=TAVILY_API_KEY)
 python_repl_tool = PythonREPLTool()
 
-# File Upload Section
+# Streamlit UI
 st.title("Multi-Agent Workflow Demonstration")
-uploaded_files = st.file_uploader("Upload your source files (TXT)", accept_multiple_files=True, type=['txt'])
+
+# Example questions for immediate testing
+example_questions = [
+    "Code hello world and print it to the terminal",
+    "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("source/*.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
+all_docs = []
+if selected_files:
+    for file_path in selected_files:
+        loader = TextLoader(file_path)
+        all_docs.extend(loader.load())
 
 if uploaded_files:
-    docs = []
     for uploaded_file in uploaded_files:
         content = uploaded_file.read().decode("utf-8")
-        docs.append(Document(page_content=content, metadata={"name": uploaded_file.name}))
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10, length_function=len)
-    new_docs = text_splitter.split_documents(documents=docs)
-    embeddings = HuggingFaceBgeEmbeddings(model_name="BAAI/bge-base-en-v1.5", model_kwargs={'device': 'cpu'}, encode_kwargs={'normalize_embeddings': True})
-    db = Chroma.from_documents(new_docs, embeddings)
-    retriever = db.as_retriever(search_kwargs={"k": 4})
-else:
-    retriever = None
-    st.warning("Please upload at least one text file to proceed.")
+        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.stop()
 
+# Process Documents
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10, length_function=len)
+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
 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.")
@@ -107,10 +126,7 @@ system_prompt = (
 options = ["FINISH"] + members
 function_def = {
     "name": "route", "description": "Select the next role.",
-    "parameters": {
-        "title": "routeSchema", "type": "object",
-        "properties": {"next": {"anyOf": [{"enum": options}]}}, "required": ["next"]
-    }
+    "parameters": {"title": "routeSchema", "type": "object", "properties": {"next": {"anyOf": [{"enum": options}]}}, "required": ["next"]}
 }
 prompt = ChatPromptTemplate.from_messages([
     ("system", system_prompt),
@@ -120,7 +136,7 @@ prompt = ChatPromptTemplate.from_messages([
 
 supervisor_chain = (prompt | llm.bind_functions(functions=[function_def], function_call="route") | JsonOutputFunctionsParser())
 
-# Build Workflow
+# Workflow
 class AgentState(TypedDict):
     messages: Annotated[Sequence[BaseMessage], operator.add]
     next: str
@@ -139,11 +155,11 @@ workflow.add_conditional_edges("supervisor", lambda x: x["next"], conditional_ma
 workflow.set_entry_point("supervisor")
 graph = workflow.compile()
 
-# Streamlit UI
+# Workflow Execution
 if 'outputs' not in st.session_state:
     st.session_state.outputs = []
 
-user_input = st.text_area("Enter your task or question:")
+user_input = st.text_area("Enter your task or question:", placeholder=example_questions[0])
 
 def run_workflow(task):
     st.session_state.outputs.clear()
@@ -159,6 +175,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)