version3 / interim.py
DrishtiSharma's picture
Create interim.py
afdff8e verified
raw
history blame
7.15 kB
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.")