Update test.py

test.py

@@ -1,17 +1,30 @@
-#add ref
-
+# to-do: make it multilingual
 import streamlit as st
 import os
+from openai import OpenAI
+import tempfile
+from langchain.chains import ConversationalRetrievalChain
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import Chroma
+from langchain_community.document_loaders import (
+    PyPDFLoader, 
+    TextLoader, 
+    CSVLoader
+)
 from datetime import datetime
 from pydub import AudioSegment
-import tempfile
 import pytz
-from openai import OpenAI
+
 from langchain.chains import ConversationalRetrievalChain
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_community.vectorstores import Chroma
 from langchain_community.document_loaders import PyPDFLoader, TextLoader, CSVLoader
+import os
+import tempfile
+from datetime import datetime
+import pytz
 
 
 class DocumentRAG:
@@ -41,10 +54,12 @@ class DocumentRAG:
         try:
             documents = []
             for uploaded_file in uploaded_files:
+                # Save uploaded file to a temporary location
                 temp_file_path = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]).name
                 with open(temp_file_path, "wb") as temp_file:
                     temp_file.write(uploaded_file.read())
 
+                # Determine the loader based on the file type
                 if temp_file_path.endswith('.pdf'):
                     loader = PyPDFLoader(temp_file_path)
                 elif temp_file_path.endswith('.txt'):
@@ -54,6 +69,7 @@ class DocumentRAG:
                 else:
                     return f"Unsupported file type: {uploaded_file.name}"
 
+                # Load the documents
                 try:
                     documents.extend(loader.load())
                 except Exception as e:
@@ -62,6 +78,7 @@ class DocumentRAG:
             if not documents:
                 return "No valid documents were processed. Please check your files."
 
+            # Split text for better processing
             text_splitter = RecursiveCharacterTextSplitter(
                 chunk_size=1000,
                 chunk_overlap=200,
@@ -69,14 +86,16 @@ class DocumentRAG:
             )
             documents = text_splitter.split_documents(documents)
 
+            # Combine text for summary
             combined_text = " ".join([doc.page_content for doc in documents])
             self.document_summary = self.generate_summary(combined_text)
 
+            # Create embeddings and initialize retrieval chain
             embeddings = OpenAIEmbeddings(api_key=self.api_key)
             self.document_store = Chroma.from_documents(
                 documents,
                 embeddings,
-                persist_directory=self.chroma_persist_dir
+                persist_directory=self.chroma_persist_dir  # Persistent directory for Chroma
             )
 
             self.qa_chain = ConversationalRetrievalChain.from_llm(
@@ -119,6 +138,8 @@ class DocumentRAG:
 
         try:
             client = OpenAI(api_key=self.api_key)
+
+            # Generate podcast script
             script_response = client.chat.completions.create(
                 model="gpt-4",
                 messages=[
@@ -137,6 +158,7 @@ class DocumentRAG:
             if not script:
                 return "Error: Failed to generate podcast script.", None
 
+            # Convert script to audio
             final_audio = AudioSegment.empty()
             is_first_speaker = True
 
@@ -180,6 +202,47 @@ class DocumentRAG:
         except Exception as e:
             return f"Error generating podcast: {str(e)}", None
 
+    def generate_summary(self, text):
+        """Generate a summary of the provided text."""
+        if not self.api_key:
+            return "API Key not set. Please set it in the environment variables."
+        try:
+            client = OpenAI(api_key=self.api_key)
+            response = client.chat.completions.create(
+                model="gpt-4",
+                messages=[
+                    {"role": "system", "content": "Summarize the document content concisely and provide 3-5 key points for discussion."},
+                    {"role": "user", "content": text[:4000]}
+                ],
+                temperature=0.3
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            return f"Error generating summary: {str(e)}"
+
+    def handle_query(self, question, history):
+        """Handle user queries."""
+        if not self.qa_chain:
+            return history + [("System", "Please process the documents first.")]
+        try:
+            preface = """
+            Instruction: Respond in English. Be professional and concise, keeping the response under 300 words.
+            If you cannot provide an answer, say: "I am not sure about this question. Please try asking something else."
+            """
+            query = f"{preface}\nQuery: {question}"
+
+            result = self.qa_chain({
+                "question": query,
+                "chat_history": [(q, a) for q, a in history]
+            })
+
+            if "answer" not in result:
+                return history + [("System", "Sorry, an error occurred.")]
+
+            history.append((question, result["answer"]))
+            return history
+        except Exception as e:
+            return history + [("System", f"Error: {str(e)}")]
 
 # Initialize RAG system in session state
 if "rag_system" not in st.session_state:
@@ -200,23 +263,70 @@ with st.sidebar:
     st.markdown("3. Ask questions.")
     st.markdown("4. Create podcasts.")
 
+# Streamlit UI
+# Sidebar
+#with st.sidebar:
+    #st.title("About")
+    #st.markdown(
+        #"""
+        #This app is inspired by the [RAG_HW HuggingFace Space](https://huggingface.co/spaces/wint543/RAG_HW).  
+        #It allows users to:
+        #1. Upload and process documents
+        #2. Generate summaries
+        #3. Ask questions
+        #4. Create podcasts
+        #"""
+    #)
+
 # Main App
 st.title("Document Analyzer and Podcast Generator")
 
+# Step 1: Upload and Process Documents
+st.subheader("Step 1: Upload and Process Documents")
 uploaded_files = st.file_uploader("Upload files (PDF, TXT, CSV)", accept_multiple_files=True)
 
 if st.button("Process Documents"):
     if uploaded_files:
+        # Process the uploaded files
         result = st.session_state.rag_system.process_documents(uploaded_files)
-        st.success(result) if "successfully" in result else st.error(result)
+        if "successfully" in result:
+            st.success(result)
+        else:
+            st.error(result)
     else:
         st.warning("No files uploaded.")
 
+# Step 2: Generate Summaries
+st.subheader("Step 2: Generate Summaries")
+if st.session_state.rag_system.document_summary:
+    st.text_area("Document Summary", st.session_state.rag_system.document_summary, height=200)
+else:
+    st.info("Please process documents first to generate summaries.")
+
+# Step 3: Ask Questions
+st.subheader("Step 3: Ask Questions")
+if st.session_state.rag_system.qa_chain:
+    history = []
+    user_question = st.text_input("Ask a question:")
+    if st.button("Submit Question"):
+        # Handle the user query
+        history = st.session_state.rag_system.handle_query(user_question, history)
+        for question, answer in history:
+            st.chat_message("user").write(question)
+            st.chat_message("assistant").write(answer)
+else:
+    st.info("Please process documents first to enable Q&A.")
+
+# Step 4: Generate Podcast
+st.subheader("Step 4: Generate Podcast")
 if st.session_state.rag_system.document_summary:
-    st.subheader("Step 2: Generate Podcast")
     if st.button("Generate Podcast"):
         script, audio_path = st.session_state.rag_system.create_podcast()
         if audio_path:
             st.text_area("Generated Podcast Script", script, height=200)
             st.audio(audio_path, format="audio/mp3")
-            st.success("Podcast generated successfully!")
+            st.success("Podcast generated successfully! You can listen to it above.")
+        else:
+            st.error(script)
+else:
+    st.info("Please process documents and generate summaries before creating a podcast.")