Update app.py

app.py

@@ -83,11 +83,7 @@ def init_astra_db():
         astra_db_endpoint = os.getenv("ASTRA_DB_ENDPOINT", "https://8e3fd85c-5f28-4e1f-8538-9dd28a3ea2b0-us-east-2.apps.astra.datastax.com")
         
         # Initialize the client
-        client = DataAPIClient(astra_db_application_token)
-        db = client.get_database_by_api_endpoint(
-            astra_db_endpoint,
-            keyspace=astra_db_keyspace
-        )
+        client = DataAPIClient(api_endpoint=astra_db_endpoint, token=astra_db_application_token)
         
         # Get or create collections
         product_embeddings = db.get_collection("product_embeddings")
@@ -259,6 +255,57 @@ def process_pdf_catalogs():
         print(f"Error processing PDF catalogs: {e}")
         return {"status": "error", "message": str(e)}
 
+# Add this function to process PDFs from URLs
+def process_pdf_from_url(url):
+    """Download and process a PDF from a URL"""
+    try:
+        # Download the PDF
+        response = requests.get(url, stream=True)
+        if response.status_code != 200:
+            return f"Error downloading PDF: HTTP status code {response.status_code}"
+        
+        # Get the content
+        pdf_content = response.content
+        
+        # Determine product type from URL or filename
+        product_type = "other"
+        for pt in ["circuit_breaker", "motor_starter", "contactor", "switch", "relay"]:
+            if pt in url.lower():
+                product_type = pt.replace("_", " ")
+                break
+        
+        # Process PDF text content
+        pdf_reader = PyPDF2.PdfReader(io.BytesIO(pdf_content))
+        text_content = ""
+        
+        # Extract text from each page
+        for page in pdf_reader.pages:
+            text_content += page.extract_text() + "\n\n"
+        
+        # Split text into smaller chunks for efficient embedding
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=1000,
+            chunk_overlap=200,
+            length_function=len,
+        )
+        chunks = text_splitter.split_text(text_content)
+        
+        # Store chunks in vector database (if available)
+        if astra_session:
+            store_chunks_in_db(chunks, product_type)
+        
+        # Extract and store images (if database available)
+        images_count = 0
+        if astra_session:
+            images_count = extract_images_from_pdf(pdf_content, product_type)
+        
+        print(f"Processed PDF from URL: {url}: {len(chunks)} text chunks and {images_count} images extracted")
+        return f"Successfully processed PDF from URL: {len(chunks)} chunks, {images_count} images"
+    
+    except Exception as e:
+        print(f"Error processing PDF from URL: {e}")
+        return f"Error processing PDF: {str(e)}"        
+
 # Function to store text chunks in Astra DB with embeddings
 def store_chunks_in_db(chunks, product_type):
     """Store text chunks with embeddings in Astra DB"""
@@ -648,22 +695,30 @@ def setup_and_update():
     
     # Initialize database and other services
     global astra_session, astra_keyspace, s3_client, embeddings_model
-    astra_db_result = init_astra_db()
-    if astra_db_result:
-        astra_session = astra_db_result.get("db")
-        astra_keyspace = astra_db_result.get("keyspace")
+    astra_result = init_astra_db()
+    
+    if astra_result:
+        astra_session = astra_result.get("db")
+        astra_keyspace = astra_result.get("keyspace")
     else:
         astra_session = None
         astra_keyspace = None
-    
+        
     s3_client = init_s3_client()
     embeddings_model = get_embeddings_model()
     
     # Return status
-    if openai_initialized and mistral_initialized:
-        return "System is ready. You can start chatting!"
-    else:
-        return "System initialization incomplete. Some features may not work properly."
+    status_msg = "System is ready. "
+    if not openai_initialized:
+        status_msg += "OpenAI API not initialized. "
+    if not mistral_initialized:
+        status_msg += "Mistral API not initialized. "
+    if not astra_session:
+        status_msg += "Astra DB not connected. "
+    if not s3_client:
+        status_msg += "S3 client not initialized. "
+    
+    return status_msg
 
 def create_gradio_app():
     # Define CSS styles for a more modern, appealing interface
@@ -824,7 +879,24 @@ def create_gradio_app():
                             s3_bucket = gr.Textbox(label="S3 Bucket Name")
                             s3_prefix = gr.Textbox(label="S3 Prefix (folder)", value="catalogs/")
                             process_btn = gr.Button("Process PDFs from S3", elem_classes="action-button")
+
+                            # Add direct PDF URL input
+                            with gr.Tab("Direct PDF URLs"):
+                                pdf_url = gr.Textbox(label="PDF URL", placeholder="https://example.com/sample.pdf")
+                                pdf_dropdown = gr.Dropdown(
+                                    label="ABB Catalog PDFs",
+                                    choices=[
+                                        "https://agent-product-discovery.s3.ap-south-1.amazonaws.com/ABB-catalog/ABB+Ability%E2%84%A2+System+800xA%C2%AE+6.2.pdf",
+                                        "https://agent-product-discovery.s3.ap-south-1.amazonaws.com/ABB-catalog/Enclosed+Softstarters.pdf",
+                                        "https://agent-product-discovery.s3.ap-south-1.amazonaws.com/ABB-catalog/Ex-Solutions.pdf",
+                                        "https://agent-product-discovery.s3.ap-south-1.amazonaws.com/ABB-catalog/Low_power_UPS_catalogue_EN.pdf"
+                                    ],
+                                    interactive=True
+                                )
+                                process_url_btn = gr.Button("Process PDF from URL", elem_classes="action-button")
+
                             result_text = gr.Textbox(label="Processing Result")
+    
 
             # Set up event handlers
             send_btn.click(
@@ -855,6 +927,22 @@ def create_gradio_app():
                 api_name="process_pdfs"
             )
 
+            # Add this event handler
+            process_url_btn.click(
+            process_pdf_from_url,
+            [pdf_url],
+            [result_text],
+            api_name="process_pdf_url"
+            )
+
+            # Add this dropdown change event
+            pdf_dropdown.change(
+            lambda x: x,
+            [pdf_dropdown],
+            [pdf_url],
+            api_name="update_pdf_url"
+            )
+
         # Add the system setup to run when the app loads
         app.load(setup_and_update, None, status_display)