Greatly simplify the template saving logic

Author: blythed

superduper/__init__.py

@@ -35,7 +35,7 @@
 from .components.plugin import Plugin
 from .components.streamlit import Streamlit
 from .components.table import Table
-from .components.template import QueryTemplate, Template
+from .components.template import Template
 from .components.vector_index import VectorIndex
 
 REQUIRES = [
@@ -65,7 +65,6 @@
     'Table',
     'Application',
     'Template',
-    'QueryTemplate',
     'Application',
     'Component',
     'pickle_serializer',

superduper/backends/base/backends.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 import typing as t
 from abc import ABC, abstractmethod
 
@@ -6,6 +7,64 @@
     from superduper.components.component import Component
 
 
+class Bookkeeping(ABC):
+    def __init__(self):
+        self.component_uuid_mapping = defaultdict(set)
+        self.uuid_component_mapping = {}
+        self.tool_uuid_mapping = defaultdict(set)
+        self.uuid_tool_mapping = {}
+        self.tools = {}
+
+    def build_tool(self, component: 'Component'):
+        pass
+
+    def get_tool(self, uuid: str):
+        tool_id = self.uuid_tool_mapping[uuid]
+        return self.tools[tool_id]
+
+    def put_component(self, component: 'Component', **kwargs):
+        tool = self.build_tool(component)
+        tool.db = self.db
+        self.component_uuid_mapping[(component.component, component.identifier)].add(component.uuid)
+        self.uuid_component_mapping[component.uuid] = (component.component, component.identifier)
+        self.uuid_tool_mapping[component.uuid] = tool.identifier
+        self.tool_uuid_mapping[tool.identifier].add(component.uuid)
+        self.tools[tool.identifier] = tool
+        tool.initialize(**kwargs)
+
+    def drop_component(self, component: str, identifier: str):
+        uuids = self.component_uuid_mapping[(component, identifier)]
+        tool_ids = []
+        for uuid in uuids:
+            del self.uuid_component_mapping[uuid]
+            tool_id = self.uuid_tool_mapping[uuid]
+            tool_ids.append(tool_id)
+            del self.uuid_tool_mapping[uuid]
+            self.tool_uuid_mapping[tool_id].remove(uuid)
+            if not self.tool_uuid_mapping[tool_id]:
+                self.tools.drop()
+                del self.tools[tool_id]
+        del self.component_uuid_mapping[(component, identifier)]
+
+    def drop(self):
+        for tool in self.tools.values():
+            tool.drop()
+        self.component_uuid_mapping = defaultdict(set)
+        self.uuid_component_mapping = {}
+        self.tool_uuid_mapping = defaultdict(set)
+        self.uuid_tool_mapping = {}
+        self.tools = {}
+
+    def list_components(self):
+        return list(self.component_uuid_mapping.keys())
+
+    def list_tools(self):
+        return list(self.tools.keys())
+
+    def list_uuids(self):
+        return list(self.uuid_component_mapping.keys())
+
+
 class BaseBackend(ABC):
     """Base backend class for cluster client."""
 
@@ -34,28 +93,13 @@ def initialize(self):
         """To be called on program start."""
         pass
 
-    def put_component(self, component: 'Component', **kwargs):
+    @abstractmethod
+    def put_component(self, component: 'Component'):
         """Add a component to the deployment.
 
         :param component: ``Component`` to put.
         :param kwargs: kwargs dictionary.
         """
-        # This is to make sure that we only have 1 version
-        # of each component implemented at any given time
-        # TODO: get identifier in string component argument.
-        identifier = ''
-        if isinstance(component, str):
-            uuid = component
-        else:
-            uuid = component.uuid
-            identifier = component.identifier
-
-        if uuid in self.list_uuids():
-            return
-        if identifier in self.list_components():
-            del self[component.identifier]
-
-        self._put(component, **kwargs)
 
     @abstractmethod
     def drop_component(self, component: str, identifier: str):

superduper/backends/base/scheduler.py

@@ -19,59 +19,6 @@
 BATCH_SIZE = 100
 
 
-def _chunked_list(lst, batch_size=BATCH_SIZE):
-    if len(lst) <= batch_size:
-        return [lst]
-    return [lst[i : i + batch_size] for i in range(0, len(lst), batch_size)]
-
-
-class BaseQueueConsumer(ABC):
-    """
-    Base class for handling consumer process.
-
-    This class is an implementation of message broker between
-    producers (superduper db client) and consumers i.e listeners.
-
-    :param uri: Uri to connect.
-    :param queue_name: Queue to consume.
-    :param callback: Callback for consumed messages.
-    """
-
-    def __init__(
-        self,
-        queue_name: str = '',
-        callback: t.Optional[t.Callable] = None,
-    ):
-        self.callback = callback
-        self.queue_name = queue_name
-        self.futures: t.DefaultDict = defaultdict(lambda: {})
-
-    @abstractmethod
-    def start_consuming(self):
-        """Abstract method to start consuming messages."""
-        pass
-
-    @abstractmethod
-    def close_connection(self):
-        """Abstract method to close connection."""
-        pass
-
-    def consume(self, *args, **kwargs):
-        """Start consuming messages from queue.
-
-        :param args: positional arguments
-        :param kwargs: keyword arguments
-        """
-        logging.info(f"Started consuming on queue: {self.queue_name}")
-        try:
-            self.start_consuming()
-        except KeyboardInterrupt:
-            logging.info("KeyboardInterrupt: Stopping consumer...")
-        finally:
-            self.close_connection()
-            logging.info(f"Stopped consuming on queue: {self.queue_name}")
-
-
 class BaseScheduler(BaseBackend):
     """
     Base class for handling publisher and consumer process.
@@ -82,10 +29,6 @@ class BaseScheduler(BaseBackend):
     :param uri: Uri to connect.
     """
 
-    def __init__(self):
-        super().__init__()
-        self.queue: t.Dict = defaultdict(lambda: [])
-
     @abstractmethod
     def publish(self, events: t.List[Event]):
         """

superduper/backends/base/vector_search.py

@@ -6,38 +6,35 @@
 import numpy
 import numpy.typing
 
-from superduper.backends.base.backends import BaseBackend
+from superduper.backends.base.backends import BaseBackend, Bookkeeping
 
 if t.TYPE_CHECKING:
     from superduper.base.datalayer import Datalayer
-    from superduper.components.vector_index import VectorIndex
+    from superduper.components.vector_index import VectorIndex, VectorItem
 
 
-class VectorSearchBackend(BaseBackend):
+class VectorSearchBackend(Bookkeeping, BaseBackend):
     """Base vector-search backend."""
 
     def __init__(self):
-        self._cache = {}
+        Bookkeeping.__init__(self)
+        BaseBackend.__init__(self)
 
-    @abstractmethod
-    def __getitem__(self, identifier):
-        pass
-
-    def add(self, identifier, vectors):
+    def add(self, uuid: str, vectors: t.List['VectorItem']):
         """Add vectors to a vector-index.
 
         :param identifier: Identifier of index.
         :param vectors: Vectors.
         """
-        self.get(identifier).add(vectors)
+        self.get_tool(uuid).add(vectors)
 
-    def delete(self, identifier, ids):
+    def delete(self, uuid, ids):
         """Delete ids from index.
 
         :param identifier: Identifier of index.
         :param ids: Ids to delete.
         """
-        self.get(identifier).delete(ids)
+        self.get_tool(uuid).delete(ids)
 
     @property
     def db(self) -> 'Datalayer':
@@ -52,13 +49,48 @@ def db(self, value: 'Datalayer'):
         """
         self._db = value
 
+    @abstractmethod
+    def find_nearest_from_array(
+        self,
+        h: numpy.typing.ArrayLike,
+        vector_index: str,
+        n: int = 100,
+        within_ids: t.Sequence[str] = (),
+    ) -> t.Tuple[t.List[str], t.List[float]]:
+        """
+        Find the nearest vectors to the given vector.
+
+        :param h: vector
+        :param n: number of nearest vectors to return
+        :param within_ids: list of ids to search within
+        """
+
+    @abstractmethod
+    def find_nearest_from_id(
+        self,
+        id: str,
+        vector_index: str,
+        n: int = 100,
+        within_ids: t.Sequence[str] = (),
+    ) -> t.Tuple[t.List[str], t.List[float]]:
+        """
+        Find the nearest vectors to the given vector.
+
+        :param id: id of the vector to search with
+        :param n: number of nearest vectors to return
+        :param within_ids: list of ids to search within
+        """
+
 
 class VectorSearcherInterface(ABC):
     """Interface for vector searchers.
 
     # noqa
     """
 
+    def __init__(self, identifier: str):
+        self.identifier = identifier
+
     @abstractmethod
     def add(self, items: t.Sequence['VectorItem']) -> None:
         """
@@ -75,31 +107,31 @@ def delete(self, ids: t.Sequence[str]) -> None:
         """
 
     @abstractmethod
-    def find_nearest_from_id(
+    def find_nearest_from_array(
         self,
-        _id,
+        h: numpy.typing.ArrayLike,
         n: int = 100,
         within_ids: t.Sequence[str] = (),
     ) -> t.Tuple[t.List[str], t.List[float]]:
         """
-        Find the nearest vectors to the vector with the given id.
+        Find the nearest vectors to the given vector.
 
-        :param _id: id of the vector
+        :param h: vector
         :param n: number of nearest vectors to return
         :param within_ids: list of ids to search within
         """
 
     @abstractmethod
-    def find_nearest_from_array(
+    def find_nearest_from_id(
         self,
-        h: numpy.typing.ArrayLike,
+        id: str,
         n: int = 100,
         within_ids: t.Sequence[str] = (),
     ) -> t.Tuple[t.List[str], t.List[float]]:
         """
         Find the nearest vectors to the given vector.
 
-        :param h: vector
+        :param id: id of the vector to search with
         :param n: number of nearest vectors to return
         :param within_ids: list of ids to search within
         """
@@ -111,6 +143,13 @@ def post_create(self):
         to perform a task after all vectors have been added
         """
 
+    def post_create(self):
+        """Post create method.
+
+        This method is used for searchers which requires
+        to perform a task after all vectors have been added
+        """
+
 
 class BaseVectorSearcher(VectorSearcherInterface):
     """Base class for vector searchers.
@@ -122,10 +161,18 @@ class BaseVectorSearcher(VectorSearcherInterface):
 
     native_service: t.ClassVar[bool] = True
 
+    @property
+    def db(self) -> 'Datalayer':
+        return self._db
+
+    @db.setter
+    def db(self, value: 'Datalayer'):
+        self._db = value
+
     @abstractmethod
     def __init__(
         self,
-        uuid: str,
+        identifier: str,
         dimensions: int,
         measure: str,
     ):
@@ -137,7 +184,7 @@ def from_component(cls, index: 'VectorIndex'):
 
         :param vi: ``VectorIndex`` instance
         """
-        return cls(uuid=index.uuid, dimensions=index.dimensions, measure=index.measure)
+        return cls(identifier=index.uuid, dimensions=index.dimensions, measure=index.measure)
 
     @abstractmethod
     def initialize(self, db):

superduper/backends/local/cache.py

@@ -20,10 +20,6 @@ def __delitem__(self, item):
         del self._cache[item]
 
     def __setitem__(self, key, value):
-        if isinstance(value, str):
-            import pdb
-
-            pdb.set_trace()
         self._cache[key] = copy.deepcopy(value)
 
     def __getitem__(self, item):

superduper/backends/local/cdc.py

@@ -32,9 +32,9 @@ def list_uuids(self):
         """List UUIDs of components."""
         return list(self._trigger_uuid_mapping.values())
 
-    def _put(self, item):
-        assert isinstance(item, CDC)
-        self.triggers.add((item.component, item.identifier))
+    def put_component(self, component):
+        assert isinstance(component, CDC)
+        self.triggers.add((component.component, component.identifier))
 
     def drop_component(self, component, identifier):
         c = self.db.load(component=component, identifier=identifier)