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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: semantic-router
-Version: 0.1.0.dev3
+Version: 0.1.0.dev4
 Summary: Super fast semantic router for AI decision making
 License: MIT
 Author: Aurelio AI

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

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "semantic-router"
-version = "0.1.0.dev3"
+version = "0.1.0.dev4"
 description = "Super fast semantic router for AI decision making"
 authors = ["Aurelio AI <hello@aurelio.ai>"]
 readme = "README.md"

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

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

{semantic_router-0.1.0.dev3 → semantic_router-0.1.0.dev4}/semantic_router/encoders/bedrock.py

@@ -17,7 +17,7 @@ Classes:
 """
 
 import json
-from typing import List, Optional, Any
+from typing import Dict, List, Optional, Any, Union
 import os
 from time import sleep
 import tiktoken
@@ -138,11 +138,14 @@ class BedrockEncoder(DenseEncoder):
             ) from err
         return bedrock_client
 
-    def __call__(self, docs: List[str]) -> List[List[float]]:
+    def __call__(
+        self, docs: List[Union[str, Dict]], model_kwargs: Optional[Dict] = None
+    ) -> List[List[float]]:
         """Generates embeddings for the given documents.
 
         Args:
             docs: A list of strings representing the documents to embed.
+            model_kwargs: A dictionary of model-specific inference parameters.
 
         Returns:
             A list of lists, where each inner list contains the embedding values for a
@@ -168,13 +171,29 @@ class BedrockEncoder(DenseEncoder):
                 embeddings = []
                 if self.name and "amazon" in self.name:
                     for doc in docs:
-                        embedding_body = json.dumps(
-                            {
-                                "inputText": doc,
-                            }
-                        )
+
+                        embedding_body = {}
+
+                        if isinstance(doc, dict):
+                            embedding_body["inputText"] = doc.get("text")
+                            embedding_body["inputImage"] = doc.get(
+                                "image"
+                            )  # expects a base64-encoded image
+                        else:
+                            embedding_body["inputText"] = doc
+
+                        # Add model-specific inference parameters
+                        if model_kwargs:
+                            embedding_body = embedding_body | model_kwargs
+
+                        # Clean up null values
+                        embedding_body = {k: v for k, v in embedding_body.items() if v}
+
+                        # Format payload
+                        embedding_body_payload: str = json.dumps(embedding_body)
+
                         response = self.client.invoke_model(
-                            body=embedding_body,
+                            body=embedding_body_payload,
                             modelId=self.name,
                             accept="application/json",
                             contentType="application/json",
@@ -184,9 +203,16 @@ class BedrockEncoder(DenseEncoder):
                 elif self.name and "cohere" in self.name:
                     chunked_docs = self.chunk_strings(docs)
                     for chunk in chunked_docs:
-                        chunk = json.dumps(
-                            {"texts": chunk, "input_type": self.input_type}
-                        )
+                        chunk = {"texts": chunk, "input_type": self.input_type}
+
+                        # Add model-specific inference parameters
+                        # Note: if specified, input_type will be overwritten by model_kwargs
+                        if model_kwargs:
+                            chunk = chunk | model_kwargs
+
+                        # Format payload
+                        chunk = json.dumps(chunk)
+
                         response = self.client.invoke_model(
                             body=chunk,
                             modelId=self.name,

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

@@ -1,3 +1,5 @@
+from datetime import datetime
+import time
 from typing import Any, List, Optional, Tuple, Union, Dict
 import json
 
@@ -157,26 +159,91 @@ class BaseIndex(BaseModel):
         logger.warning("This method should be implemented by subclasses.")
         self.index = None
 
-    def _read_hash(self) -> ConfigParameter:
-        """
-        Read the hash of the previously written index.
+    def _read_config(self, field: str, scope: str | None = None) -> ConfigParameter:
+        """Read a config parameter from the index.
 
-        This method should be implemented by subclasses.
+        :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
         """
         logger.warning("This method should be implemented by subclasses.")
         return ConfigParameter(
-            field="sr_hash",
+            field=field,
             value="",
-            namespace="",
+            scope=scope,
         )
 
-    def _write_config(self, config: ConfigParameter):
+    def _read_hash(self) -> ConfigParameter:
+        """Read the hash of the previously written index.
+
+        :return: The config parameter that was read.
+        :rtype: ConfigParameter
         """
-        Write a config parameter to the index.
+        return self._read_config(field="sr_hash")
 
-        This method should be implemented by subclasses.
+    def _write_config(self, config: ConfigParameter) -> ConfigParameter:
+        """Write a config parameter to the index.
+
+        :param config: The config parameter to write.
+        :type config: ConfigParameter
+        :return: The config parameter that was written.
+        :rtype: ConfigParameter
         """
         logger.warning("This method should be implemented by subclasses.")
+        return config
+
+    def lock(
+        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.
+
+        :param scope: The scope to lock.
+        :type scope: str | None
+        :param wait: The number of seconds to wait for the index to be unlocked, if
+        set to 0, will raise an error if index is already locked/unlocked.
+        :type wait: int
+        :return: The config parameter that was locked.
+        :rtype: ConfigParameter
+        """
+        start_time = datetime.now()
+        while True:
+            if self._is_locked(scope=scope) != value:
+                # in this case, we can set the lock value
+                break
+            if (datetime.now() - start_time).total_seconds() < wait:
+                # wait for 2.5 seconds before checking again
+                time.sleep(2.5)
+            else:
+                raise ValueError(
+                    f"Index is already {'locked' if value else 'unlocked'}."
+                )
+        lock_param = ConfigParameter(
+            field="sr_lock",
+            value=str(value),
+            scope=scope,
+        )
+        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
+        """
+        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}")
 
     def _get_all(self, prefix: Optional[str] = None, include_metadata: bool = False):
         """

{semantic_router-0.1.0.dev3 → semantic_router-0.1.0.dev4}/semantic_router/index/pinecone.py

@@ -405,39 +405,43 @@ class PineconeIndex(BaseIndex):
         route_names = [result["metadata"]["sr_route"] for result in results["matches"]]
         return np.array(scores), route_names
 
-    def _read_hash(self) -> ConfigParameter:
+    def _read_config(self, field: str, scope: str | None = None) -> ConfigParameter:
+        scope = scope or self.namespace
         if self.index is None:
             return ConfigParameter(
-                field="sr_hash",
+                field=field,
                 value="",
-                namespace=self.namespace,
+                scope=scope,
             )
-        hash_id = f"sr_hash#{self.namespace}"
-        hash_record = self.index.fetch(
-            ids=[hash_id],
+        config_id = f"{field}#{scope}"
+        config_record = self.index.fetch(
+            ids=[config_id],
             namespace="sr_config",
         )
-        if hash_record["vectors"]:
+        if config_record["vectors"]:
             return ConfigParameter(
-                field="sr_hash",
-                value=hash_record["vectors"][hash_id]["metadata"]["value"],
-                created_at=hash_record["vectors"][hash_id]["metadata"]["created_at"],
-                namespace=self.namespace,
+                field=field,
+                value=config_record["vectors"][config_id]["metadata"]["value"],
+                created_at=config_record["vectors"][config_id]["metadata"][
+                    "created_at"
+                ],
+                scope=scope,
             )
         else:
-            logger.warning("Configuration for hash parameter not found in index.")
+            logger.warning(f"Configuration for {field} parameter not found in index.")
             return ConfigParameter(
-                field="sr_hash",
+                field=field,
                 value="",
-                namespace=self.namespace,
+                scope=scope,
             )
 
-    def _write_config(self, config: ConfigParameter) -> None:
+    def _write_config(self, config: ConfigParameter) -> ConfigParameter:
         """Method to write a config parameter to the remote Pinecone index.
 
         :param config: The config parameter to write to the index.
         :type config: ConfigParameter
         """
+        config.scope = config.scope or self.namespace
         if self.index is None:
             raise ValueError("Index has not been initialized.")
         if self.dimensions is None:
@@ -446,6 +450,7 @@ class PineconeIndex(BaseIndex):
             vectors=[config.to_pinecone(dimensions=self.dimensions)],
             namespace="sr_config",
         )
+        return config
 
     async def aquery(
         self,

{semantic_router-0.1.0.dev3 → semantic_router-0.1.0.dev4}/semantic_router/routers/base.py

@@ -4,6 +4,7 @@ import os
 import random
 import hashlib
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing_extensions import deprecated
 from pydantic import BaseModel, Field
 
 import numpy as np
@@ -280,6 +281,20 @@ class RouterConfig:
         )
 
 
+def xq_reshape(xq: List[float] | np.ndarray) -> np.ndarray:
+    # convert to numpy array if not already
+    if not isinstance(xq, np.ndarray):
+        xq = np.array(xq)
+    # check if vector is 1D and expand to 2D if necessary
+    if len(xq.shape) == 1:
+        xq = np.expand_dims(xq, axis=0)
+    if xq.shape[0] != 1:
+        raise ValueError(
+            f"Expected (1, x) dimensional input for query, got {xq.shape}."
+        )
+    return xq
+
+
 class BaseRouter(BaseModel):
     encoder: DenseEncoder = Field(default_factory=OpenAIEncoder)
     index: BaseIndex = Field(default_factory=BaseIndex)
@@ -402,7 +417,7 @@ class BaseRouter(BaseModel):
     def __call__(
         self,
         text: Optional[str] = None,
-        vector: Optional[List[float]] = None,
+        vector: Optional[List[float] | np.ndarray] = None,
         simulate_static: bool = False,
         route_filter: Optional[List[str]] = None,
     ) -> RouteChoice:
@@ -411,6 +426,9 @@ class BaseRouter(BaseModel):
             if text is None:
                 raise ValueError("Either text or vector must be provided")
             vector = self._encode(text=[text])
+        # convert to numpy array if not already
+        vector = xq_reshape(vector)
+        # calculate semantics
         route, top_class_scores = self._retrieve_top_route(vector, route_filter)
         passed = self._check_threshold(top_class_scores, route)
         if passed and route is not None and not simulate_static:
@@ -444,7 +462,7 @@ class BaseRouter(BaseModel):
     async def acall(
         self,
         text: Optional[str] = None,
-        vector: Optional[List[float]] = None,
+        vector: Optional[List[float] | np.ndarray] = None,
         simulate_static: bool = False,
         route_filter: Optional[List[str]] = None,
     ) -> RouteChoice:
@@ -453,7 +471,9 @@ class BaseRouter(BaseModel):
             if text is None:
                 raise ValueError("Either text or vector must be provided")
             vector = await self._async_encode(text=[text])
-
+        # convert to numpy array if not already
+        vector = xq_reshape(vector)
+        # calculate semantics
         route, top_class_scores = await self._async_retrieve_top_route(
             vector, route_filter
         )
@@ -483,19 +503,21 @@ class BaseRouter(BaseModel):
             # if no route passes threshold, return empty route choice
             return RouteChoice()
 
+    # TODO: add multiple routes return to __call__ and acall
+    @deprecated("This method is deprecated. Use `__call__` instead.")
     def retrieve_multiple_routes(
         self,
         text: Optional[str] = None,
-        vector: Optional[List[float]] = None,
+        vector: Optional[List[float] | np.ndarray] = None,
     ) -> List[RouteChoice]:
         if vector is None:
             if text is None:
                 raise ValueError("Either text or vector must be provided")
-            vector_arr = self._encode(text=[text])
-        else:
-            vector_arr = np.array(vector)
+            vector = self._encode(text=[text])
+        # convert to numpy array if not already
+        vector = xq_reshape(vector)
         # get relevant utterances
-        results = self._retrieve(xq=vector_arr)
+        results = self._retrieve(xq=vector)
         # decide most relevant routes
         categories_with_scores = self._semantic_classify_multiple_routes(results)
         return [
@@ -514,16 +536,14 @@ class BaseRouter(BaseModel):
         # return route_choices
 
     def _retrieve_top_route(
-        self, vector: List[float], route_filter: Optional[List[str]] = None
+        self, vector: np.ndarray, route_filter: Optional[List[str]] = None
     ) -> Tuple[Optional[Route], List[float]]:
         """
         Retrieve the top matching route based on the given vector.
         Returns a tuple of the route (if any) and the scores of the top class.
         """
         # get relevant results (scores and routes)
-        results = self._retrieve(
-            xq=np.array(vector), top_k=self.top_k, route_filter=route_filter
-        )
+        results = self._retrieve(xq=vector, top_k=self.top_k, route_filter=route_filter)
         # decide most relevant routes
         top_class, top_class_scores = self._semantic_classify(results)
         # TODO do we need this check?
@@ -531,11 +551,11 @@ class BaseRouter(BaseModel):
         return route, top_class_scores
 
     async def _async_retrieve_top_route(
-        self, vector: List[float], route_filter: Optional[List[str]] = None
+        self, vector: np.ndarray, route_filter: Optional[List[str]] = None
     ) -> Tuple[Optional[Route], List[float]]:
         # get relevant results (scores and routes)
         results = await self._async_retrieve(
-            xq=np.array(vector), top_k=self.top_k, route_filter=route_filter
+            xq=vector, top_k=self.top_k, route_filter=route_filter
         )
         # decide most relevant routes
         top_class, top_class_scores = await self._async_semantic_classify(results)
@@ -543,7 +563,7 @@ class BaseRouter(BaseModel):
         route = self.check_for_matching_routes(top_class)
         return route, top_class_scores
 
-    def sync(self, sync_mode: str, force: bool = False) -> List[str]:
+    def sync(self, sync_mode: str, force: bool = False, wait: int = 0) -> List[str]:
         """Runs a sync of the local routes with the remote index.
 
         :param sync_mode: The mode to sync the routes with the remote index.
@@ -551,6 +571,10 @@ class BaseRouter(BaseModel):
         :param force: Whether to force the sync even if the local and remote
             hashes already match. Defaults to False.
         :type force: bool, optional
+        :param wait: The number of seconds to wait for the index to be unlocked
+        before proceeding with the sync. If set to 0, will raise an error if
+        index is already locked/unlocked.
+        :type wait: int
         :return: A list of diffs describing the addressed differences between
             the local and remote route layers.
         :rtype: List[str]
@@ -565,7 +589,9 @@ class BaseRouter(BaseModel):
                 remote_utterances=local_utterances,
             )
             return diff.to_utterance_str()
-        # otherwise we continue with the sync, first creating a diff
+        # otherwise we continue with the sync, first locking the index
+        _ = self.index.lock(value=True, wait=wait)
+        # first creating a diff
         local_utterances = self.to_config().to_utterances()
         remote_utterances = self.index.get_utterances()
         diff = UtteranceDiff.from_utterances(
@@ -576,6 +602,8 @@ class BaseRouter(BaseModel):
         sync_strategy = diff.get_sync_strategy(sync_mode=sync_mode)
         # and execute
         self._execute_sync_strategy(sync_strategy)
+        # unlock index after sync
+        _ = self.index.lock(value=False)
         return diff.to_utterance_str()
 
     def _execute_sync_strategy(self, strategy: Dict[str, Dict[str, List[Utterance]]]):
@@ -781,6 +809,9 @@ class BaseRouter(BaseModel):
         :param route_name: the name of the route to be deleted
         :type str:
         """
+        # ensure index is not locked
+        if self.index._is_locked():
+            raise ValueError("Index is locked. Cannot delete route.")
         current_local_hash = self._get_hash()
         current_remote_hash = self.index._read_hash()
         if current_remote_hash.value == "":
@@ -928,7 +959,7 @@ class BaseRouter(BaseModel):
         """Given a query vector, retrieve the top_k most similar records."""
         # get scores and routes
         scores, routes = self.index.query(
-            vector=xq, top_k=top_k, route_filter=route_filter
+            vector=xq[0], top_k=top_k, route_filter=route_filter
         )
         return [{"route": d, "score": s.item()} for d, s in zip(routes, scores)]
 
@@ -938,7 +969,7 @@ class BaseRouter(BaseModel):
         """Given a query vector, retrieve the top_k most similar records."""
         # get scores and routes
         scores, routes = await self.index.aquery(
-            vector=xq, top_k=top_k, route_filter=route_filter
+            vector=xq[0], top_k=top_k, route_filter=route_filter
         )
         return [{"route": d, "score": s.item()} for d, s in zip(routes, scores)]
 

{semantic_router-0.1.0.dev3 → semantic_router-0.1.0.dev4}/semantic_router/routers/hybrid.py

@@ -14,7 +14,7 @@ from semantic_router.route import Route
 from semantic_router.index import BaseIndex, HybridLocalIndex
 from semantic_router.schema import RouteChoice, SparseEmbedding, Utterance
 from semantic_router.utils.logger import logger
-from semantic_router.routers.base import BaseRouter
+from semantic_router.routers.base import BaseRouter, xq_reshape
 from semantic_router.llms import BaseLLM
 
 
@@ -197,18 +197,19 @@ class HybridRouter(BaseRouter):
     def __call__(
         self,
         text: Optional[str] = None,
-        vector: Optional[List[float]] = None,
+        vector: Optional[List[float] | np.ndarray] = None,
         simulate_static: bool = False,
         route_filter: Optional[List[str]] = None,
         sparse_vector: dict[int, float] | SparseEmbedding | None = None,
     ) -> RouteChoice:
-        vector_arr: np.ndarray | None = None
         potential_sparse_vector: List[SparseEmbedding] | None = None
         # if no vector provided, encode text to get vector
         if vector is None:
             if text is None:
                 raise ValueError("Either text or vector must be provided")
-            vector_arr, potential_sparse_vector = self._encode(text=[text])
+            vector, potential_sparse_vector = self._encode(text=[text])
+        # convert to numpy array if not already
+        vector = xq_reshape(vector)
         if sparse_vector is None:
             if text is None:
                 raise ValueError("Either text or sparse_vector must be provided")
@@ -217,10 +218,9 @@ class HybridRouter(BaseRouter):
             )
         if sparse_vector is None:
             raise ValueError("Sparse vector is required for HybridLocalIndex.")
-        vector_arr = vector_arr if vector_arr is not None else np.array(vector)
         # TODO: add alpha as a parameter
         scores, route_names = self.index.query(
-            vector=vector_arr,
+            vector=vector,
             top_k=self.top_k,
             route_filter=route_filter,
             sparse_vector=sparse_vector,

{semantic_router-0.1.0.dev3 → semantic_router-0.1.0.dev4}/semantic_router/routers/semantic.py

@@ -40,14 +40,12 @@ class SemanticRouter(BaseRouter):
         """Given some text, encode it."""
         # create query vector
         xq = np.array(self.encoder(text))
-        xq = np.squeeze(xq)  # Reduce to 1d array.
         return xq
 
     async def _async_encode(self, text: list[str]) -> Any:
         """Given some text, encode it."""
         # create query vector
         xq = np.array(await self.encoder.acall(docs=text))
-        xq = np.squeeze(xq)  # Reduce to 1d array.
         return xq
 
     def add(self, routes: List[Route] | Route):

{semantic_router-0.1.0.dev3 → semantic_router-0.1.0.dev4}/semantic_router/schema.py

@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, timezone
 from difflib import Differ
 from enum import Enum
 import json
@@ -62,12 +62,13 @@ class Message(BaseModel):
 class ConfigParameter(BaseModel):
     field: str
     value: str
-    namespace: Optional[str] = None
-    created_at: str = Field(default_factory=lambda: datetime.utcnow().isoformat())
+    scope: Optional[str] = None
+    created_at: str = Field(
+        default_factory=lambda: datetime.now(timezone.utc).isoformat()
+    )
 
     def to_pinecone(self, dimensions: int):
-        if self.namespace is None:
-            namespace = ""
+        namespace = self.scope or ""
         return {
             "id": f"{self.field}#{namespace}",
             "values": [0.1] * dimensions,