Pablo276's picture
Update app.py
51458f9 verified
Raw
History Blame
7.56 kB
import gradio as gr
import os
import time
from dotenv import load_dotenv
from pinecone import Pinecone, ServerlessSpec
from langchain_pinecone import PineconeVectorStore
from langchain_openai import OpenAIEmbeddings
from langchain_core.documents import Document
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
import io
import pandas as pd
import re # Make sure to add this import at the top of your file
# 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"""
try:
pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY"))
index_name = os.environ.get("PINECONE_INDEX_NAME")
existing_indexes = [index_info["name"] for index_info in pc.list_indexes()]
if index_name not in existing_indexes:
return []
index = pc.Index(index_name)
# Query with a dummy vector to fetch metadata. Increase top_k if you have more files.
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 hasattr(match, 'metadata') and match.metadata 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."""
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)
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 embedding of the uploaded PDF file."""
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")
existing_indexes = [index_info["name"] for index_info in pc.list_indexes()]
if index_name not in existing_indexes:
pc.create_index(
name=index_name,
dimension=3072,
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)
embeddings = OpenAIEmbeddings(model="text-embedding-3-large", api_key=os.environ.get("OPENAI_API_KEY"))
vector_store = PineconeVectorStore(index=index, embedding=embeddings)
loader = PyPDFLoader(uploaded_file_path)
raw_documents = loader.load()
for doc in raw_documents:
doc.metadata['source'] = original_filename
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=800,
chunk_overlap=400,
length_function=len,
is_separator_regex=False,
)
documents = text_splitter.split_documents(raw_documents)
# --- INIZIO: CODICE CORRETTO ---
# 1. Rimuove l'estensione, gli spazi iniziali/finali e converte in minuscolo.
sanitized_filename = original_filename.replace('.pdf', '').strip().lower()
# 2. Sostituisce qualsiasi carattere non alfanumerico con un trattino.
sanitized_filename = re.sub(r'[^a-z0-9]', '-', sanitized_filename)
# 3. Sostituisce trattini multipli consecutivi con un singolo trattino.
sanitized_filename = re.sub(r'-+', '-', sanitized_filename)
# 4. (Opzionale ma consigliato) Rimuove eventuali trattini all'inizio o alla fine.
sanitized_filename = sanitized_filename.strip('-')
# 5. Usa il nome sanificato per generare gli ID.
uuids = [f"{sanitized_filename}-{i}" for i in range(len(documents))]
# --- FINE: CODICE CORRETTO ---
batch_size = 100
for i in range(0, len(documents), batch_size):
batch_docs = documents[i:i+batch_size]
batch_ids = uuids[i:i+batch_size]
vector_store.add_documents(documents=batch_docs, ids=batch_ids)
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():
files = get_stored_files()
if files:
return pd.DataFrame({"Stored Files": files})
else:
return pd.DataFrame({"Stored Files": []})
# CORRECTION: Updated function to handle the select event correctly
def handle_file_selection(evt: gr.SelectData):
"""
Handles the file selection event from the DataFrame.
evt.value contains the value of the selected cell.
"""
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]
)
# CORRECTION: Removed the 'inputs' argument.
# The event data 'evt' is now passed automatically to 'handle_file_selection'.
file_df.select(
fn=handle_file_selection,
inputs=None, # Explicitly setting to None or removing this line works
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()