diff --git a/semantic_router/indices/base.py b/semantic_router/indices/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..f5c3e8d6193c33551cf21dc914b65690f80d24d2
--- /dev/null
+++ b/semantic_router/indices/base.py
@@ -0,0 +1,5 @@
+from pydantic.v1 import BaseModel
+
+class BaseIndex(BaseModel):
+
+ pass
\ No newline at end of file
diff --git a/semantic_router/indices/local_index.py b/semantic_router/indices/local_index.py
index 9a72b7efae93dcb3cf5535268bfc67fc73e40a90..784c99673c38c45adf71f01d41438a7e961ced4e 100644
--- a/semantic_router/indices/local_index.py
+++ b/semantic_router/indices/local_index.py
@@ -1,6 +1,8 @@
import numpy as np
from typing import List, Any
from .base import BaseIndex
+from semantic_router.linear import similarity_matrix, top_scores
+from typing import Tuple
class LocalIndex(BaseIndex):
"""
@@ -9,35 +11,32 @@ class LocalIndex(BaseIndex):
def __init__(self):
self.index = None
- self.categories = None
- def add(self, items: List[Any], categories: List[str]):
+ def add(self, embeds: List[Any]):
"""
- Add items to the index with their corresponding categories.
+ Add items to the index.
"""
- embeds = np.array(items)
+ embeds = np.array(embeds)
if self.index is None:
self.index = embeds
- self.categories = np.array(categories)
else:
self.index = np.concatenate([self.index, embeds])
- self.categories = np.concatenate([self.categories, np.array(categories)])
- def remove(self, category: str):
+ def remove(self, indices_to_remove: List[int]):
"""
Remove all items of a specific category from the index.
"""
- if self.categories is not None:
- indices_to_remove = np.where(self.categories == category)[0]
- self.index = np.delete(self.index, indices_to_remove, axis=0)
- self.categories = np.delete(self.categories, indices_to_remove, axis=0)
+ self.index = np.delete(self.index, indices_to_remove, axis=0)
- def search(self, query: Any, top_k: int = 5) -> List[Any]:
+ def is_index_populated(self):
+ return self.index is not None and len(self.index) > 0
+
+ def search(self, query_vector: Any, top_k: int = 5) -> Tuple[np.ndarray, np.ndarray]:
"""
Search the index for the query and return top_k results.
"""
if self.index is None:
- return []
- sim = np.dot(self.index, query) / (np.linalg.norm(self.index, axis=1) * np.linalg.norm(query))
- idx = np.argsort(sim)[-top_k:]
- return [(self.categories[i], sim[i]) for i in idx[::-1]]
\ No newline at end of file
+ raise ValueError("Index is not populated.")
+ sim = similarity_matrix(query_vector, self.index)
+ return top_scores(sim, top_k)
+
diff --git a/semantic_router/layer.py b/semantic_router/layer.py
index 08afccfab19b590a1fc975e9325c69d6cb961d0b..487d8b24554bc059e9898d8610fac39ece5186e2 100644
--- a/semantic_router/layer.py
+++ b/semantic_router/layer.py
@@ -11,9 +11,10 @@ from semantic_router.encoders import BaseEncoder, OpenAIEncoder
from semantic_router.linear import similarity_matrix, top_scores
from semantic_router.llms import BaseLLM, OpenAILLM
from semantic_router.route import Route
-from semantic_router.schema import Encoder, EncoderType, RouteChoice
+from semantic_router.schema import Encoder, EncoderType, RouteChoice, Index
from semantic_router.utils.logger import logger
+IndexType = Union[LocalIndex, None]
def is_valid(layer_config: str) -> bool:
"""Make sure the given string is json format and contains the 3 keys: ["encoder_name", "encoder_type", "routes"]"""
@@ -155,15 +156,17 @@ class RouteLayer:
categories: Optional[np.ndarray] = None
score_threshold: float
encoder: BaseEncoder
+ index: IndexType = None
def __init__(
self,
encoder: Optional[BaseEncoder] = None,
llm: Optional[BaseLLM] = None,
routes: Optional[List[Route]] = None,
+ index_name: Optional[str] = None,
):
- logger.info("Initializing RouteLayer")
- self.index = None
+ logger.info("local")
+ self.index = Index.get_by_name(index_name="index")
self.categories = None
if encoder is None:
logger.warning(
@@ -281,11 +284,7 @@ class RouteLayer:
str_arr = np.array([route.name] * len(embeds))
self.categories = np.concatenate([self.categories, str_arr])
# create utterance array (the index)
- if self.index is None:
- self.index = np.array(embeds)
- else:
- embed_arr = np.array(embeds)
- self.index = np.concatenate([self.index, embed_arr])
+ self.index.add(embeds)
# add route to routes list
self.routes.append(route)
@@ -301,13 +300,13 @@ class RouteLayer:
self.routes = [route for route in self.routes if route.name != name]
logger.info(f"Removed route `{name}`")
# Also remove from index and categories
- if self.categories is not None and self.index is not None:
+ if self.categories is not None and self.index.is_index_populated():
indices_to_remove = [
i
for i, route_name in enumerate(self.categories)
if route_name == name
]
- self.index = np.delete(self.index, indices_to_remove, axis=0)
+ self.index.remove(indices_to_remove)
self.categories = np.delete(self.categories, indices_to_remove, axis=0)
def _add_routes(self, routes: List[Route]):
@@ -325,14 +324,7 @@ class RouteLayer:
if self.categories is not None
else route_array
)
-
- # create utterance array (the index)
- embed_utterance_arr = np.array(embedded_utterance)
- self.index = (
- np.concatenate([self.index, embed_utterance_arr])
- if self.index is not None
- else embed_utterance_arr
- )
+ self.index.add(embedded_utterance)
def _encode(self, text: str) -> Any:
"""Given some text, encode it."""
@@ -343,10 +335,9 @@ class RouteLayer:
def _retrieve(self, xq: Any, top_k: int = 5) -> List[dict]:
"""Given a query vector, retrieve the top_k most similar records."""
- if self.index is not None:
+ if self.index.is_index_populated():
# calculate similarity matrix
- sim = similarity_matrix(xq, self.index)
- scores, idx = top_scores(sim, top_k)
+ scores, idx = self.index.search(xq, top_k)
# get the utterance categories (route names)
routes = self.categories[idx] if self.categories is not None else []
return [{"route": d, "score": s.item()} for d, s in zip(routes, scores)]
diff --git a/semantic_router/schema.py b/semantic_router/schema.py
index 46ee7f590c3275b314b9ecb3e52adc54f16012fe..917055cb69afee66a2f1581b683c013d756b33e7 100644
--- a/semantic_router/schema.py
+++ b/semantic_router/schema.py
@@ -11,6 +11,8 @@ from semantic_router.encoders import (
OpenAIEncoder,
)
+from semantic_router.indices.local_index import LocalIndex
+
class EncoderType(Enum):
HUGGINGFACE = "huggingface"
@@ -73,3 +75,13 @@ class DocumentSplit(BaseModel):
docs: List[str]
is_triggered: bool = False
triggered_score: Optional[float] = None
+
+
+class Index:
+ @classmethod
+ def get_by_name(cls, index_name: str):
+ if index_name == "local" or index_name is None:
+ return LocalIndex()
+ # TODO: Later we'll add more index options.
+ else:
+ raise ValueError(f"Invalid index name: {index_name}")
\ No newline at end of file