Spaces:
Running
Running
| import streamlit as st | |
| import PyPDF2 | |
| import pptx | |
| import os | |
| from langchain.llms import OpenAI | |
| from langchain.vectorstores.cassandra import Cassandra | |
| from langchain.indexes.vectorstore import VectorStoreIndexWrapper | |
| from langchain.embeddings import OpenAIEmbeddings | |
| import cassio | |
| from langchain.text_splitter import CharacterTextSplitter | |
| from huggingface_hub import login | |
| # Secure API keys (ensure they are set) | |
| ASTRA_DB_APPLICATION_TOKEN = os.getenv("ASTRA_DB_APPLICATION_TOKEN") | |
| ASTRA_DB_ID = os.getenv("ASTRA_DB_ID") | |
| HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACE_API_KEY") | |
| OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") | |
| llm = OpenAI(openai_api_key=OPENAI_API_KEY) | |
| if not ASTRA_DB_APPLICATION_TOKEN or not ASTRA_DB_ID: | |
| st.error("Astra DB credentials are missing. Set the environment variables.") | |
| st.stop() | |
| if not HUGGINGFACE_API_KEY: | |
| st.error("Hugging Face API key is missing. Set the HUGGINGFACE_API_KEY environment variable.") | |
| st.stop() | |
| if not OPENAI_API_KEY: | |
| st.error("OpenAI API key is missing. Set the OPENAI_API_KEY environment variable.") | |
| st.stop() | |
| # Initialize Astra DB connection | |
| cassio.init(token=ASTRA_DB_APPLICATION_TOKEN, database_id=ASTRA_DB_ID) | |
| # Initialize LLM & Embeddings | |
| login(token=HUGGINGFACE_API_KEY) | |
| embedding = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY) | |
| # Initialize vector store | |
| astra_vector_store = Cassandra(embedding=embedding, table_name="qa_mini_demo") | |
| def extract_text_from_pdf(uploaded_file): | |
| """Extract text from a PDF file.""" | |
| text = "" | |
| try: | |
| pdf_reader = PyPDF2.PdfReader(uploaded_file) | |
| for page in pdf_reader.pages: | |
| page_text = page.extract_text() | |
| if page_text: # Avoid NoneType error | |
| text += page_text + "\n" | |
| except Exception as e: | |
| st.error(f"Error reading PDF: {e}") | |
| return text | |
| def extract_text_from_ppt(uploaded_file): | |
| """Extract text from a PowerPoint file.""" | |
| text = "" | |
| try: | |
| presentation = pptx.Presentation(uploaded_file) | |
| for slide in presentation.slides: | |
| for shape in slide.shapes: | |
| if hasattr(shape, "text"): | |
| text += shape.text + "\n" | |
| except Exception as e: | |
| st.error(f"Error reading PPT: {e}") | |
| return text | |
| def main(): | |
| st.title("Chat with Documents") | |
| uploaded_file = st.file_uploader("Upload a PDF or PPT file", type=["pdf", "pptx"]) | |
| extract_button = st.button("Extract Text") | |
| extracted_text = "" | |
| if extract_button and uploaded_file is not None: | |
| if uploaded_file.name.endswith(".pdf"): | |
| extracted_text = extract_text_from_pdf(uploaded_file) | |
| elif uploaded_file.name.endswith(".pptx"): | |
| extracted_text = extract_text_from_ppt(uploaded_file) | |
| if extracted_text: | |
| text_splitter = CharacterTextSplitter(separator="\n", chunk_size=800, chunk_overlap=200, length_function=len) | |
| texts = text_splitter.split_text(extracted_text) | |
| astra_vector_store.add_texts(texts) | |
| st.success("Text extracted and stored successfully!") | |
| # Ensure the vector store index is initialized properly | |
| astra_vector_index = VectorStoreIndexWrapper(vectorstore=astra_vector_store) | |
| query = st.text_input("Enter your query") | |
| submit_query = st.button("Submit Query") | |
| if submit_query and query: | |
| response = astra_vector_index.query(query, llm =llm) | |
| st.write(f"Response: {response}") | |
| if __name__ == "__main__": | |
| main() | |