Update lancedb_client.py

lancedb_client.py

@@ -1,100 +1,129 @@
-import numpy as np
-import lancedb
-import pyarrow as pa
-import logging
-from dotenv import load_dotenv
-import os
-import ast
-
-# Load env vars
-load_dotenv(os.path.join(os.getcwd(), ".env"),override = True)
-metadata_keys_raw = os.getenv("_DEFAULT_PARSE_METADATA", "").split(",")
-metadata_keys = [key.replace(" ", "").replace(")", "").strip("'") for key in metadata_keys_raw]
-
-# Setup logger
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-class LanceDBManager:
-
-    def __init__(self, db_uri="lancedb", embedding_dim=1536):
-        self.db = lancedb.connect(db_uri)
-        self.embedding_dim = embedding_dim
-        logger.info(f"Connected to LanceDB at {db_uri}")
-
-    def _build_schema(self):
-        """Build LanceDB schema with dynamic metadata fields and embedding vector."""
-        fields = [
-            pa.field("id", pa.int64()),  
-            pa.field("item_id", pa.string()),           
-            pa.field("images_urls", pa.string()),
-            pa.field("text", pa.string()),
-            pa.field("Cluster", pa.string()),
-            pa.field("Topic", pa.string()),
-            pa.field("embeddings", pa.list_(pa.float32(), self.embedding_dim)),
-            pa.field("umap_embeddings", pa.list_(pa.float32(), 2)),           
-        ]
-
-        # Add fields from metadata
-        for key in metadata_keys:
-            sanitized_key = key.split(":")[1].strip().capitalize() # remove the vocabulary prefix in key label and capitalize
-            fields.append(pa.field(sanitized_key, pa.string()))
-
-        return pa.schema(fields)
-
-    def create_table(self, table_name):
-        """Create table using dynamic schema."""
-        try:
-            schema = self._build_schema()
-            table = self.db.create_table(table_name, schema=schema)
-            logger.info(f"Created LanceDB table '{table_name}'")
-            return table
-        except Exception as e:
-            logger.error(f"Failed to create table '{table_name}': {e}")
-            raise
-
-    def retrieve_table(self, table_name):
-        try:
-            table = self.db.open_table(table_name)
-            logger.info(f"Opened existing LanceDB table '{table_name}'")
-            return table
-        except Exception as e:
-            logger.error(f"Failed to open table '{table_name}': {e}")
-            raise
-
-    def initialize_table(self, table_name):
-        try:
-            return self.retrieve_table(table_name)
-        except Exception:
-            logger.info(f"Table '{table_name}' not found. Creating new.")
-            return self.create_table(table_name)
-
-    def add_entry(self, table_name, items):
-        table = self.initialize_table(table_name)
-        table.add(items)
-        logger.info(f"Added items to table '{table_name}'")
-        
-    def list_tables(self):
-        """List all existing tables in the LanceDB instance."""
-        try:
-            tables = self.db.table_names()
-            logger.info("Retrieved list of tables.")
-            return tables
-        except Exception as e:
-            logger.error(f"Failed to list tables: {e}")
-            raise
-        
-    def get_content_table(self, table_name):
-        table = self.initialize_table(table_name)
-        return table.to_pandas()
-    
-    def drop_table(self, table_name):
-        """remove an existing table by name."""
-        try:
-            table = self.db.drop_table(table_name)
-            logger.info(f"Remove existing LanceDB table '{table_name}' successfully.")
-            return table
-        except Exception as e:
-            logger.error(f"Failed to remove existing table '{table_name}': {e}")
-            raise
+import numpy as np
+import lancedb
+import pyarrow as pa
+import logging
+from dotenv import load_dotenv
+import os
+import ast
+
+# Load env vars
+load_dotenv(os.path.join(os.getcwd(), ".env"),override = True)
+metadata_keys_raw = os.getenv("_DEFAULT_PARSE_METADATA", "").split(",")
+metadata_keys = [key.replace(" ", "").replace(")", "").strip("'") for key in metadata_keys_raw]
+
+# Setup logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class LanceDBManager:
+
+    def __init__(self, db_uri="lancedb", embedding_dim=768):
+        self.db = lancedb.connect(db_uri)
+        self.embedding_dim = embedding_dim
+        logger.info(f"Connected to LanceDB at {db_uri}")
+
+    def _build_schema(self):
+        """Build LanceDB schema with dynamic metadata fields and embedding vector."""
+        fields = [
+            pa.field("id", pa.int64()),  
+            pa.field("item_id", pa.string()),           
+            pa.field("images_urls", pa.string()),
+            pa.field("text", pa.string()),
+            pa.field("Cluster", pa.string()),
+            pa.field("Topic", pa.string()),
+            pa.field("embeddings", pa.list_(pa.float32(), self.embedding_dim)),
+            pa.field("umap_embeddings", pa.list_(pa.float32(), 2)),           
+        ]
+
+        # Add fields from metadata
+        for key in metadata_keys:
+            sanitized_key = key.split(":")[1].strip().capitalize() # remove the vocabulary prefix in key label and capitalize
+            fields.append(pa.field(sanitized_key, pa.string()))
+
+        return pa.schema(fields)
+
+    def create_table(self, table_name):
+        """Create table using dynamic schema."""
+        try:
+            schema = self._build_schema()
+            table = self.db.create_table(table_name, schema=schema)
+            logger.info(f"Created LanceDB table '{table_name}'")
+            return table
+        except Exception as e:
+            logger.error(f"Failed to create table '{table_name}': {e}")
+            raise
+
+    def retrieve_table(self, table_name):
+        try:
+            table = self.db.open_table(table_name)
+            logger.info(f"Opened existing LanceDB table '{table_name}'")
+            return table
+        except Exception as e:
+            logger.error(f"Failed to open table '{table_name}': {e}")
+            raise
+
+    def initialize_table(self, table_name):
+        try:
+            return self.retrieve_table(table_name)
+        except Exception:
+            logger.info(f"Table '{table_name}' not found. Creating new.")
+            return self.create_table(table_name)
+
+    def add_entry(self, table_name, items):
+        table = self.initialize_table(table_name)
+        table.add(items)
+        logger.info(f"Added items to table '{table_name}'")
+        
+    def list_tables(self):
+        """List all existing tables in the LanceDB instance."""
+        try:
+            tables = self.db.table_names()
+            logger.info("Retrieved list of tables.")
+            return tables
+        except Exception as e:
+            logger.error(f"Failed to list tables: {e}")
+            raise
+        
+    def get_content_table(self, table_name):
+        table = self.initialize_table(table_name)
+        return table.to_pandas()
+    
+    def drop_table(self, table_name):
+        """remove an existing table by name."""
+        try:
+            table = self.db.drop_table(table_name)
+            logger.info(f"Remove existing LanceDB table '{table_name}' successfully.")
+            return table
+        except Exception as e:
+            logger.error(f"Failed to remove existing table '{table_name}': {e}")
+            raise
+        
+    def semantic_search(self, table_name, query_embed, limit):
+        """
+        Perform a semantic search using a provided text query or image.
+
+        Args:
+            query_text (str): The text query for the search.
+            query_image_path (str): The path to the image for the search.
+            limit (int): The maximum number of results to return.
+
+        Returns:
+            str: JSON string of search results.
+        """
+        table = self.initialize_table(table_name)
+        #https://lancedb.github.io/lancedb/notebooks/DisappearingEmbeddingFunction/
+
+        try:
+            # Perform the search in LanceDB
+            results = (table
+                    .search(query_embed,vector_column_name="embeddings")
+                    .distance_type("cosine")
+                    .select(["id"])
+                    .limit(limit)
+                    .to_pandas()
+                    #.sort_values(by='_distance', ascending=True)
+                    .to_json(orient="records")
+            )
+            return results
+        except Exception as e:
+            raise