Merge pull request #625 from aurelio-labs/josh/semantic-router/local-index/metadata

feat: added metadata attribute to local index

Author: jamescalam

semantic_router/encoders/openai.py

@@ -137,6 +137,7 @@ def __call__(self, docs: List[str], truncate: bool = True) -> List[List[float]]:
         # Exponential backoff
         for j in range(self.max_retries + 1):
             try:
+                logger.debug(f"Creating embeddings for {len(docs)} docs")
                 embeds = self._client.embeddings.create(
                     input=docs,
                     model=self.name,

semantic_router/index/hybrid_local.py

@@ -15,6 +15,7 @@ class HybridLocalIndex(LocalIndex):
 
     def __init__(self, **data):
         super().__init__(**data)
+        self.metadata = None
 
     def add(
         self,
@@ -64,12 +65,32 @@ def add(
             ]  # TODO: switch back to using SparseEmbedding later
             self.routes = routes_arr
             self.utterances = utterances_arr
+            self.metadata = (
+                np.array(metadata_list, dtype=object)
+                if metadata_list
+                else np.array([{} for _ in utterances], dtype=object)
+            )
         else:
             # TODO: we should probably switch to an `upsert` method and standardize elsewhere
             self.index = np.concatenate([self.index, embeds])
             self.sparse_index.extend([x.to_dict() for x in sparse_embeddings])
             self.routes = np.concatenate([self.routes, routes_arr])
             self.utterances = np.concatenate([self.utterances, utterances_arr])
+            if self.metadata is not None:
+                self.metadata = np.concatenate(
+                    [
+                        self.metadata,
+                        np.array(metadata_list, dtype=object)
+                        if metadata_list
+                        else np.array([{} for _ in utterances], dtype=object),
+                    ]
+                )
+            else:
+                self.metadata = (
+                    np.array(metadata_list, dtype=object)
+                    if metadata_list
+                    else np.array([{} for _ in utterances], dtype=object)
+                )
 
     async def aadd(
         self,
@@ -119,7 +140,20 @@ def get_utterances(self, include_metadata: bool = False) -> 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)]
+        if include_metadata and self.metadata is not None:
+            return [
+                Utterance(
+                    route=route,
+                    utterance=utterance,
+                    function_schemas=None,
+                    metadata=metadata,
+                )
+                for route, utterance, metadata in zip(
+                    self.routes, self.utterances, self.metadata
+                )
+            ]
+        else:
+            return [Utterance.from_tuple(x) for x in zip(self.routes, self.utterances)]
 
     def _sparse_dot_product(
         self, vec_a: dict[int, float], vec_b: dict[int, float]
@@ -260,6 +294,8 @@ def delete(self, route_name: str):
             self.index = np.delete(self.index, delete_idx, axis=0)
             self.routes = np.delete(self.routes, delete_idx, axis=0)
             self.utterances = np.delete(self.utterances, delete_idx, axis=0)
+            if self.metadata is not None:
+                self.metadata = np.delete(self.metadata, delete_idx, axis=0)
         else:
             raise ValueError(
                 "Attempted to delete route records but either index, routes or "
@@ -275,6 +311,7 @@ def delete_index(self):
         self.index = None
         self.routes = None
         self.utterances = None
+        self.metadata = None
 
     def _get_indices_for_route(self, route_name: str):
         """Gets an array of indices for a specific route.

semantic_router/index/local.py

@@ -1,7 +1,7 @@
 from typing import Any, ClassVar, Dict, List, Optional, Tuple
 
 import numpy as np
-from pydantic import ConfigDict
+from pydantic import ConfigDict, Field
 
 from semantic_router.index.base import BaseIndex, IndexConfig
 from semantic_router.linear import similarity_matrix, top_scores
@@ -11,9 +11,12 @@
 
 class LocalIndex(BaseIndex):
     type: str = "local"
+    metadata: Optional[np.ndarray] = Field(default=None, exclude=True)
 
     def __init__(self, **data):
         super().__init__(**data)
+        if self.metadata is None:
+            self.metadata = None
 
     # Stop pydantic from complaining about Optional[np.ndarray]type hints.
     model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
@@ -50,10 +53,30 @@ def add(
             self.index = embeds  # type: ignore
             self.routes = routes_arr
             self.utterances = utterances_arr
+            self.metadata = (
+                np.array(metadata_list, dtype=object)
+                if metadata_list
+                else np.array([{} for _ in utterances], dtype=object)
+            )
         else:
             self.index = np.concatenate([self.index, embeds])
             self.routes = np.concatenate([self.routes, routes_arr])
             self.utterances = np.concatenate([self.utterances, utterances_arr])
+            if self.metadata is not None:
+                self.metadata = np.concatenate(
+                    [
+                        self.metadata,
+                        np.array(metadata_list, dtype=object)
+                        if metadata_list
+                        else np.array([{} for _ in utterances], dtype=object),
+                    ]
+                )
+            else:
+                self.metadata = (
+                    np.array(metadata_list, dtype=object)
+                    if metadata_list
+                    else np.array([{} for _ in utterances], dtype=object)
+                )
 
     def _remove_and_sync(self, routes_to_delete: dict) -> np.ndarray:
         """Remove and sync the index.
@@ -80,21 +103,35 @@ def _remove_and_sync(self, routes_to_delete: dict) -> np.ndarray:
         self.index = self.index[mask]
         self.routes = self.routes[mask]
         self.utterances = self.utterances[mask]
+        if self.metadata is not None:
+            self.metadata = self.metadata[mask]
         # return what was removed
         return route_utterances[~mask]
 
     def get_utterances(self, include_metadata: bool = False) -> List[Utterance]:
         """Gets a list of route and utterance objects currently stored in the index.
 
         :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.
+        the returned Utterance objects - LocalIndex now includes metadata if present.
         :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)]
+        if include_metadata and self.metadata is not None:
+            return [
+                Utterance(
+                    route=route,
+                    utterance=utterance,
+                    function_schemas=None,
+                    metadata=metadata,
+                )
+                for route, utterance, metadata in zip(
+                    self.routes, self.utterances, self.metadata
+                )
+            ]
+        else:
+            return [Utterance.from_tuple(x) for x in zip(self.routes, self.utterances)]
 
     def describe(self) -> IndexConfig:
         """Describe the index.
@@ -235,6 +272,8 @@ def delete(self, route_name: str):
             self.index = np.delete(self.index, delete_idx, axis=0)
             self.routes = np.delete(self.routes, delete_idx, axis=0)
             self.utterances = np.delete(self.utterances, delete_idx, axis=0)
+            if self.metadata is not None:
+                self.metadata = np.delete(self.metadata, delete_idx, axis=0)
         else:
             raise ValueError(
                 "Attempted to delete route records but either index, routes or "
@@ -260,6 +299,7 @@ def delete_index(self):
         self.index = None
         self.routes = None
         self.utterances = None
+        self.metadata = None
 
     async def adelete_index(self):
         """Deletes the index, effectively clearing it and setting it to None. Note that this just points
@@ -272,6 +312,7 @@ async def adelete_index(self):
         self.index = None
         self.routes = None
         self.utterances = None
+        self.metadata = None
 
     def _get_indices_for_route(self, route_name: str):
         """Gets an array of indices for a specific route.

semantic_router/utils/logger.py

@@ -43,12 +43,11 @@ def setup_custom_logger(name):
     )
     log_level = log_level.upper()
 
-    if not logger.hasHandlers():
-        add_coloured_handler(logger)
-        logger.setLevel(log_level)
-        logger.propagate = False
+    add_coloured_handler(logger)
+    logger.setLevel(log_level)
+    logger.propagate = False
 
     return logger
 
 
-logger: logging.Logger = setup_custom_logger(__name__)
+logger: logging.Logger = setup_custom_logger("semantic_router")