diff --git a/poetry.lock b/poetry.lock
index d5e1153b3e2611cb158cf0184a4b1a37c36e982d..361cedea80e4a8bf5fe4bd6af59df6819c15ba2d 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -3218,7 +3218,7 @@ files = [
name = "types-requests"
version = "2.31.0.20240125"
description = "Typing stubs for requests"
-optional = true
+optional = false
python-versions = ">=3.8"
files = [
{file = "types-requests-2.31.0.20240125.tar.gz", hash = "sha256:03a28ce1d7cd54199148e043b2079cdded22d6795d19a2c2a6791a4b2b5e2eb5"},
@@ -3404,4 +3404,4 @@ pinecone = ["pinecone-client"]
[metadata]
lock-version = "2.0"
python-versions = ">=3.9,<3.13"
-content-hash = "52ce34492a7d4827a3c2b96332e7285369209dbe9ec2a9488d8eac2c13d4d0c6"
+content-hash = "23b8995b7eee4df19bb2242d6a81de8557a855053b4346a532efa63be2ea303f"
diff --git a/pyproject.toml b/pyproject.toml
index 1491b6d7a267fa493e93162e8c55e924b7003ee8..e1137235108ab96511a5b6dc3102da0784715331 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,6 +45,7 @@ pytest-cov = "^4.1.0"
pytest-xdist = "^3.5.0"
mypy = "^1.7.1"
types-pyyaml = "^6.0.12.12"
+types-requests = "^2.31.0"
[build-system]
requires = ["poetry-core"]
diff --git a/semantic_router/index/base.py b/semantic_router/index/base.py
index 37550281e77acd87ad180dc025452fcfba5fe21d..c3c9ed0fb740e7195e723965a23dddedeb2948bc 100644
--- a/semantic_router/index/base.py
+++ b/semantic_router/index/base.py
@@ -18,7 +18,9 @@ class BaseIndex(BaseModel):
dimensions: Union[int, None] = None
type: str = "base"
- def add(self, embeddings: List[float], routes: List[str], utterances: List[str]):
+ def add(
+ self, embeddings: List[List[float]], routes: List[str], utterances: List[str]
+ ):
"""
Add embeddings to the index.
This method should be implemented by subclasses.
@@ -32,20 +34,20 @@ class BaseIndex(BaseModel):
"""
raise NotImplementedError("This method should be implemented by subclasses.")
- def describe(self) -> bool:
+ def describe(self) -> dict:
"""
Returns a dictionary 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 query(self, vector: np.ndarray, top_k: int = 5) -> Tuple[np.ndarray, np.ndarray]:
+ def query(self, vector: np.ndarray, top_k: int = 5) -> Tuple[np.ndarray, List[str]]:
"""
Search the index for the query_vector and return top_k results.
This method should be implemented by subclasses.
"""
raise NotImplementedError("This method should be implemented by subclasses.")
-
+
def delete_index(self):
"""
Deletes or resets the index.
@@ -53,6 +55,5 @@ class BaseIndex(BaseModel):
"""
raise NotImplementedError("This method should be implemented by subclasses.")
-
class Config:
arbitrary_types_allowed = True
diff --git a/semantic_router/index/local.py b/semantic_router/index/local.py
index 72c474ae7de395d19a9b472e03d6b2c57a1804c0..02a1f209fc17052e1c8fc52c6198347a347e80de 100644
--- a/semantic_router/index/local.py
+++ b/semantic_router/index/local.py
@@ -1,19 +1,23 @@
import numpy as np
-from typing import List, Any, Tuple, Optional
+from typing import List, Tuple, Optional
from semantic_router.linear import similarity_matrix, top_scores
from semantic_router.index.base import BaseIndex
class LocalIndex(BaseIndex):
-
def __init__(self):
- super().__init__()
+ super().__init__()
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
- def add(self, embeddings: List[List[float]], routes: List[str], utterances: List[str]):
+ def add(
+ self, embeddings: List[List[float]], routes: List[str], utterances: List[str]
+ ):
embeds = np.array(embeddings) # type: ignore
routes_arr = np.array(routes)
utterances_arr = np.array(utterances)
@@ -27,44 +31,50 @@ class LocalIndex(BaseIndex):
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.
- """
- idx = [
- i for i, route in enumerate(self.routes)
- if route == route_name
- ]
+ """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):
"""
Delete all records of a specific route from the index.
"""
- if self.index is not None:
+ 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."
+ )
- def describe(self):
+ 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
+ "vectors": self.index.shape[0] if self.index is not None else 0,
}
def query(self, vector: np.ndarray, top_k: int = 5) -> Tuple[np.ndarray, List[str]]:
"""
Search the index for the query and return top_k results.
"""
- if self.index is None:
- raise ValueError("Index is not populated.")
+ if self.index is None or self.routes is None:
+ raise ValueError("Index or routes are not populated.")
sim = similarity_matrix(vector, self.index)
# extract the index values of top scoring vectors
scores, idx = top_scores(sim, top_k)
# get routes from index values
route_names = self.routes[idx].copy()
return scores, route_names
-
+
def delete_index(self):
"""
Deletes the index, effectively clearing it and setting it to None.
diff --git a/semantic_router/index/pinecone.py b/semantic_router/index/pinecone.py
index 391fb9a65b5ceb46627006b3b666e1328784cfe1..d2b21d95f9ff176493f91d91050eed204125aa45 100644
--- a/semantic_router/index/pinecone.py
+++ b/semantic_router/index/pinecone.py
@@ -12,6 +12,7 @@ import numpy as np
def clean_route_name(route_name: str) -> str:
return route_name.strip().replace(" ", "-")
+
class PineconeRecord(BaseModel):
id: str = ""
values: List[float]
@@ -29,10 +30,7 @@ class PineconeRecord(BaseModel):
return {
"id": self.id,
"values": self.values,
- "metadata": {
- "sr_route": self.route,
- "sr_utterance": self.utterance
- }
+ "metadata": {"sr_route": self.route, "sr_utterance": self.utterance},
}
@@ -42,14 +40,14 @@ class PineconeIndex(BaseIndex):
dimensions: Union[int, None] = None
metric: str = "cosine"
cloud: str = "aws"
- region: str = "us-west-2"
+ region: str = "us-west-2"
host: str = ""
client: Any = Field(default=None, exclude=True)
index: Optional[Any] = Field(default=None, exclude=True)
ServerlessSpec: Any = Field(default=None, exclude=True)
def __init__(self, **data):
- super().__init__(**data)
+ super().__init__(**data)
self._initialize_client()
self.type = "pinecone"
@@ -62,6 +60,7 @@ class PineconeIndex(BaseIndex):
def _initialize_client(self, api_key: Optional[str] = None):
try:
from pinecone import Pinecone, ServerlessSpec
+
self.ServerlessSpec = ServerlessSpec
except ImportError:
raise ImportError(
@@ -81,13 +80,10 @@ class PineconeIndex(BaseIndex):
# if the index doesn't exist and we have dimension value
# we create the index
self.client.create_index(
- name=self.index_name,
- dimension=self.dimensions,
+ name=self.index_name,
+ dimension=self.dimensions,
metric=self.metric,
- spec=self.ServerlessSpec(
- cloud=self.cloud,
- region=self.region
- )
+ spec=self.ServerlessSpec(cloud=self.cloud, region=self.region),
)
# wait for index to be created
while not self.client.describe_index(self.index_name).status["ready"]:
@@ -100,7 +96,9 @@ class PineconeIndex(BaseIndex):
# grab the dimensions from the index
self.dimensions = index.describe_index_stats()["dimension"]
elif force_create and not dimensions_given:
- raise ValueError("Cannot create an index without specifying the dimensions.")
+ raise ValueError(
+ "Cannot create an index without specifying the dimensions."
+ )
else:
# if the index doesn't exist and we don't have the dimensions
# we return None
@@ -109,8 +107,10 @@ class PineconeIndex(BaseIndex):
if index is not None:
self.host = self.client.describe_index(self.index_name)["host"]
return index
-
- def add(self, embeddings: List[List[float]], routes: List[str], utterances: List[str]):
+
+ def add(
+ self, embeddings: List[List[float]], routes: List[str], utterances: List[str]
+ ):
if self.index is None:
self.dimensions = self.dimensions or len(embeddings[0])
# we set force_create to True as we MUST have an index to add data
@@ -119,35 +119,46 @@ class PineconeIndex(BaseIndex):
for vector, route, utterance in zip(embeddings, routes, utterances):
record = PineconeRecord(values=vector, route=route, utterance=utterance)
vectors_to_upsert.append(record.to_dict())
- self.index.upsert(vectors=vectors_to_upsert)
+ if self.index is not None:
+ self.index.upsert(vectors=vectors_to_upsert)
+ else:
+ raise ValueError("Index is None could not upsert.")
def _get_route_vecs(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"]}
+ headers={"Api-Key": os.environ["PINECONE_API_KEY"]},
)
return [vec["id"] for vec in res.json()["vectors"]]
def delete(self, route_name: str):
route_vec_ids = self._get_route_vecs(route_name=route_name)
- self.index.delete(ids=route_vec_ids)
+ if self.index is not None:
+ self.index.delete(ids=route_vec_ids)
+ else:
+ raise ValueError("Index is None, could not delete.")
def delete_all(self):
self.index.delete(delete_all=True)
- def describe(self) -> bool:
- stats = self.index.describe_index_stats()
- return {
- "type": self.type,
- "dimensions": stats["dimension"],
- "vectors": stats["total_vector_count"]
- }
-
+ def describe(self) -> dict:
+ if self.index is not None:
+ stats = self.index.describe_index_stats()
+ return {
+ "type": self.type,
+ "dimensions": stats["dimension"],
+ "vectors": stats["total_vector_count"],
+ }
+ else:
+ raise ValueError("Index is None, cannot describe index stats.")
+
def query(self, vector: np.ndarray, top_k: int = 5) -> Tuple[np.ndarray, List[str]]:
+ if self.index is None:
+ raise ValueError("Index is not populated.")
query_vector_list = vector.tolist()
results = self.index.query(
- vector=[query_vector_list],
+ vector=[query_vector_list],
top_k=top_k,
include_metadata=True,
)
@@ -159,4 +170,4 @@ class PineconeIndex(BaseIndex):
self.client.delete_index(self.index_name)
def __len__(self):
- return self.index.describe_index_stats()["total_vector_count"]
\ No newline at end of file
+ return self.index.describe_index_stats()["total_vector_count"]
diff --git a/semantic_router/layer.py b/semantic_router/layer.py
index 21d208f1d5a2c6813d5c24190f8c3bed4d24c2ae..573c24d669c69726d7091aee30e1d4f3da22d081 100644
--- a/semantic_router/layer.py
+++ b/semantic_router/layer.py
@@ -161,10 +161,10 @@ class RouteLayer:
encoder: Optional[BaseEncoder] = None,
llm: Optional[BaseLLM] = None,
routes: Optional[List[Route]] = None,
- index: Optional[BaseIndex] = LocalIndex,
+ index: Optional[BaseIndex] = LocalIndex, # type: ignore
):
logger.info("local")
- self.index: BaseIndex = index
+ self.index: BaseIndex = index if index is not None else LocalIndex()
if encoder is None:
logger.warning(
"No encoder provided. Using default OpenAIEncoder. Ensure "
@@ -283,13 +283,13 @@ class RouteLayer:
def list_route_names(self) -> List[str]:
return [route.name for route in self.routes]
-
+
def update(self, route_name: str, utterances: List[str]):
raise NotImplementedError("This method has not yet been implemented.")
def delete(self, route_name: str):
"""Deletes a route given a specific route name.
-
+
:param route_name: the name of the route to be deleted
:type str:
"""
@@ -303,7 +303,9 @@ class RouteLayer:
def _add_routes(self, routes: List[Route]):
# create embeddings for all routes
- all_utterances = [utterance for route in routes for utterance in route.utterances]
+ all_utterances = [
+ utterance for route in routes for utterance in route.utterances
+ ]
embedded_utterances = self.encoder(all_utterances)
# create route array
route_names = [route.name for route in routes for _ in route.utterances]
@@ -311,7 +313,7 @@ class RouteLayer:
self.index.add(
embeddings=embedded_utterances,
routes=route_names,
- utterances=all_utterances
+ utterances=all_utterances,
)
def _encode(self, text: str) -> Any:
diff --git a/semantic_router/schema.py b/semantic_router/schema.py
index 2ac7280a8c25398047b916bbbb495b24dcc22b72..46ee7f590c3275b314b9ecb3e52adc54f16012fe 100644
--- a/semantic_router/schema.py
+++ b/semantic_router/schema.py
@@ -11,9 +11,6 @@ from semantic_router.encoders import (
OpenAIEncoder,
)
-from semantic_router.index.local import LocalIndex
-from semantic_router.index.pinecone import PineconeIndex
-from semantic_router.index.base import BaseIndex
class EncoderType(Enum):
HUGGINGFACE = "huggingface"