Spaces:
Running
Running
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 | |
import pytz | |
# DocumentRAG class with environment variable support for API Key | |
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.") | |
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: | |
# 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'): | |
loader = TextLoader(temp_file_path) | |
elif temp_file_path.endswith('.csv'): | |
loader = CSVLoader(temp_file_path) | |
else: | |
continue | |
# Load the documents | |
try: | |
documents.extend(loader.load()) | |
except Exception as e: | |
print(f"Error loading {temp_file_path}: {str(e)}") | |
continue | |
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, | |
length_function=len | |
) | |
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) | |
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 handle_query(self, question, history): | |
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)}")] | |
# Streamlit UI | |
st.title("Document Analyzer and Podcast Generator") | |
# Fetch the API key status | |
if "OPENAI_API_KEY" not in os.environ or not os.getenv("OPENAI_API_KEY"): | |
st.error("The 'OPENAI_API_KEY' environment variable is not set. Please configure it in your hosting environment.") | |
else: | |
st.success("API Key successfully loaded from environment variable.") | |
# Initialize RAG system | |
try: | |
rag_system = DocumentRAG() | |
except ValueError as e: | |
st.error(str(e)) | |
st.stop() | |
# File upload | |
st.subheader("Step 1: Upload 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 = rag_system.process_documents(uploaded_files) | |
# Ensure that result is a string and display appropriately | |
if isinstance(result, str): | |
if "successfully" in result: | |
st.success(result) | |
else: | |
st.error(result) | |
else: | |
st.error("An unexpected error occurred during document processing.") | |
else: | |
st.warning("No files uploaded.") | |
# Document Q&A | |
st.subheader("Step 2: Ask Questions") | |
if rag_system.qa_chain: | |
history = [] | |
user_question = st.text_input("Ask a question:") | |
if st.button("Submit Question"): | |
history = 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 before asking questions.") | |