Soumil RathiandClaude Opus 4.6
b9e26acf4d
Merge 190 commits from main including: - limit→top_k rename across all vector store APIs - custom_fact_extraction_prompt→custom_instructions rename - enable_graph flag removal (use self.graph truthiness) - US/Pacific→timezone.utc normalization - 18 Category A bug fixes (partial updates, graph cleanup on delete, telemetry gating, metadata mutation prevention, filter operator merging, delete timestamps, actor_id preservation, etc.) Conflict resolutions: - openai.py: merged _pass_dimensions_to_api with embed_batch - milvus.py: merged partial-update logic with BM25 text field - qdrant.py: took main's enhanced _create_filter, merged BM25 sparse vectors with partial-update endpoints - test_main.py: adopted main's test structure, removed obsolete test_custom_prompts, kept hybrid search over-fetch assertion AsyncMemory v3 port (full parity): - Rewrote AsyncMemory._add_to_vector_store with the same 8-phase batch pipeline as sync Memory (single LLM call, batch embed, hash dedup, batch persist, batch entity linking, message saving) - All blocking calls wrapped in await asyncio.to_thread() - Moved _build_session_scope to module level for shared access Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
333 lines
10 KiB
Python
333 lines
10 KiB
Python
import logging
|
|
from typing import Dict, List, Optional
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from mem0.vector_stores.base import VectorStoreBase
|
|
|
|
try:
|
|
from upstash_vector import Index
|
|
except ImportError:
|
|
raise ImportError("The 'upstash_vector' library is required. Please install it using 'pip install upstash_vector'.")
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class OutputData(BaseModel):
|
|
id: Optional[str] # memory id
|
|
score: Optional[float] # is None for `get` method
|
|
payload: Optional[Dict] # metadata
|
|
|
|
|
|
class UpstashVector(VectorStoreBase):
|
|
def __init__(
|
|
self,
|
|
collection_name: str,
|
|
url: Optional[str] = None,
|
|
token: Optional[str] = None,
|
|
client: Optional[Index] = None,
|
|
enable_embeddings: bool = False,
|
|
):
|
|
"""
|
|
Initialize the UpstashVector vector store.
|
|
|
|
Args:
|
|
url (str, optional): URL for Upstash Vector index. Defaults to None.
|
|
token (int, optional): Token for Upstash Vector index. Defaults to None.
|
|
client (Index, optional): Existing `upstash_vector.Index` client instance. Defaults to None.
|
|
namespace (str, optional): Default namespace for the index. Defaults to None.
|
|
"""
|
|
if client:
|
|
self.client = client
|
|
elif url and token:
|
|
self.client = Index(url, token)
|
|
else:
|
|
raise ValueError("Either a client or URL and token must be provided.")
|
|
|
|
self.collection_name = collection_name
|
|
|
|
self.enable_embeddings = enable_embeddings
|
|
|
|
def insert(
|
|
self,
|
|
vectors: List[list],
|
|
payloads: Optional[List[Dict]] = None,
|
|
ids: Optional[List[str]] = None,
|
|
):
|
|
"""
|
|
Insert vectors
|
|
|
|
Args:
|
|
vectors (list): List of vectors to insert.
|
|
payloads (list, optional): List of payloads corresponding to vectors. These will be passed as metadatas to the Upstash Vector client. Defaults to None.
|
|
ids (list, optional): List of IDs corresponding to vectors. Defaults to None.
|
|
"""
|
|
logger.info(f"Inserting {len(vectors)} vectors into namespace {self.collection_name}")
|
|
|
|
if self.enable_embeddings:
|
|
if not payloads or any("data" not in m or m["data"] is None for m in payloads):
|
|
raise ValueError("When embeddings are enabled, all payloads must contain a 'data' field.")
|
|
processed_vectors = [
|
|
{
|
|
"id": ids[i] if ids else None,
|
|
"data": payloads[i]["data"],
|
|
"metadata": payloads[i],
|
|
}
|
|
for i, v in enumerate(vectors)
|
|
]
|
|
else:
|
|
processed_vectors = [
|
|
{
|
|
"id": ids[i] if ids else None,
|
|
"vector": vectors[i],
|
|
"metadata": payloads[i] if payloads else None,
|
|
}
|
|
for i, v in enumerate(vectors)
|
|
]
|
|
|
|
self.client.upsert(
|
|
vectors=processed_vectors,
|
|
namespace=self.collection_name,
|
|
)
|
|
|
|
def _stringify(self, x):
|
|
return f'"{x}"' if isinstance(x, str) else x
|
|
|
|
def search(
|
|
self,
|
|
query: str,
|
|
vectors: List[list],
|
|
top_k: int = 5,
|
|
filters: Optional[Dict] = None,
|
|
) -> List[OutputData]:
|
|
"""
|
|
Search for similar vectors.
|
|
|
|
Args:
|
|
query (list): Query vector.
|
|
top_k (int, optional): Number of results to return. Defaults to 5.
|
|
filters (Dict, optional): Filters to apply to the search.
|
|
|
|
Returns:
|
|
List[OutputData]: Search results.
|
|
"""
|
|
|
|
filters_str = " AND ".join([f"{k} = {self._stringify(v)}" for k, v in filters.items()]) if filters else None
|
|
|
|
response = []
|
|
|
|
if self.enable_embeddings:
|
|
response = self.client.query(
|
|
data=query,
|
|
top_k=top_k,
|
|
filter=filters_str or "",
|
|
include_metadata=True,
|
|
namespace=self.collection_name,
|
|
)
|
|
else:
|
|
queries = [
|
|
{
|
|
"vector": v,
|
|
"top_k": top_k,
|
|
"filter": filters_str or "",
|
|
"include_metadata": True,
|
|
"namespace": self.collection_name,
|
|
}
|
|
for v in vectors
|
|
]
|
|
responses = self.client.query_many(queries=queries)
|
|
# flatten
|
|
response = [res for res_list in responses for res in res_list]
|
|
|
|
return [
|
|
OutputData(
|
|
id=res.id,
|
|
score=res.score,
|
|
payload=res.metadata,
|
|
)
|
|
for res in response
|
|
]
|
|
|
|
def keyword_search(self, query, top_k=5, filters=None):
|
|
"""
|
|
Perform keyword-based search using Upstash's BM25 sparse search.
|
|
|
|
Args:
|
|
query (str): The text query to search for.
|
|
top_k (int, optional): Number of results to return. Defaults to 5.
|
|
filters (Dict, optional): Filters to apply to the search.
|
|
|
|
Returns:
|
|
List[OutputData]: Search results, or None if sparse/BM25 search is not supported.
|
|
"""
|
|
try:
|
|
filters_str = (
|
|
" AND ".join([f"{k} = {self._stringify(v)}" for k, v in filters.items()])
|
|
if filters
|
|
else None
|
|
)
|
|
|
|
response = self.client.query(
|
|
data=query,
|
|
top_k=top_k,
|
|
filter=filters_str or "",
|
|
include_metadata=True,
|
|
namespace=self.collection_name,
|
|
)
|
|
|
|
return [
|
|
OutputData(
|
|
id=res.id,
|
|
score=res.score,
|
|
payload=res.metadata,
|
|
)
|
|
for res in response
|
|
]
|
|
except Exception as e:
|
|
logger.error(f"Error during keyword search for query '{query}': {e}")
|
|
return None
|
|
|
|
def delete(self, vector_id: int):
|
|
"""
|
|
Delete a vector by ID.
|
|
|
|
Args:
|
|
vector_id (int): ID of the vector to delete.
|
|
"""
|
|
self.client.delete(
|
|
ids=[str(vector_id)],
|
|
namespace=self.collection_name,
|
|
)
|
|
|
|
def update(
|
|
self,
|
|
vector_id: int,
|
|
vector: Optional[list] = None,
|
|
payload: Optional[dict] = None,
|
|
):
|
|
"""
|
|
Update a vector and its payload.
|
|
|
|
Args:
|
|
vector_id (int): ID of the vector to update.
|
|
vector (list, optional): Updated vector. Defaults to None.
|
|
payload (dict, optional): Updated payload. Defaults to None.
|
|
"""
|
|
self.client.update(
|
|
id=str(vector_id),
|
|
vector=vector,
|
|
data=payload.get("data") if payload else None,
|
|
metadata=payload,
|
|
namespace=self.collection_name,
|
|
)
|
|
|
|
def get(self, vector_id: int) -> Optional[OutputData]:
|
|
"""
|
|
Retrieve a vector by ID.
|
|
|
|
Args:
|
|
vector_id (int): ID of the vector to retrieve.
|
|
|
|
Returns:
|
|
dict: Retrieved vector.
|
|
"""
|
|
response = self.client.fetch(
|
|
ids=[str(vector_id)],
|
|
namespace=self.collection_name,
|
|
include_metadata=True,
|
|
)
|
|
if len(response) == 0:
|
|
return None
|
|
vector = response[0]
|
|
if not vector:
|
|
return None
|
|
return OutputData(id=vector.id, score=None, payload=vector.metadata)
|
|
|
|
def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> List[List[OutputData]]:
|
|
"""
|
|
List all memories.
|
|
Args:
|
|
filters (Dict, optional): Filters to apply to the search. Defaults to None.
|
|
top_k (int, optional): Number of results to return. Defaults to 100.
|
|
Returns:
|
|
List[OutputData]: Search results.
|
|
"""
|
|
filters_str = " AND ".join([f"{k} = {self._stringify(v)}" for k, v in filters.items()]) if filters else None
|
|
|
|
info = self.client.info()
|
|
ns_info = info.namespaces.get(self.collection_name)
|
|
|
|
if not ns_info or ns_info.vector_count == 0:
|
|
return [[]]
|
|
|
|
random_vector = [1.0] * self.client.info().dimension
|
|
|
|
results, query = self.client.resumable_query(
|
|
vector=random_vector,
|
|
filter=filters_str or "",
|
|
include_metadata=True,
|
|
namespace=self.collection_name,
|
|
top_k=100,
|
|
)
|
|
with query:
|
|
while True:
|
|
if len(results) >= top_k:
|
|
break
|
|
res = query.fetch_next(100)
|
|
if not res:
|
|
break
|
|
results.extend(res)
|
|
|
|
parsed_result = [
|
|
OutputData(
|
|
id=res.id,
|
|
score=res.score,
|
|
payload=res.metadata,
|
|
)
|
|
for res in results
|
|
]
|
|
return [parsed_result]
|
|
|
|
def create_col(self, name, vector_size, distance):
|
|
"""
|
|
Upstash Vector has namespaces instead of collections. A namespace is created when the first vector is inserted.
|
|
|
|
This method is a placeholder to maintain the interface.
|
|
"""
|
|
pass
|
|
|
|
def list_cols(self) -> List[str]:
|
|
"""
|
|
Lists all namespaces in the Upstash Vector index.
|
|
Returns:
|
|
List[str]: List of namespaces.
|
|
"""
|
|
return self.client.list_namespaces()
|
|
|
|
def delete_col(self):
|
|
"""
|
|
Delete the namespace and all vectors in it.
|
|
"""
|
|
self.client.reset(namespace=self.collection_name)
|
|
pass
|
|
|
|
def col_info(self):
|
|
"""
|
|
Return general information about the Upstash Vector index.
|
|
|
|
- Total number of vectors across all namespaces
|
|
- Total number of vectors waiting to be indexed across all namespaces
|
|
- Total size of the index on disk in bytes
|
|
- Vector dimension
|
|
- Similarity function used
|
|
- Per-namespace vector and pending vector counts
|
|
"""
|
|
return self.client.info()
|
|
|
|
def reset(self):
|
|
"""
|
|
Reset the Upstash Vector index.
|
|
"""
|
|
self.delete_col()
|