|
|
|
@ -1,12 +1,12 @@
|
|
|
|
|
from typing import Optional
|
|
|
|
|
from swarms.memory.base import BaseVector
|
|
|
|
|
from swarms.memory.base import BaseVectorStore
|
|
|
|
|
import pinecone
|
|
|
|
|
from attr import define, field
|
|
|
|
|
from swarms.utils.hash import str_to_hash
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@define
|
|
|
|
|
class PineconeVectorStoreStore(BaseVector):
|
|
|
|
|
class PineconeVectorStoreStore(BaseVectorStore):
|
|
|
|
|
"""
|
|
|
|
|
PineconeVectorStore is a vector storage driver that uses Pinecone as the underlying storage engine.
|
|
|
|
|
|
|
|
|
@ -24,11 +24,11 @@ class PineconeVectorStoreStore(BaseVector):
|
|
|
|
|
Methods:
|
|
|
|
|
upsert_vector(vector: list[float], vector_id: Optional[str] = None, namespace: Optional[str] = None, meta: Optional[dict] = None, **kwargs) -> str:
|
|
|
|
|
Upserts a vector into the index.
|
|
|
|
|
load_entry(vector_id: str, namespace: Optional[str] = None) -> Optional[BaseVector.Entry]:
|
|
|
|
|
load_entry(vector_id: str, namespace: Optional[str] = None) -> Optional[BaseVectorStore.Entry]:
|
|
|
|
|
Loads a single vector from the index.
|
|
|
|
|
load_entries(namespace: Optional[str] = None) -> list[BaseVector.Entry]:
|
|
|
|
|
load_entries(namespace: Optional[str] = None) -> list[BaseVectorStore.Entry]:
|
|
|
|
|
Loads all vectors from the index.
|
|
|
|
|
query(query: str, count: Optional[int] = None, namespace: Optional[str] = None, include_vectors: bool = False, include_metadata=True, **kwargs) -> list[BaseVector.QueryResult]:
|
|
|
|
|
query(query: str, count: Optional[int] = None, namespace: Optional[str] = None, include_vectors: bool = False, include_metadata=True, **kwargs) -> list[BaseVectorStore.QueryResult]:
|
|
|
|
|
Queries the index for vectors similar to the given query string.
|
|
|
|
|
create_index(name: str, **kwargs) -> None:
|
|
|
|
|
Creates a new index.
|
|
|
|
@ -121,7 +121,7 @@ class PineconeVectorStoreStore(BaseVector):
|
|
|
|
|
|
|
|
|
|
def load_entry(
|
|
|
|
|
self, vector_id: str, namespace: Optional[str] = None
|
|
|
|
|
) -> Optional[BaseVector.Entry]:
|
|
|
|
|
) -> Optional[BaseVectorStore.Entry]:
|
|
|
|
|
"""Load entry"""
|
|
|
|
|
result = self.index.fetch(
|
|
|
|
|
ids=[vector_id], namespace=namespace
|
|
|
|
@ -131,7 +131,7 @@ class PineconeVectorStoreStore(BaseVector):
|
|
|
|
|
if len(vectors) > 0:
|
|
|
|
|
vector = vectors[0]
|
|
|
|
|
|
|
|
|
|
return BaseVector.Entry(
|
|
|
|
|
return BaseVectorStore.Entry(
|
|
|
|
|
id=vector["id"],
|
|
|
|
|
meta=vector["metadata"],
|
|
|
|
|
vector=vector["values"],
|
|
|
|
@ -142,7 +142,7 @@ class PineconeVectorStoreStore(BaseVector):
|
|
|
|
|
|
|
|
|
|
def load_entries(
|
|
|
|
|
self, namespace: Optional[str] = None
|
|
|
|
|
) -> list[BaseVector.Entry]:
|
|
|
|
|
) -> list[BaseVectorStore.Entry]:
|
|
|
|
|
"""Load entries"""
|
|
|
|
|
# This is a hacky way to query up to 10,000 values from Pinecone. Waiting on an official API for fetching
|
|
|
|
|
# all values from a namespace:
|
|
|
|
@ -156,7 +156,7 @@ class PineconeVectorStoreStore(BaseVector):
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
return [
|
|
|
|
|
BaseVector.Entry(
|
|
|
|
|
BaseVectorStore.Entry(
|
|
|
|
|
id=r["id"],
|
|
|
|
|
vector=r["values"],
|
|
|
|
|
meta=r["metadata"],
|
|
|
|
@ -174,12 +174,12 @@ class PineconeVectorStoreStore(BaseVector):
|
|
|
|
|
# PineconeVectorStoreStorageDriver-specific params:
|
|
|
|
|
include_metadata=True,
|
|
|
|
|
**kwargs,
|
|
|
|
|
) -> list[BaseVector.QueryResult]:
|
|
|
|
|
) -> list[BaseVectorStore.QueryResult]:
|
|
|
|
|
"""Query vectors"""
|
|
|
|
|
vector = self.embedding_driver.embed_string(query)
|
|
|
|
|
|
|
|
|
|
params = {
|
|
|
|
|
"top_k": count if count else BaseVector.DEFAULT_QUERY_COUNT,
|
|
|
|
|
"top_k": count if count else BaseVectorStore.DEFAULT_QUERY_COUNT,
|
|
|
|
|
"namespace": namespace,
|
|
|
|
|
"include_values": include_vectors,
|
|
|
|
|
"include_metadata": include_metadata,
|
|
|
|
@ -188,7 +188,7 @@ class PineconeVectorStoreStore(BaseVector):
|
|
|
|
|
results = self.index.query(vector, **params)
|
|
|
|
|
|
|
|
|
|
return [
|
|
|
|
|
BaseVector.QueryResult(
|
|
|
|
|
BaseVectorStore.QueryResult(
|
|
|
|
|
id=r["id"],
|
|
|
|
|
vector=r["values"],
|
|
|
|
|
score=r["score"],
|
|
|
|
|