import gradio as gr import os import time import re import pandas as pd from dotenv import load_dotenv from pinecone import Pinecone, ServerlessSpec # --- Langchain components for document processing and embedding --- from langchain_openai import OpenAIEmbeddings from langchain_core.documents import Document from langchain_community.document_loaders import PyPDFLoader from langchain_text_splitters import RecursiveCharacterTextSplitter # Load environment variables from a .env file load_dotenv() # --- Backend Functions --- def get_stored_files(): """Retrieve a list of files currently stored in the Pinecone index using the pinecone library.""" try: PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY") pc = Pinecone(api_key=PINECONE_API_KEY) index_name = os.environ.get("PINECONE_INDEX_NAME") print(PINECONE_API_KEY, index_name) # Check if the index exists. If not, there are no files. if index_name not in [index_info["name"] for index_info in pc.list_indexes()]: return [] index = pc.Index(index_name) # A dummy query to fetch metadata from all vectors. # Increase top_k if you expect to have more than 1000 chunks. results = index.query(vector=[0.0] * 3072, top_k=10000, include_metadata=True) unique_files = set() if results.matches: for match in results.matches: if 'metadata' in match and 'source' in match.metadata: unique_files.add(match.metadata['source']) return sorted(list(unique_files)) except Exception as e: print(f"Error retrieving stored files: {str(e)}") return [] def delete_file_from_vectorstore(filename): """Deletes all vectors associated with a specific filename from Pinecone using the pinecone library.""" if not filename: return "No file selected for deletion.", get_files_df() try: pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY")) index_name = os.environ.get("PINECONE_INDEX_NAME") index = pc.Index(index_name) # Use metadata filtering to delete all vectors associated with the file. index.delete(filter={"source": {"$eq": filename}}) return f"Successfully deleted {filename}.", get_files_df() except Exception as e: return f"Error while deleting the file: {str(e)}", get_files_df() def embedder(uploaded_file_path): """ Handles the embedding of the uploaded PDF file using langchain for processing and the pinecone library for vector store operations. """ if uploaded_file_path is None: return "No file uploaded. Please upload a PDF.", get_files_df() try: original_filename = os.path.basename(uploaded_file_path) pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY")) index_name = os.environ.get("PINECONE_INDEX_NAME") embedding_dimension = 3072 # As specified for text-embedding-3-large # Create the index if it doesn't exist if index_name not in [index_info["name"] for index_info in pc.list_indexes()]: pc.create_index( name=index_name, dimension=embedding_dimension, metric="cosine", spec=ServerlessSpec(cloud="aws", region="us-east-1"), ) while not pc.describe_index(index_name).status["ready"]: time.sleep(1) index = pc.Index(index_name) # 1. Load and Split Document loader = PyPDFLoader(uploaded_file_path) raw_documents = loader.load() text_splitter = RecursiveCharacterTextSplitter( chunk_size=800, chunk_overlap=400, length_function=len, ) documents = text_splitter.split_documents(raw_documents) # 2. Create Embeddings embeddings_model = OpenAIEmbeddings(model="text-embedding-3-large", api_key=os.environ.get("OPENAI_API_KEY")) texts_to_embed = [doc.page_content for doc in documents] embeddings = embeddings_model.embed_documents(texts_to_embed) # 3. Sanitize filename and prepare vectors for upsert sanitized_filename = re.sub(r'[^a-z0-9]', '-', original_filename.replace('.pdf', '').strip().lower()) sanitized_filename = re.sub(r'-+', '-', sanitized_filename).strip('-') vectors_to_upsert = [] for i, (doc, vec) in enumerate(zip(documents, embeddings)): vector_id = f"{sanitized_filename}-{i}" metadata = { "text": doc.page_content, "source": original_filename } vectors_to_upsert.append({"id": vector_id, "values": vec, "metadata": metadata}) # 4. Upsert vectors to Pinecone in batches batch_size = 100 for i in range(0, len(vectors_to_upsert), batch_size): batch = vectors_to_upsert[i:i+batch_size] index.upsert(vectors=batch) return f"File '{original_filename}' successfully embedded!", get_files_df() except Exception as e: return f"Unable to create embeddings: {str(e)}", get_files_df() # --- Gradio Interface Functions --- def get_files_df(): """Creates a DataFrame from the list of stored files for Gradio display.""" files = get_stored_files() if files: return pd.DataFrame({"Stored Files": files}) else: return pd.DataFrame({"Stored Files": []}) def handle_file_selection(evt: gr.SelectData): """Handles the file selection event from the DataFrame.""" if evt.value: return evt.value return "" # --- Gradio UI --- with gr.Blocks(theme=gr.themes.Soft(), title="PDF Uploader") as demo: gr.Markdown("# PDF File Uploader for Chatbot") gr.Markdown("Upload PDF files to add their content to the chatbot's knowledge base.") with gr.Row(): with gr.Column(scale=1): gr.Markdown("## 📤 Upload New File") file_uploader = gr.File( label="Upload your PDF file", file_types=[".pdf"], type="filepath" ) upload_button = gr.Button("Upload to Chatbot Memory", variant="primary") upload_status = gr.Markdown("") with gr.Column(scale=1): gr.Markdown("## 🗂️ Stored Files") refresh_button = gr.Button("Refresh File List") file_df = gr.DataFrame( value=get_files_df, headers=["Stored Files"], interactive=True ) selected_file_text = gr.Textbox( label="Selected File", interactive=False, placeholder="Click on a file above to select it" ) delete_button = gr.Button("🗑️ Delete Selected File", variant="stop") delete_status = gr.Markdown("") # --- Event Handlers --- upload_button.click( fn=embedder, inputs=[file_uploader], outputs=[upload_status, file_df] ) refresh_button.click( fn=get_files_df, inputs=[], outputs=[file_df] ) file_df.select( fn=handle_file_selection, outputs=[selected_file_text] ) delete_button.click( fn=delete_file_from_vectorstore, inputs=[selected_file_text], outputs=[delete_status, file_df] ) if __name__ == "__main__": demo.launch()