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,65 @@
     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 = {}
+
+    @abstractmethod
+    def build_tool(self, component: str, uuid: str):
+        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'):
+        tool = self.build_tool(component)
+        tool.db = self.db
+        self.component_uuid_mapping[component.component].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()
+
+    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 +94,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/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/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)