Update app.py

app.py

@@ -79,16 +79,14 @@ def init_astra_db():
         astra_db_region = os.getenv("ASTRA_DB_REGION")
         astra_db_keyspace = os.getenv("ASTRA_DB_KEYSPACE")
         astra_db_application_token = os.getenv("ASTRA_DB_APPLICATION_TOKEN")
+        astra_db_endpoint = os.getenv("ASTRA_DB_ENDPOINT")
         
-        # Setup the connection
-        cloud_config = {
-            'secure_connect_bundle': 'secure-connect-' + astra_db_id + '.zip'
-        }
-        
-        auth_provider = PlainTextAuthProvider(
-            'token', astra_db_application_token)
-        cluster = Cluster(cloud=cloud_config, auth_provider=auth_provider)
-        session = cluster.connect()
+        # Initialize the client
+        client = DataAPIClient(astra_db_token)
+        db = client.get_database_by_api_endpoint(
+            astra_db_endpoint,
+            keyspace=astra_db_keyspace
+        )
         
         # Create keyspace if it doesn't exist
         session.execute(f"""
@@ -107,35 +105,43 @@ def init_astra_db():
             )
         """)
         
-        # Create table for query analytics
-        session.execute(f"""
-            CREATE TABLE IF NOT EXISTS {astra_db_keyspace}.query_analytics (
-                id text PRIMARY KEY,
-                query text,
-                product_type text,
-                timestamp timestamp,
-                response_time float
-            )
-        """)
-        
-        # Create table for product images
-        session.execute(f"""
-            CREATE TABLE IF NOT EXISTS {astra_db_keyspace}.product_images (
-                id text PRIMARY KEY,
-                product_type text,
-                image_data blob,
-                page_number int,
-                image_index int,
-                metadata text
-            )
-        """)
+        # Get or create collections
+        product_embeddings = db.get_collection("product_embeddings")
+        query_analytics = db.get_collection("query_analytics")
+        product_images = db.get_collection("product_images")
         
-        print("Astra DB connection established")
-        return session, astra_db_keyspace
+        # Create vector index for product_embeddings if it doesn't exist
+        try:
+            # Check if vector search is already set up
+            vector_indexes = product_embeddings.get_vector_indexes()
+            has_vector_index = any(index.get("name") == "embedding_vector_index" for index in vector_indexes)
+            
+            if not has_vector_index:
+                # Create vector index
+                product_embeddings.create_vector_index(
+                    vector_field="$vector", 
+                    dimension=1536,  # Adjust dimension based on your embedding model
+                    index_name="embedding_vector_index"
+                )
+                print("Vector index created for product_embeddings collection")
+        except Exception as e:
+            print(f"Warning: Could not create vector index: {e}")
+        
+        print(f"Connected to Astra DB: {db.list_collection_names()}")
+        
+        # Return DB client and collections for use in application
+        return {
+            "db": db,
+            "keyspace": astra_db_keyspace,
+            "collections": {
+                "product_embeddings": product_embeddings,
+                "query_analytics": query_analytics,
+                "product_images": product_images
+            }
+        }
     except Exception as e:
         print(f"Error connecting to Astra DB: {e}")
-        # Return None values to allow the app to run without DB connection
-        return None, None
+        return None
 
 # Initialize AWS S3 client for accessing product catalogs
 def init_s3_client():