Create app.py

app.py

@@ -0,0 +1,125 @@
+import gradio as gr
+import chromadb
+import os
+import tempfile
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import Chroma
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.document_loaders import PyPDFLoader
+
+def process_pdf(file_binary):
+    log = []
+    status_message = ""
+
+    if not file_binary:
+        return "No file uploaded.", "Error: No file was provided."
+
+    try:
+        log.append("Starting PDF upload and processing...")
+
+        # Write uploaded PDF bytes to a temporary file
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
+            temp_file.write(file_binary)
+            temp_path = temp_file.name
+        log.append(f"Temporary PDF path: {temp_path}")
+
+        # Load and extract text from the PDF
+        try:
+            loader = PyPDFLoader(temp_path)
+            documents = loader.load()
+            log.append(f"Loaded {len(documents)} page(s) from PDF.")
+        except Exception as e:
+            raise RuntimeError(f"Error loading PDF: {e}")
+
+        # Split text into chunks
+        try:
+            text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+            splits = text_splitter.split_documents(documents)
+            log.append(f"Text split into {len(splits)} chunk(s).")
+        except Exception as e:
+            raise RuntimeError(f"Error splitting text: {e}")
+
+        # Create an in-memory Chroma client (ephemeral)
+        try:
+            log.append("Initializing in-memory ChromaDB...")
+            chroma_client = chromadb.Client()  # in-memory, no local storage
+            embeddings = HuggingFaceEmbeddings(
+                model_name="sentence-transformers/all-MiniLM-L6-v2"
+            )
+            Chroma.from_documents(
+                splits,
+                embeddings,
+                client=chroma_client
+            )
+            log.append("Successfully stored PDF chunks in ChromaDB.")
+        except Exception as e:
+            raise RuntimeError(f"Error creating ChromaDB vector store: {e}")
+
+        status_message = "PDF processed and stored in (ephemeral) ChromaDB successfully!"
+        log.append(status_message)
+
+    except Exception as e:
+        status_message = "Error"
+        log.append(f"Exception occurred: {str(e)}")
+
+    return status_message, "\n".join(log)
+
+
+def retrieve_context(query):
+    log = []
+    if not query:
+        return "Error: No query provided."
+
+    try:
+        log.append("Retrieving context from in-memory ChromaDB...")
+
+        # Re-initialize the in-memory Chroma client each time
+        chroma_client = chromadb.Client()  # ephemeral
+        embeddings = HuggingFaceEmbeddings(
+            model_name="sentence-transformers/all-MiniLM-L6-v2"
+        )
+        vectorstore = Chroma(embedding_function=embeddings, client=chroma_client)
+
+        # Perform similarity search
+        results = vectorstore.similarity_search(query, k=3)
+        if results:
+            log.append(f"Found {len(results)} matching chunk(s).")
+            return "\n\n".join([doc.page_content for doc in results])
+        else:
+            log.append("No matching context found in the current in-memory DB.")
+            return "No relevant context found. Have you processed a PDF yet?"
+
+    except Exception as e:
+        log.append(f"Error retrieving context: {str(e)}")
+        return "\n".join(log)
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("## PDF Context Retriever with ChromaDB (In-Memory)")
+
+    with gr.Row():
+        # Use type 'binary' to receive file data as binary
+        pdf_upload = gr.File(label="Upload PDF", type="binary")
+        process_button = gr.Button("Process PDF")
+
+    output_text = gr.Textbox(label="Processing Status")
+    log_output = gr.Textbox(label="Log Output", interactive=False)
+
+    # Outputs: [status_message, log_output]
+    process_button.click(
+        fn=process_pdf, 
+        inputs=pdf_upload, 
+        outputs=[output_text, log_output]
+    )
+
+    query_input = gr.Textbox(label="Enter your query")
+    retrieve_button = gr.Button("Retrieve Context")
+    context_output = gr.Textbox(label="Retrieved Context")
+
+    retrieve_button.click(
+        fn=retrieve_context, 
+        inputs=query_input, 
+        outputs=context_output
+    )
+
+demo.launch()