diff --git a/semantic_router/index/__init__.py b/semantic_router/index/__init__.py
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..1ad70df4228e35b6212d2b749e5c1b123f539c25 100644
--- a/semantic_router/index/__init__.py
+++ b/semantic_router/index/__init__.py
@@ -0,0 +1,9 @@
+from semantic_router.index.base import BaseIndex
+from semantic_router.index.local import LocalIndex
+from semantic_router.index.pinecone import PineconeIndex
+
+__all__ = [
+ "BaseIndex",
+ "LocalIndex",
+ "PineconeIndex",
+]
diff --git a/semantic_router/index/local.py b/semantic_router/index/local.py
index 02a1f209fc17052e1c8fc52c6198347a347e80de..08cd088591850f2b5da2aefa5a430fbf7870cbb2 100644
--- a/semantic_router/index/local.py
+++ b/semantic_router/index/local.py
@@ -5,12 +5,16 @@ from semantic_router.index.base import BaseIndex
class LocalIndex(BaseIndex):
- def __init__(self):
- super().__init__()
+ def __init__(
+ self,
+ index: Optional[np.ndarray] = None,
+ routes: Optional[List[str]] = None,
+ utterances: Optional[List[str]] = None,
+ ):
+ super().__init__(
+ index=index, routes=routes, utterances=utterances
+ )
self.type = "local"
- self.index: Optional[np.ndarray] = None
- self.routes: Optional[np.ndarray] = None
- self.utterances: Optional[np.ndarray] = None
class Config: # Stop pydantic from complaining about Optional[np.ndarray] type hints.
arbitrary_types_allowed = True
@@ -30,30 +34,16 @@ class LocalIndex(BaseIndex):
self.routes = np.concatenate([self.routes, routes_arr])
self.utterances = np.concatenate([self.utterances, utterances_arr])
- def _get_indices_for_route(self, route_name: str):
- """Gets an array of indices for a specific route."""
- if self.routes is None:
- raise ValueError("Routes are not populated.")
- idx = [i for i, route in enumerate(self.routes) if route == route_name]
- return idx
-
- def delete(self, route_name: str):
+ def get_routes(self) -> List[Tuple]:
"""
- Delete all records of a specific route from the index.
+ Gets a list of route and utterance objects currently stored in the index.
+
+ Returns:
+ List[Tuple]: A list of (route_name, utterance) objects.
"""
- if (
- self.index is not None
- and self.routes is not None
- and self.utterances is not None
- ):
- delete_idx = self._get_indices_for_route(route_name=route_name)
- 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)
- else:
- raise ValueError(
- "Attempted to delete route records but either indx, routes or utterances is None."
- )
+ if self.route_names is None or self.utterances is None:
+ raise ValueError("No routes have been added to the index.")
+ return list(zip(self.route_names, self.utterances))
def describe(self) -> dict:
return {
@@ -74,9 +64,40 @@ class LocalIndex(BaseIndex):
# get routes from index values
route_names = self.routes[idx].copy()
return scores, route_names
+
+ def delete(self, route_name: str):
+ """
+ Delete all records of a specific route from the index.
+ """
+ if (
+ self.index is not None
+ and self.routes is not None
+ and self.utterances is not None
+ ):
+ delete_idx = self._get_indices_for_route(route_name=route_name)
+ 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)
+ else:
+ raise ValueError(
+ "Attempted to delete route records but either index, routes or utterances is None."
+ )
def delete_index(self):
"""
Deletes the index, effectively clearing it and setting it to None.
"""
self.index = None
+
+ def _get_indices_for_route(self, route_name: str):
+ """Gets an array of indices for a specific route."""
+ if self.routes is None:
+ raise ValueError("Routes are not populated.")
+ idx = [i for i, route in enumerate(self.routes) if route == route_name]
+ return idx
+
+ def __len__(self):
+ if self.index is not None:
+ return self.index.shape[0]
+ else:
+ return 0
diff --git a/semantic_router/index/pinecone.py b/semantic_router/index/pinecone.py
index d2b21d95f9ff176493f91d91050eed204125aa45..4594962f5a3d612dadcf7161c8bb6001d5d1edce 100644
--- a/semantic_router/index/pinecone.py
+++ b/semantic_router/index/pinecone.py
@@ -124,16 +124,71 @@ class PineconeIndex(BaseIndex):
else:
raise ValueError("Index is None could not upsert.")
- def _get_route_vecs(self, route_name: str):
+ def _get_route_ids(self, route_name: str):
clean_route = clean_route_name(route_name)
- res = requests.get(
- f"https://{self.host}/vectors/list?prefix={clean_route}#",
- headers={"Api-Key": os.environ["PINECONE_API_KEY"]},
- )
- return [vec["id"] for vec in res.json()["vectors"]]
+ ids, _ = self._get_all(prefix=f"{clean_route}#")
+ return ids
+
+ def _get_all(self, prefix: Optional[str] = None, include_metadata: bool = False):
+ """
+ Retrieves all vector IDs from the Pinecone index using pagination.
+ """
+ all_vector_ids = []
+ next_page_token = None
+
+ if prefix:
+ prefix_str = f"?prefix={prefix}"
+ else:
+ prefix_str = ""
+
+ # Construct the request URL for listing vectors. Adjust parameters as needed.
+ list_url = f"https://{self.host}/vectors/list{prefix_str}"
+ params = {}
+ headers = {"Api-Key": os.getenv("PINECONE_API_KEY")}
+ metadata = []
+
+ while True:
+ if next_page_token:
+ params["paginationToken"] = next_page_token
+
+ # Make the request to list vectors. Adjust headers and parameters as needed.
+ response = requests.get(list_url, params=params, headers=headers)
+ response_data = response.json()
+ print(response_data)
+
+ # Extract vector IDs from the response and add them to the list
+ vector_ids = [vec["id"] for vec in response_data.get("vectors", [])]
+ all_vector_ids.extend(vector_ids)
+
+ # if we need metadata, we fetch it
+ if include_metadata:
+ res_meta = self.index.fetch(ids=vector_ids)
+ # extract metadata only
+ metadata.extend(
+ [x["metadata"] for x in res_meta["vectors"].values()]
+ )
+
+ # Check if there's a next page token; if not, break the loop
+ next_page_token = response_data.get("pagination", {}).get("next")
+ if not next_page_token:
+ break
+
+ return all_vector_ids, metadata
+
+ def get_routes(self) -> List[Tuple]:
+ """
+ Gets a list of route and utterance objects currently stored in the index.
+
+ Returns:
+ List[Tuple]: A list of (route_name, utterance) objects.
+ """
+ # Get all records
+ _, metadata = self._get_all(include_metadata=True)
+ route_tuples = [(x["sr_route"], x["sr_utterance"]) for x in metadata]
+ return route_tuples
def delete(self, route_name: str):
- route_vec_ids = self._get_route_vecs(route_name=route_name)
+ route_vec_ids = self._get_route_ids(route_name=route_name)
if self.index is not None:
self.index.delete(ids=route_vec_ids)
else:
diff --git a/semantic_router/layer.py b/semantic_router/layer.py
index d4aa5bcbaab77a03c8d7025d691006d4608db359..bec876748caa563f2941bb7d800c240466d9930b 100644
--- a/semantic_router/layer.py
+++ b/semantic_router/layer.py
@@ -161,7 +161,7 @@ class RouteLayer:
encoder: Optional[BaseEncoder] = None,
llm: Optional[BaseLLM] = None,
routes: Optional[List[Route]] = None,
- index: Optional[BaseIndex] = LocalIndex(), # type: ignore
+ index: Optional[BaseIndex] = None, # type: ignore
):
logger.info("local")
self.index: BaseIndex = index if index is not None else LocalIndex()
@@ -280,6 +280,7 @@ class RouteLayer:
routes=[route.name] * len(route.utterances),
utterances=route.utterances,
)
+ self.routes.append(route)
def list_route_names(self) -> List[str]:
return [route.name for route in self.routes]
@@ -301,6 +302,25 @@ class RouteLayer:
self.routes = [route for route in self.routes if route.name != route_name]
self.index.delete(route_name=route_name)
+ def _refresh_routes(self):
+ """Pulls out the latest routes from the index.
+ """
+ raise NotImplementedError("This method has not yet been implemented.")
+ route_mapping = {route.name: route for route in self.routes}
+ index_routes = self.index.get_routes()
+ new_routes_names = []
+ new_routes = []
+ for route_name, utterance in index_routes:
+ if route_name in route_mapping:
+ if route_name not in new_routes_names:
+ existing_route = route_mapping[route_name]
+ new_routes.append(existing_route)
+
+ new_routes.append(Route(name=route_name, utterances=[utterance]))
+ route = route_mapping[route_name]
+ self.routes.append(route)
+
+
def _add_routes(self, routes: List[Route]):
# create embeddings for all routes
all_utterances = [
@@ -438,6 +458,9 @@ class RouteLayer:
correct += 1
accuracy = correct / len(Xq)
return accuracy
+
+ def _get_route_names(self) -> List[str]:
+ return [route.name for route in self.routes]
def threshold_random_search(
diff --git a/tests/unit/test_layer.py b/tests/unit/test_layer.py
index 3ccabea0a829376cf9439b588962a42ad6634134..9ddb1fe0f7a14644a8df7ca1a07354154fc3a10b 100644
--- a/tests/unit/test_layer.py
+++ b/tests/unit/test_layer.py
@@ -124,8 +124,8 @@ class TestRouteLayer:
assert route_layer.score_threshold == 0.82
assert len(route_layer.index) if route_layer.index is not None else 0 == 5
assert (
- len(set(route_layer.categories))
- if route_layer.categories is not None
+ len(set(route_layer._get_route_names()))
+ if route_layer._get_route_names() is not None
else 0 == 2
)
@@ -153,15 +153,18 @@ class TestRouteLayer:
def test_add_route(self, openai_encoder):
route_layer = RouteLayer(encoder=openai_encoder)
route1 = Route(name="Route 1", utterances=["Yes", "No"])
+ route2 = Route(name="Route 2", utterances=["Maybe", "Sure"])
- # Initially, the routes list should be empty and index should have no vectors
+ # Initially, the routes list should be empty
assert route_layer.routes == []
- assert route_layer.index.describe()['vectors'] == 0, "Index should initially have no vectors"
-
# Add route1 and check
route_layer.add(route=route1)
- assert route_layer.routes == [route1], "Route1 should be correctly added to routes list"
- assert route_layer.index.describe()['vectors'] == len(route1.utterances), "Index should reflect the correct number of vectors after adding route1"
+ assert route_layer.routes == [route1]
+ # TODO add length check
+ # Add route2 and check
+ route_layer.add(route=route2)
+ assert route_layer.routes == [route1, route2]
+ # TODO add length check
def test_list_route_names(self, openai_encoder, routes):
route_layer = RouteLayer(encoder=openai_encoder, routes=routes)
@@ -190,19 +193,16 @@ class TestRouteLayer:
# Attempt to remove a route that does not exist
non_existent_route = "non-existent-route"
with pytest.raises(ValueError) as excinfo:
- route_layer.remove(non_existent_route)
+ route_layer.delete(non_existent_route)
assert (
str(excinfo.value) == f"Route `{non_existent_route}` not found"
), "Attempting to remove a non-existent route should raise a ValueError."
def test_add_multiple_routes(self, openai_encoder, routes):
route_layer = RouteLayer(encoder=openai_encoder)
- # Assuming 'routes' fixture provides two routes with a total of 5 utterances
- total_utterances = sum(len(route.utterances) for route in routes)
-
route_layer._add_routes(routes=routes)
assert route_layer.index is not None
- assert route_layer.index.describe()['vectors'] == total_utterances, f"Index should contain {total_utterances} vectors after adding multiple routes"
+ # TODO add length check
def test_query_and_classification(self, openai_encoder, routes):
route_layer = RouteLayer(encoder=openai_encoder, routes=routes)
@@ -271,7 +271,7 @@ class TestRouteLayer:
route_layer_from_file = RouteLayer.from_json(temp.name)
assert (
route_layer_from_file.index is not None
- and route_layer_from_file.categories is not None
+ and route_layer_from_file._get_route_names() is not None
)
os.remove(temp.name)
@@ -284,7 +284,7 @@ class TestRouteLayer:
route_layer_from_file = RouteLayer.from_yaml(temp.name)
assert (
route_layer_from_file.index is not None
- and route_layer_from_file.categories is not None
+ and route_layer_from_file._get_route_names() is not None
)
os.remove(temp.name)
@@ -296,8 +296,8 @@ class TestRouteLayer:
assert layer_config.routes == routes
# now load from config and confirm it's the same
route_layer_from_config = RouteLayer.from_config(layer_config)
- assert (route_layer_from_config.index == route_layer.index).all()
- assert (route_layer_from_config.categories == route_layer.categories).all()
+ assert route_layer_from_config.index == route_layer.index
+ assert route_layer_from_config._get_route_names() == route_layer._get_route_names()
assert route_layer_from_config.score_threshold == route_layer.score_threshold
def test_get_thresholds(self, openai_encoder, routes):