Create test.py

test.py

@@ -0,0 +1,220 @@
+import streamlit as st
+import os
+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
+
+
+class DocumentRAG:
+    def __init__(self):
+        self.document_store = None
+        self.qa_chain = None
+        self.document_summary = ""
+        self.chat_history = []
+        self.last_processed_time = None
+        self.api_key = os.getenv("OPENAI_API_KEY")  # Fetch the API key from environment variable
+        self.init_time = datetime.now(pytz.UTC)
+
+        if not self.api_key:
+            raise ValueError("API Key not found. Make sure to set the 'OPENAI_API_KEY' environment variable.")
+
+        # Persistent directory for Chroma to avoid tenant-related errors
+        self.chroma_persist_dir = "./chroma_storage"
+        os.makedirs(self.chroma_persist_dir, exist_ok=True)
+
+    def process_documents(self, uploaded_files):
+        """Process uploaded files by saving them temporarily and extracting content."""
+        if not self.api_key:
+            return "Please set the OpenAI API key in the environment variables."
+        if not uploaded_files:
+            return "Please upload documents first."
+
+        try:
+            documents = []
+            for uploaded_file in uploaded_files:
+                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())
+
+                if temp_file_path.endswith('.pdf'):
+                    loader = PyPDFLoader(temp_file_path)
+                elif temp_file_path.endswith('.txt'):
+                    loader = TextLoader(temp_file_path)
+                elif temp_file_path.endswith('.csv'):
+                    loader = CSVLoader(temp_file_path)
+                else:
+                    return f"Unsupported file type: {uploaded_file.name}"
+
+                try:
+                    documents.extend(loader.load())
+                except Exception as e:
+                    return f"Error loading {uploaded_file.name}: {str(e)}"
+
+            if not documents:
+                return "No valid documents were processed. Please check your files."
+
+            text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=1000,
+                chunk_overlap=200,
+                length_function=len
+            )
+            documents = text_splitter.split_documents(documents)
+
+            combined_text = " ".join([doc.page_content for doc in documents])
+            self.document_summary = self.generate_summary(combined_text)
+
+            embeddings = OpenAIEmbeddings(api_key=self.api_key)
+            self.document_store = Chroma.from_documents(
+                documents,
+                embeddings,
+                persist_directory=self.chroma_persist_dir
+            )
+
+            self.qa_chain = ConversationalRetrievalChain.from_llm(
+                ChatOpenAI(temperature=0, model_name='gpt-4', api_key=self.api_key),
+                self.document_store.as_retriever(search_kwargs={'k': 6}),
+                return_source_documents=True,
+                verbose=False
+            )
+
+            self.last_processed_time = datetime.now(pytz.UTC)
+            return "Documents processed successfully!"
+        except Exception as e:
+            return f"Error processing documents: {str(e)}"
+
+    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 create_podcast(self):
+        """Generate a podcast script and audio based on the document summary."""
+        if not self.document_summary:
+            return "Please process documents before generating a podcast.", None
+
+        if not self.api_key:
+            return "Please set the OpenAI API key in the environment variables.", None
+
+        try:
+            client = OpenAI(api_key=self.api_key)
+            script_response = client.chat.completions.create(
+                model="gpt-4",
+                messages=[
+                    {"role": "system", "content": "You are a professional podcast producer. Create a natural dialogue based on the provided document summary."},
+                    {"role": "user", "content": f"""Based on the following document summary, create a 1-2 minute podcast script:
+                    1. Clearly label the dialogue as 'Host 1:' and 'Host 2:'
+                    2. Keep the content engaging and insightful.
+                    3. Use conversational language suitable for a podcast.
+                    4. Ensure the script has a clear opening and closing.
+                    Document Summary: {self.document_summary}"""}
+                ],
+                temperature=0.7
+            )
+
+            script = script_response.choices[0].message.content
+            if not script:
+                return "Error: Failed to generate podcast script.", None
+
+            final_audio = AudioSegment.empty()
+            is_first_speaker = True
+
+            lines = [line.strip() for line in script.split("\n") if line.strip()]
+            for line in lines:
+                if ":" not in line:
+                    continue
+
+                speaker, text = line.split(":", 1)
+                if not text.strip():
+                    continue
+
+                try:
+                    voice = "nova" if is_first_speaker else "onyx"
+                    audio_response = client.audio.speech.create(
+                        model="tts-1",
+                        voice=voice,
+                        input=text.strip()
+                    )
+
+                    temp_audio_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
+                    audio_response.stream_to_file(temp_audio_file.name)
+
+                    segment = AudioSegment.from_file(temp_audio_file.name)
+                    final_audio += segment
+                    final_audio += AudioSegment.silent(duration=300)
+
+                    is_first_speaker = not is_first_speaker
+                except Exception as e:
+                    print(f"Error generating audio for line: {text}")
+                    print(f"Details: {e}")
+                    continue
+
+            if len(final_audio) == 0:
+                return "Error: No audio could be generated.", None
+
+            output_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3").name
+            final_audio.export(output_file, format="mp3")
+            return script, output_file
+
+        except Exception as e:
+            return f"Error generating podcast: {str(e)}", None
+
+
+# Initialize RAG system in session state
+if "rag_system" not in st.session_state:
+    st.session_state.rag_system = DocumentRAG()
+
+# 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 upload documents, generate summaries, ask questions, and create podcasts.
+        """
+    )
+    st.markdown("### Steps:")
+    st.markdown("1. Upload documents.")
+    st.markdown("2. Generate summaries.")
+    st.markdown("3. Ask questions.")
+    st.markdown("4. Create podcasts.")
+
+# Main App
+st.title("Document Analyzer and Podcast Generator")
+
+uploaded_files = st.file_uploader("Upload files (PDF, TXT, CSV)", accept_multiple_files=True)
+
+if st.button("Process Documents"):
+    if uploaded_files:
+        result = st.session_state.rag_system.process_documents(uploaded_files)
+        st.success(result) if "successfully" in result else st.error(result)
+    else:
+        st.warning("No files uploaded.")
+
+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!")