llama-index
143 строки · 4.3 Кб
1"""
2Upstash vector store index.
3
4An index that is built with Upstash Vector.
5
6https://upstash.com/docs/vector/overall/getstarted
7"""
8
9import logging10from typing import Any, List11
12from llama_index.legacy.schema import BaseNode13from llama_index.legacy.utils import iter_batch14from llama_index.legacy.vector_stores.types import (15VectorStore,16VectorStoreQuery,17VectorStoreQueryMode,18VectorStoreQueryResult,19)
20from llama_index.legacy.vector_stores.utils import (21metadata_dict_to_node,22node_to_metadata_dict,23)
24
25logger = logging.getLogger(__name__)26
27DEFAULT_BATCH_SIZE = 12828
29
30class UpstashVectorStore(VectorStore):31"""32Upstash Vector Store.
33"""
34
35stores_text: bool = True36flat_metadata: bool = False37
38@classmethod39def class_name(cls) -> str:40return "UpstashVectorStore"41
42@property43def client(self) -> Any:44"""Return the Upstash client."""45return self._index46
47def __init__(48self, url: str, token: str, batch_size: int = DEFAULT_BATCH_SIZE49) -> None:50"""51Create a UpstashVectorStore. The index can be created using the Upstash console.
52
53Args:
54url (String): URL of the Upstash Vector instance, found in the Upstash console.
55token (String): Token for the Upstash Vector Index, found in the Upstash console.
56batch_size (Optional[int]): Batch size for adding nodes to the vector store.
57
58Raises:
59ImportError: If the upstash-vector python package is not installed.
60"""
61self.batch_size = batch_size62
63try:64from upstash_vector import Index65except ImportError:66raise ImportError(67"Could not import upstash_vector.Index, Please install it with `pip install upstash-vector`"68)69
70self._index = Index(url=url, token=token)71
72def add(self, nodes: List[BaseNode], **add_kwargs: Any) -> List[str]:73"""74Add nodes to the vector store.
75
76Args:
77nodes: List of nodes to add to the vector store.
78add_kwargs: Additional arguments to pass to the add method.
79
80Returns:
81List of ids of the added nodes.
82"""
83ids = []84vectors = []85for node_batch in iter_batch(nodes, self.batch_size):86for node in node_batch:87metadata_dict = node_to_metadata_dict(node)88ids.append(node.node_id)89vectors.append((node.node_id, node.embedding, metadata_dict))90
91self.client.upsert(vectors=vectors)92
93return ids94
95def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:96"""97Delete node from the vector store.
98
99Args:
100ref_doc_id: Reference doc id of the node to delete.
101delete_kwargs: Additional arguments to pass to the delete method.
102"""
103raise NotImplementedError(104"Delete is not currently supported, but will be in the future."105)106
107def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResult:108"""109Query the vector store.
110
111Args:
112query: Query to run against the vector store.
113kwargs: Additional arguments to pass to the query method.
114
115Returns:
116Query result.
117"""
118if query.mode != VectorStoreQueryMode.DEFAULT:119raise ValueError(f"Query mode {query.mode} not supported")120
121if query.filters:122raise ValueError("Metadata filtering not supported")123
124res = self.client.query(125vector=query.query_embedding,126top_k=query.similarity_top_k,127include_vectors=True,128include_metadata=True,129)130
131top_k_nodes = []132top_k_ids = []133top_k_scores = []134for vector in res:135node = metadata_dict_to_node(vector.metadata)136node.embedding = vector.vector137top_k_nodes.append(node)138top_k_ids.append(vector.id)139top_k_scores.append(vector.score)140
141return VectorStoreQueryResult(142nodes=top_k_nodes, similarities=top_k_scores, ids=top_k_ids143)144