Skip to content

Commit

Permalink
Merge pull request #165 from siddicky/fix/batch-upsert-issue
Browse files Browse the repository at this point in the history
fix: PineconeIndex upsert request size limit
  • Loading branch information
jamescalam committed Feb 26, 2024
2 parents 01eb8d7 + cfa13de commit 2f7d92c
Showing 1 changed file with 22 additions and 10 deletions.
32 changes: 22 additions & 10 deletions semantic_router/index/pinecone.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,21 +105,33 @@ def _init_index(self, force_create: bool = False) -> Union[Any, None]:
self.host = self.client.describe_index(self.index_name)["host"]
return index

def _batch_upsert(self, batch: List[dict]):
"""Helper method for upserting a single batch of records."""
if self.index is not None:
self.index.upsert(vectors=batch)
else:
raise ValueError("Index is None, could not upsert.")

def add(
self, embeddings: List[List[float]], routes: List[str], utterances: List[str]
self,
embeddings: List[List[float]],
routes: List[str],
utterances: List[str],
batch_size: int = 100,
):
"""Add vectors to Pinecone in batches."""
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
self.index = self._init_index(force_create=True)
vectors_to_upsert = []
for vector, route, utterance in zip(embeddings, routes, utterances):
record = PineconeRecord(values=vector, route=route, utterance=utterance)
vectors_to_upsert.append(record.to_dict())
if self.index is not None:
self.index.upsert(vectors=vectors_to_upsert)
else:
raise ValueError("Index is None could not upsert.")

vectors_to_upsert = [
PineconeRecord(values=vector, route=route, utterance=utterance).to_dict()
for vector, route, utterance in zip(embeddings, routes, utterances)
]

for i in range(0, len(vectors_to_upsert), batch_size):
batch = vectors_to_upsert[i : i + batch_size]
self._batch_upsert(batch)

def _get_route_ids(self, route_name: str):
clean_route = clean_route_name(route_name)
Expand Down

0 comments on commit 2f7d92c

Please sign in to comment.