[FEAT][BaseVectorStore]

pull/207/head^2
Kye 1 year ago
parent ffb8e5a1d5
commit dba245a725

@ -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"],

Loading…
Cancel
Save