semantic-router 0.1.0.dev4__tar.gz → 0.1.0.dev6__tar.gz

{semantic_router-0.1.0.dev4 → semantic_router-0.1.0.dev6}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: semantic-router
-Version: 0.1.0.dev4
+Version: 0.1.0.dev6
 Summary: Super fast semantic router for AI decision making
 License: MIT
 Author: Aurelio AI
@@ -40,7 +40,7 @@ Requires-Dist: numpy (>=1.25.2,<2.0.0)
 Requires-Dist: openai (>=1.10.0,<2.0.0)
 Requires-Dist: pillow (>=10.2.0,<11.0.0) ; (python_version < "3.13") and (extra == "vision")
 Requires-Dist: pinecone (>=5.0.0) ; extra == "pinecone"
-Requires-Dist: pinecone-text (>=0.7.1,<0.10.0) ; extra == "hybrid"
+Requires-Dist: pinecone-text (>=0.7.1) ; extra == "hybrid"
 Requires-Dist: psycopg2-binary (>=2.9.9,<3.0.0)
 Requires-Dist: pydantic (>=2.10.2,<3.0.0)
 Requires-Dist: pyyaml (>=6.0.1,<7.0.0)

{semantic_router-0.1.0.dev4 → semantic_router-0.1.0.dev6}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "semantic-router"
-version = "0.1.0.dev4"
+version = "0.1.0.dev6"
 description = "Super fast semantic router for AI decision making"
 authors = ["Aurelio AI <hello@aurelio.ai>"]
 readme = "README.md"
@@ -17,7 +17,7 @@ numpy = "^1.25.2"
 colorlog = "^6.8.0"
 pyyaml = "^6.0.1"
 aurelio-sdk = {version = "^0.0.16"}
-pinecone-text = {version = ">=0.7.1,<0.10.0", optional = true}
+pinecone-text = {version = ">=0.7.1", optional = true}
 torch = {version = ">=2.1.0,<2.6.0", optional = true, python = "<3.13" }
 transformers = {version = ">=4.36.2", optional = true, python = "<3.13" }
 tokenizers = {version = ">=0.19", optional = true, python = "<3.13" }

{semantic_router-0.1.0.dev4 → semantic_router-0.1.0.dev6}/semantic_router/__init__.py

@@ -3,4 +3,4 @@ from semantic_router.route import Route
 
 __all__ = ["SemanticRouter", "HybridRouter", "Route", "RouterConfig"]
 
-__version__ = "0.1.0.dev4"
+__version__ = "0.1.0.dev6"

{semantic_router-0.1.0.dev4 → semantic_router-0.1.0.dev6}/semantic_router/index/base.py

@@ -1,3 +1,4 @@
+import asyncio
 from datetime import datetime
 import time
 from typing import Any, List, Optional, Tuple, Union, Dict
@@ -11,6 +12,15 @@ from semantic_router.route import Route
 from semantic_router.utils.logger import logger
 
 
+RETRY_WAIT_TIME = 2.5
+
+
+class IndexConfig(BaseModel):
+    type: str
+    dimensions: int
+    vectors: int
+
+
 class BaseIndex(BaseModel):
     """
     Base class for indices using Pydantic's BaseModel.
@@ -34,26 +44,73 @@ class BaseIndex(BaseModel):
         utterances: List[Any],
         function_schemas: Optional[List[Dict[str, Any]]] = None,
         metadata_list: List[Dict[str, Any]] = [],
+        **kwargs,
     ):
-        """
-        Add embeddings to the index.
+        """Add embeddings to the index.
         This method should be implemented by subclasses.
         """
         raise NotImplementedError("This method should be implemented by subclasses.")
 
-    def get_utterances(self) -> List[Utterance]:
+    async def aadd(
+        self,
+        embeddings: List[List[float]],
+        routes: List[str],
+        utterances: List[str],
+        function_schemas: Optional[Optional[List[Dict[str, Any]]]] = None,
+        metadata_list: List[Dict[str, Any]] = [],
+        **kwargs,
+    ):
+        """Add vectors to the index asynchronously.
+        This method should be implemented by subclasses.
+        """
+        logger.warning("Async method not implemented.")
+        return self.add(
+            embeddings=embeddings,
+            routes=routes,
+            utterances=utterances,
+            function_schemas=function_schemas,
+            metadata_list=metadata_list,
+            **kwargs,
+        )
+
+    def get_utterances(self, include_metadata: bool = False) -> List[Utterance]:
         """Gets a list of route and utterance objects currently stored in the
         index, including additional metadata.
 
-        :return: A list of tuples, each containing route, utterance, function
-        schema and additional metadata.
-        :rtype: List[Tuple]
+        :param include_metadata: Whether to include function schemas and metadata in
+        the returned Utterance objects.
+        :type include_metadata: bool
+        :return: A list of Utterance objects.
+        :rtype: List[Utterance]
         """
         if self.index is None:
             logger.warning("Index is None, could not retrieve utterances.")
             return []
-        _, metadata = self._get_all(include_metadata=True)
+        _, metadata = self._get_all(include_metadata=True)  # include_metadata required
         route_tuples = parse_route_info(metadata=metadata)
+        if not include_metadata:
+            # we remove the metadata from the tuples (ie only keep 0, 1 items)
+            route_tuples = [x[:2] for x in route_tuples]
+        return [Utterance.from_tuple(x) for x in route_tuples]
+
+    async def aget_utterances(self, include_metadata: bool = False) -> List[Utterance]:
+        """Gets a list of route and utterance objects currently stored in the
+        index, including additional metadata.
+
+        :param include_metadata: Whether to include function schemas and metadata in
+        the returned Utterance objects.
+        :type include_metadata: bool
+        :return: A list of Utterance objects.
+        :rtype: List[Utterance]
+        """
+        if self.index is None:
+            logger.warning("Index is None, could not retrieve utterances.")
+            return []
+        _, metadata = await self._async_get_all(include_metadata=True)
+        route_tuples = parse_route_info(metadata=metadata)
+        if not include_metadata:
+            # we remove the metadata from the tuples (ie only keep 0, 1 items)
+            route_tuples = [x[:2] for x in route_tuples]
         return [Utterance.from_tuple(x) for x in route_tuples]
 
     def get_routes(self) -> List[Route]:
@@ -62,7 +119,7 @@ class BaseIndex(BaseModel):
         :return: A list of Route objects.
         :rtype: List[Route]
         """
-        utterances = self.get_utterances()
+        utterances = self.get_utterances(include_metadata=True)
         routes_dict: Dict[str, Route] = {}
         # first create a dictionary of route names to Route objects
         for utt in utterances:
@@ -90,6 +147,14 @@ class BaseIndex(BaseModel):
         """
         raise NotImplementedError("This method should be implemented by subclasses.")
 
+    async def _async_remove_and_sync(self, routes_to_delete: dict):
+        """
+        Remove embeddings in a routes syncing process from the index asynchronously.
+        This method should be implemented by subclasses.
+        """
+        logger.warning("Async method not implemented.")
+        return self._remove_and_sync(routes_to_delete=routes_to_delete)
+
     def delete(self, route_name: str):
         """
         Deletes route by route name.
@@ -97,10 +162,17 @@ class BaseIndex(BaseModel):
         """
         raise NotImplementedError("This method should be implemented by subclasses.")
 
-    def describe(self) -> Dict:
+    def describe(self) -> IndexConfig:
         """
-        Returns a dictionary with index details such as type, dimensions, and total
-        vector count.
+        Returns an IndexConfig object with index details such as type, dimensions, and
+        total vector count.
+        This method should be implemented by subclasses.
+        """
+        raise NotImplementedError("This method should be implemented by subclasses.")
+
+    def is_ready(self) -> bool:
+        """
+        Checks if the index is ready to be used.
         This method should be implemented by subclasses.
         """
         raise NotImplementedError("This method should be implemented by subclasses.")
@@ -159,6 +231,10 @@ class BaseIndex(BaseModel):
         logger.warning("This method should be implemented by subclasses.")
         self.index = None
 
+    # ___________________________ CONFIG ___________________________
+    # When implementing a new index, the following methods should be implemented
+    # to enable synchronization of remote indexes.
+
     def _read_config(self, field: str, scope: str | None = None) -> ConfigParameter:
         """Read a config parameter from the index.
 
@@ -176,13 +252,20 @@ class BaseIndex(BaseModel):
             scope=scope,
         )
 
-    def _read_hash(self) -> ConfigParameter:
-        """Read the hash of the previously written index.
+    async def _async_read_config(
+        self, field: str, scope: str | None = None
+    ) -> ConfigParameter:
+        """Read a config parameter from the index asynchronously.
 
+        :param field: The field to read.
+        :type field: str
+        :param scope: The scope to read.
+        :type scope: str | None
         :return: The config parameter that was read.
         :rtype: ConfigParameter
         """
-        return self._read_config(field="sr_hash")
+        logger.warning("_async_read_config method not implemented.")
+        return self._read_config(field=field, scope=scope)
 
     def _write_config(self, config: ConfigParameter) -> ConfigParameter:
         """Write a config parameter to the index.
@@ -195,6 +278,67 @@ class BaseIndex(BaseModel):
         logger.warning("This method should be implemented by subclasses.")
         return config
 
+    async def _async_write_config(self, config: ConfigParameter) -> ConfigParameter:
+        """Write a config parameter to the index asynchronously.
+
+        :param config: The config parameter to write.
+        :type config: ConfigParameter
+        :return: The config parameter that was written.
+        :rtype: ConfigParameter
+        """
+        logger.warning("Async method not implemented.")
+        return self._write_config(config=config)
+
+    # _________________________ END CONFIG _________________________
+
+    def _read_hash(self) -> ConfigParameter:
+        """Read the hash of the previously written index.
+
+        :return: The config parameter that was read.
+        :rtype: ConfigParameter
+        """
+        return self._read_config(field="sr_hash")
+
+    async def _async_read_hash(self) -> ConfigParameter:
+        """Read the hash of the previously written index asynchronously.
+
+        :return: The config parameter that was read.
+        :rtype: ConfigParameter
+        """
+        return await self._async_read_config(field="sr_hash")
+
+    def _is_locked(self, scope: str | None = None) -> bool:
+        """Check if the index is locked for a given scope (if applicable).
+
+        :param scope: The scope to check.
+        :type scope: str | None
+        :return: True if the index is locked, False otherwise.
+        :rtype: bool
+        """
+        lock_config = self._read_config(field="sr_lock", scope=scope)
+        if lock_config.value == "True":
+            return True
+        elif lock_config.value == "False" or not lock_config.value:
+            return False
+        else:
+            raise ValueError(f"Invalid lock value: {lock_config.value}")
+
+    async def _ais_locked(self, scope: str | None = None) -> bool:
+        """Check if the index is locked for a given scope (if applicable).
+
+        :param scope: The scope to check.
+        :type scope: str | None
+        :return: True if the index is locked, False otherwise.
+        :rtype: bool
+        """
+        lock_config = await self._async_read_config(field="sr_lock", scope=scope)
+        if lock_config.value == "True":
+            return True
+        elif lock_config.value == "False" or not lock_config.value:
+            return False
+        else:
+            raise ValueError(f"Invalid lock value: {lock_config.value}")
+
     def lock(
         self, value: bool, wait: int = 0, scope: str | None = None
     ) -> ConfigParameter:
@@ -214,9 +358,14 @@ class BaseIndex(BaseModel):
             if self._is_locked(scope=scope) != value:
                 # in this case, we can set the lock value
                 break
+            elif not value:
+                # if unlocking, we can break immediately — often with Pinecone the
+                # lock/unlocked state takes a few seconds to update, so locking then
+                # unlocking quickly will fail without this check
+                break
             if (datetime.now() - start_time).total_seconds() < wait:
-                # wait for 2.5 seconds before checking again
-                time.sleep(2.5)
+                # wait for a few seconds before checking again
+                time.sleep(RETRY_WAIT_TIME)
             else:
                 raise ValueError(
                     f"Index is already {'locked' if value else 'unlocked'}."
@@ -229,21 +378,31 @@ class BaseIndex(BaseModel):
         self._write_config(lock_param)
         return lock_param
 
-    def _is_locked(self, scope: str | None = None) -> bool:
-        """Check if the index is locked for a given scope (if applicable).
-
-        :param scope: The scope to check.
-        :type scope: str | None
-        :return: True if the index is locked, False otherwise.
-        :rtype: bool
+    async def alock(
+        self, value: bool, wait: int = 0, scope: str | None = None
+    ) -> ConfigParameter:
+        """Lock/unlock the index for a given scope (if applicable). If index
+        already locked/unlocked, raises ValueError.
         """
-        lock_config = self._read_config(field="sr_lock", scope=scope)
-        if lock_config.value == "True":
-            return True
-        elif lock_config.value == "False" or not lock_config.value:
-            return False
-        else:
-            raise ValueError(f"Invalid lock value: {lock_config.value}")
+        start_time = datetime.now()
+        while True:
+            if await self._ais_locked(scope=scope) != value:
+                # in this case, we can set the lock value
+                break
+            if (datetime.now() - start_time).total_seconds() < wait:
+                # wait for a few seconds before checking again
+                await asyncio.sleep(RETRY_WAIT_TIME)
+            else:
+                raise ValueError(
+                    f"Index is already {'locked' if value else 'unlocked'}."
+                )
+        lock_param = ConfigParameter(
+            field="sr_lock",
+            value=str(value),
+            scope=scope,
+        )
+        await self._async_write_config(lock_param)
+        return lock_param
 
     def _get_all(self, prefix: Optional[str] = None, include_metadata: bool = False):
         """

{semantic_router-0.1.0.dev4 → semantic_router-0.1.0.dev6}/semantic_router/index/hybrid_local.py

@@ -25,6 +25,7 @@ class HybridLocalIndex(LocalIndex):
         function_schemas: Optional[List[Dict[str, Any]]] = None,
         metadata_list: List[Dict[str, Any]] = [],
         sparse_embeddings: Optional[List[SparseEmbedding]] = None,
+        **kwargs,
     ):
         if sparse_embeddings is None:
             raise ValueError("Sparse embeddings are required for HybridLocalIndex.")
@@ -56,23 +57,20 @@ class HybridLocalIndex(LocalIndex):
             self.routes = np.concatenate([self.routes, routes_arr])
             self.utterances = np.concatenate([self.utterances, utterances_arr])
 
-    def get_utterances(self) -> List[Utterance]:
+    def get_utterances(self, include_metadata: bool = False) -> List[Utterance]:
         """Gets a list of route and utterance objects currently stored in the index.
 
-        Returns:
-            List[Tuple]: A list of (route_name, utterance) objects.
+        :param include_metadata: Whether to include function schemas and metadata in
+        the returned Utterance objects - HybridLocalIndex doesn't include metadata so
+        this parameter is ignored.
+        :type include_metadata: bool
+        :return: A list of Utterance objects.
+        :rtype: List[Utterance]
         """
         if self.routes is None or self.utterances is None:
             return []
         return [Utterance.from_tuple(x) for x in zip(self.routes, self.utterances)]
 
-    def describe(self) -> Dict:
-        return {
-            "type": self.type,
-            "dimensions": self.index.shape[1] if self.index is not None else 0,
-            "vectors": self.index.shape[0] if self.index is not None else 0,
-        }
-
     def _sparse_dot_product(
         self, vec_a: dict[int, float], vec_b: dict[int, float]
     ) -> float:

{semantic_router-0.1.0.dev4 → semantic_router-0.1.0.dev6}/semantic_router/index/local.py

@@ -3,7 +3,7 @@ from typing import List, Optional, Tuple, Dict
 import numpy as np
 
 from semantic_router.schema import ConfigParameter, SparseEmbedding, Utterance
-from semantic_router.index.base import BaseIndex
+from semantic_router.index.base import BaseIndex, IndexConfig
 from semantic_router.linear import similarity_matrix, top_scores
 from semantic_router.utils.logger import logger
 from typing import Any
@@ -26,6 +26,7 @@ class LocalIndex(BaseIndex):
         utterances: List[str],
         function_schemas: Optional[List[Dict[str, Any]]] = None,
         metadata_list: List[Dict[str, Any]] = [],
+        **kwargs,
     ):
         embeds = np.array(embeddings)  # type: ignore
         routes_arr = np.array(routes)
@@ -63,23 +64,31 @@ class LocalIndex(BaseIndex):
         # return what was removed
         return route_utterances[~mask]
 
-    def get_utterances(self) -> List[Utterance]:
-        """
-        Gets a list of route and utterance objects currently stored in the index.
+    def get_utterances(self, include_metadata: bool = False) -> List[Utterance]:
+        """Gets a list of route and utterance objects currently stored in the index.
 
-        Returns:
-            List[Tuple]: A list of (route_name, utterance) objects.
+        :param include_metadata: Whether to include function schemas and metadata in
+        the returned Utterance objects - LocalIndex doesn't include metadata so this
+        parameter is ignored.
+        :return: A list of Utterance objects.
+        :rtype: List[Utterance]
         """
         if self.routes is None or self.utterances is None:
             return []
         return [Utterance.from_tuple(x) for x in zip(self.routes, self.utterances)]
 
-    def describe(self) -> Dict:
-        return {
-            "type": self.type,
-            "dimensions": self.index.shape[1] if self.index is not None else 0,
-            "vectors": self.index.shape[0] if self.index is not None else 0,
-        }
+    def describe(self) -> IndexConfig:
+        return IndexConfig(
+            type=self.type,
+            dimensions=self.index.shape[1] if self.index is not None else 0,
+            vectors=self.index.shape[0] if self.index is not None else 0,
+        )
+
+    def is_ready(self) -> bool:
+        """
+        Checks if the index is ready to be used.
+        """
+        return self.index is not None and self.routes is not None
 
     def query(
         self,