fix(sdk): removing deprecating param from our sdk and docs changes with it (#4740)

This commit is contained in:
Kartik
2026-04-12 00:34:58 +05:30
committed by GitHub
parent 3882af7450
commit e44b46ef2e
138 changed files with 1186 additions and 2840 deletions
+8 -8
View File
@@ -205,14 +205,14 @@ class AzureAISearch(VectorStoreBase):
filter_expression = " and ".join(filter_conditions)
return filter_expression
def search(self, query, vectors, limit=5, filters=None):
def search(self, query, vectors, top_k=5, filters=None):
"""
Search for similar vectors.
Args:
query (str): Query.
vectors (List[float]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -222,13 +222,13 @@ class AzureAISearch(VectorStoreBase):
if filters:
filter_expression = self._build_filter_expression(filters)
vector_query = VectorizedQuery(vector=vectors, k_nearest_neighbors=limit, fields="vector")
vector_query = VectorizedQuery(vector=vectors, k_nearest_neighbors=top_k, fields="vector")
if self.hybrid_search:
search_results = self.search_client.search(
search_text=query,
vector_queries=[vector_query],
filter=filter_expression,
top=limit,
top=top_k,
vector_filter_mode=self.vector_filter_mode,
search_fields=["payload"],
)
@@ -236,7 +236,7 @@ class AzureAISearch(VectorStoreBase):
search_results = self.search_client.search(
vector_queries=[vector_query],
filter=filter_expression,
top=limit,
top=top_k,
vector_filter_mode=self.vector_filter_mode,
)
@@ -327,13 +327,13 @@ class AzureAISearch(VectorStoreBase):
index = self.index_client.get_index(self.index_name)
return {"name": index.name, "fields": index.fields}
def list(self, filters=None, limit=100):
def list(self, filters=None, top_k=100):
"""
List all vectors in the index.
Args:
filters (dict, optional): Filters to apply to the list.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors.
@@ -342,7 +342,7 @@ class AzureAISearch(VectorStoreBase):
if filters:
filter_expression = self._build_filter_expression(filters)
search_results = self.search_client.search(search_text="*", filter=filter_expression, top=limit)
search_results = self.search_client.search(search_text="*", filter=filter_expression, top=top_k)
results = []
for result in search_results:
payload = json.loads(extract_json(result["payload"]))
+8 -8
View File
@@ -7,8 +7,8 @@ from pydantic import BaseModel
try:
import pymysql
from pymysql.cursors import DictCursor
from dbutils.pooled_db import PooledDB
from pymysql.cursors import DictCursor
except ImportError:
raise ImportError(
"Azure MySQL vector store requires PyMySQL and DBUtils. "
@@ -243,7 +243,7 @@ class AzureMySQL(VectorStoreBase):
self,
query: str,
vectors: List[float],
limit: int = 5,
top_k: int = 5,
filters: Optional[Dict] = None,
) -> List[OutputData]:
"""
@@ -252,7 +252,7 @@ class AzureMySQL(VectorStoreBase):
Args:
query (str): Query string (not used in vector search)
vectors (List[float]): Query vector
limit (int): Number of results to return
top_k (int): Number of results to return
filters (Dict, optional): Filters to apply to the search
Returns:
@@ -291,9 +291,9 @@ class AzureMySQL(VectorStoreBase):
distance = 1 - similarity
scored_results.append((row['id'], distance, row['payload']))
# Sort by distance and limit
# Sort by distance and apply limit
scored_results.sort(key=lambda x: x[1])
scored_results = scored_results[:limit]
scored_results = scored_results[:top_k]
return [
OutputData(id=r[0], score=float(r[1]), payload=json.loads(r[2]) if isinstance(r[2], str) else r[2])
@@ -406,14 +406,14 @@ class AzureMySQL(VectorStoreBase):
def list(
self,
filters: Optional[Dict] = None,
limit: int = 100
top_k: int = 100
) -> List[List[OutputData]]:
"""
List all vectors in the collection.
Args:
filters (Dict, optional): Filters to apply
limit (int): Number of vectors to return
top_k (int): Number of vectors to return
Returns:
List[List[OutputData]]: List of vectors
@@ -436,7 +436,7 @@ class AzureMySQL(VectorStoreBase):
{filter_clause}
LIMIT %s
""",
(*filter_params, limit)
(*filter_params, top_k)
)
results = cur.fetchall()
+6 -6
View File
@@ -185,14 +185,14 @@ class BaiduDB(VectorStoreBase):
row = Row(id=idx, vector=vector, metadata=metadata)
self._table.upsert(rows=[row])
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> list:
def search(self, query: str, vectors: list, top_k: int = 5, filters: dict = None) -> list:
"""
Search for similar vectors.
Args:
query (str): Query string.
vectors (List[float]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -207,7 +207,7 @@ class BaiduDB(VectorStoreBase):
request = VectorTopkSearchRequest(
vector_field="vector",
vector=FloatVector(vectors),
limit=limit,
limit=top_k,
filter=search_filter,
config=VectorSearchConfig(ef=200),
)
@@ -313,20 +313,20 @@ class BaiduDB(VectorStoreBase):
"""
return self._table.stats()
def list(self, filters: dict = None, limit: int = 100) -> list:
def list(self, filters: dict = None, top_k: int = 100) -> list:
"""
List all vectors in the table.
Args:
filters (Dict, optional): Filters to apply to the list.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors.
"""
projections = ["id", "metadata"]
list_filter = self._create_filter(filters) if filters else None
result = self._table.select(filter=list_filter, projections=projections, limit=limit)
result = self._table.select(filter=list_filter, projections=projections, limit=top_k)
memories = []
for row in result.rows:
+2 -2
View File
@@ -13,7 +13,7 @@ class VectorStoreBase(ABC):
pass
@abstractmethod
def search(self, query, vectors, limit=5, filters=None):
def search(self, query, vectors, top_k=5, filters=None):
"""Search for similar vectors."""
pass
@@ -48,7 +48,7 @@ class VectorStoreBase(ABC):
pass
@abstractmethod
def list(self, filters=None, limit=None):
def list(self, filters=None, top_k=None):
"""List all memories."""
pass
+8 -8
View File
@@ -7,8 +7,8 @@ import numpy as np
from pydantic import BaseModel
try:
from cassandra.cluster import Cluster
from cassandra.auth import PlainTextAuthProvider
from cassandra.cluster import Cluster
except ImportError:
raise ImportError(
"Apache Cassandra vector store requires cassandra-driver. "
@@ -214,7 +214,7 @@ class CassandraDB(VectorStoreBase):
self,
query: str,
vectors: List[float],
limit: int = 5,
top_k: int = 5,
filters: Optional[Dict] = None,
) -> List[OutputData]:
"""
@@ -223,7 +223,7 @@ class CassandraDB(VectorStoreBase):
Args:
query (str): Query string (not used in vector search)
vectors (List[float]): Query vector
limit (int): Number of results to return
top_k (int): Number of results to return
filters (Dict, optional): Filters to apply to the search
Returns:
@@ -263,9 +263,9 @@ class CassandraDB(VectorStoreBase):
scored_results.append((row.id, distance, row.payload))
# Sort by distance and limit
# Sort by distance and apply limit
scored_results.sort(key=lambda x: x[1])
scored_results = scored_results[:limit]
scored_results = scored_results[:top_k]
return [
OutputData(
@@ -427,14 +427,14 @@ class CassandraDB(VectorStoreBase):
def list(
self,
filters: Optional[Dict] = None,
limit: int = 100
top_k: int = 100
) -> List[List[OutputData]]:
"""
List all vectors in the collection.
Args:
filters (Dict, optional): Filters to apply
limit (int): Number of vectors to return
top_k (int): Number of vectors to return
Returns:
List[List[OutputData]]: List of vectors
@@ -443,7 +443,7 @@ class CassandraDB(VectorStoreBase):
query = f"""
SELECT id, vector, payload
FROM {self.keyspace}.{self.collection_name}
LIMIT {limit}
LIMIT {top_k}
"""
rows = self.session.execute(query)
+6 -6
View File
@@ -141,7 +141,7 @@ class ChromaDB(VectorStoreBase):
self.collection.add(ids=ids, embeddings=vectors, metadatas=payloads)
def search(
self, query: str, vectors: List[list], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[list], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
@@ -149,14 +149,14 @@ class ChromaDB(VectorStoreBase):
Args:
query (str): Query.
vectors (List[list]): List of vectors to search.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Optional[Dict], optional): Filters to apply to the search. Defaults to None.
Returns:
List[OutputData]: Search results.
"""
where_clause = self._generate_where_clause(filters) if filters else None
results = self.collection.query(query_embeddings=vectors, where=where_clause, n_results=limit)
results = self.collection.query(query_embeddings=vectors, where=where_clause, n_results=top_k)
final_results = self._parse_output(results)
return final_results
@@ -222,19 +222,19 @@ class ChromaDB(VectorStoreBase):
"""
return self.client.get_collection(name=self.collection_name)
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[OutputData]:
def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> List[OutputData]:
"""
List all vectors in a collection.
Args:
filters (Optional[Dict], optional): Filters to apply to the list. Defaults to None.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors.
"""
where_clause = self._generate_where_clause(filters) if filters else None
results = self.collection.get(where=where_clause, limit=limit)
results = self.collection.get(where=where_clause, limit=top_k)
return [self._parse_output(results)]
def reset(self):
+19 -11
View File
@@ -2,20 +2,28 @@ import json
import logging
import re
import uuid
from typing import Optional, List
from datetime import datetime, date
from databricks.sdk.service.catalog import ColumnInfo, ColumnTypeName, TableType, DataSourceFormat
from databricks.sdk.service.catalog import TableConstraint, PrimaryKeyConstraint
from datetime import date, datetime
from typing import List, Optional
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.catalog import (
ColumnInfo,
ColumnTypeName,
DataSourceFormat,
PrimaryKeyConstraint,
TableConstraint,
TableType,
)
from databricks.sdk.service.sql import StatementParameterListItem
from databricks.sdk.service.vectorsearch import (
VectorIndexType,
DeltaSyncVectorIndexSpecRequest,
DirectAccessVectorIndexSpec,
EmbeddingSourceColumn,
EmbeddingVectorColumn,
VectorIndexType,
)
from pydantic import BaseModel
from mem0.memory.utils import extract_json
from mem0.vector_stores.base import VectorStoreBase
@@ -446,14 +454,14 @@ class Databricks(VectorStoreBase):
logger.error(f"Insert operation failed: {e}")
raise
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> List[MemoryResult]:
def search(self, query: str, vectors: list, top_k: int = 5, filters: dict = None) -> List[MemoryResult]:
"""
Search for similar vectors or text using the Databricks Vector Search index.
Args:
query (str): Search query text (for text-based search).
vectors (list): Query vector (for vector-based search).
limit (int): Maximum number of results.
top_k (int): Maximum number of results.
filters (dict): Filters to apply.
Returns:
@@ -468,7 +476,7 @@ class Databricks(VectorStoreBase):
query_kwargs = {
"index_name": self.fully_qualified_index_name,
"columns": self.column_names,
"num_results": limit,
"num_results": top_k,
"query_type": self.query_type,
"filters_json": filters_json,
}
@@ -719,20 +727,20 @@ class Databricks(VectorStoreBase):
logger.error(f"Failed to get info for index '{name or self.index_name}': {e}")
raise
def list(self, filters: dict = None, limit: int = None) -> list[MemoryResult]:
def list(self, filters: dict = None, top_k: int = None) -> list[MemoryResult]:
"""
List all recent created memories from the vector store.
Args:
filters (dict, optional): Filters to apply.
limit (int, optional): Maximum number of results.
top_k (int, optional): Maximum number of results.
Returns:
List containing list of MemoryResult objects.
"""
try:
filters_json = json.dumps(filters) if filters else None
num_results = limit or 100
num_results = top_k or 100
columns = self.column_names
# Use query_text for Delta Sync with model endpoint, query_vector otherwise
query_kwargs = {
+6 -6
View File
@@ -129,7 +129,7 @@ class ElasticsearchDB(VectorStoreBase):
return results
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search with two options:
@@ -137,10 +137,10 @@ class ElasticsearchDB(VectorStoreBase):
2. Use KNN search on vectors with pre-filtering if no custom search query is provided
"""
if self.custom_search_query:
search_query = self.custom_search_query(vectors, limit, filters)
search_query = self.custom_search_query(vectors, top_k, filters)
else:
search_query = {
"knn": {"field": "vector", "query_vector": vectors, "k": limit, "num_candidates": limit * 2}
"knn": {"field": "vector", "query_vector": vectors, "k": top_k, "num_candidates": top_k * 2}
}
if filters:
filter_conditions = []
@@ -203,7 +203,7 @@ class ElasticsearchDB(VectorStoreBase):
"""Get information about a collection (index)."""
return self.client.indices.get(index=name)
def list(self, filters: Optional[Dict] = None, limit: Optional[int] = None) -> List[List[OutputData]]:
def list(self, filters: Optional[Dict] = None, top_k: Optional[int] = None) -> List[List[OutputData]]:
"""List all memories."""
query: Dict[str, Any] = {"query": {"match_all": {}}}
@@ -213,8 +213,8 @@ class ElasticsearchDB(VectorStoreBase):
filter_conditions.append({"term": {f"metadata.{key}": value}})
query["query"] = {"bool": {"must": filter_conditions}}
if limit:
query["size"] = limit
if top_k:
query["size"] = top_k
response = self.client.search(index=self.collection_name, body=query)
+15 -16
View File
@@ -2,14 +2,13 @@ import logging
import os
import pickle
import uuid
import warnings
from pathlib import Path
from typing import Dict, List, Optional
import numpy as np
from pydantic import BaseModel
import warnings
try:
# Suppress SWIG deprecation warnings from FAISS
warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*SwigPy.*")
@@ -115,23 +114,23 @@ class FAISS(VectorStoreBase):
except Exception as e:
logger.warning(f"Failed to save FAISS index: {e}")
def _parse_output(self, scores, ids, limit=None) -> List[OutputData]:
def _parse_output(self, scores, ids, top_k=None) -> List[OutputData]:
"""
Parse the output data.
Args:
scores: Similarity scores from FAISS.
ids: Indices from FAISS.
limit: Maximum number of results to return.
top_k: Maximum number of results to return.
Returns:
List[OutputData]: Parsed output data.
"""
if limit is None:
limit = len(ids)
if top_k is None:
top_k = len(ids)
results = []
for i in range(min(len(ids), limit)):
for i in range(min(len(ids), top_k)):
if ids[i] == -1: # FAISS returns -1 for empty results
continue
@@ -225,7 +224,7 @@ class FAISS(VectorStoreBase):
logger.info(f"Inserted {len(vectors)} vectors into collection {self.collection_name}")
def search(
self, query: str, vectors: List[list], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[list], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
@@ -233,7 +232,7 @@ class FAISS(VectorStoreBase):
Args:
query (str): Query (not used, kept for API compatibility).
vectors (List[list]): List of vectors to search.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Optional[Dict], optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -250,19 +249,19 @@ class FAISS(VectorStoreBase):
if self.normalize_L2 and self.distance_strategy.lower() == "euclidean":
faiss.normalize_L2(query_vectors)
fetch_k = limit * 2 if filters else limit
fetch_k = top_k * 2 if filters else top_k
scores, indices = self.index.search(query_vectors, fetch_k)
results = self._parse_output(scores[0], indices[0], limit)
results = self._parse_output(scores[0], indices[0], top_k)
if filters:
filtered_results = []
for result in results:
if self._apply_filters(result.payload, filters):
filtered_results.append(result)
if len(filtered_results) >= limit:
if len(filtered_results) >= top_k:
break
results = filtered_results[:limit]
results = filtered_results[:top_k]
return results
@@ -450,13 +449,13 @@ class FAISS(VectorStoreBase):
"distance": self.distance_strategy,
}
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[OutputData]:
def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> List[OutputData]:
"""
List all vectors in a collection.
Args:
filters (Optional[Dict], optional): Filters to apply to the list. Defaults to None.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors.
@@ -482,7 +481,7 @@ class FAISS(VectorStoreBase):
)
count += 1
if count >= limit:
if count >= top_k:
break
return [results]
+5 -5
View File
@@ -91,15 +91,15 @@ class Langchain(VectorStoreBase):
texts = [payload.get("data", "") for payload in payloads] if payloads else [""] * len(vectors)
self.client.add_texts(texts=texts, metadatas=payloads, ids=ids)
def search(self, query: str, vectors: List[List[float]], limit: int = 5, filters: Optional[Dict] = None):
def search(self, query: str, vectors: List[List[float]], top_k: int = 5, filters: Optional[Dict] = None):
"""
Search for similar vectors in LangChain.
"""
# For each vector, perform a similarity search
if filters:
results = self.client.similarity_search_by_vector(embedding=vectors, k=limit, filter=filters)
results = self.client.similarity_search_by_vector(embedding=vectors, k=top_k, filter=filters)
else:
results = self.client.similarity_search_by_vector(embedding=vectors, k=limit)
results = self.client.similarity_search_by_vector(embedding=vectors, k=top_k)
final_results = self._parse_output(results)
return final_results
@@ -152,7 +152,7 @@ class Langchain(VectorStoreBase):
"""
return {"name": self.collection_name}
def list(self, filters=None, limit=None):
def list(self, filters=None, top_k=None):
"""
List all vectors in a collection.
"""
@@ -164,7 +164,7 @@ class Langchain(VectorStoreBase):
# Handle all filters, not just user_id
where_clause = filters
result = self.client._collection.get(where=where_clause, limit=limit)
result = self.client._collection.get(where=where_clause, limit=top_k)
# Convert the result to the expected format
if result and isinstance(result, dict):
+6 -6
View File
@@ -139,14 +139,14 @@ class MilvusDB(VectorStoreBase):
return memory
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> list:
def search(self, query: str, vectors: list, top_k: int = 5, filters: dict = None) -> list:
"""
Search for similar vectors.
Args:
query (str): Query.
vectors (List[float]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -156,7 +156,7 @@ class MilvusDB(VectorStoreBase):
hits = self.client.search(
collection_name=self.collection_name,
data=[vectors],
limit=limit,
limit=top_k,
filter=query_filter,
output_fields=["*"],
)
@@ -235,19 +235,19 @@ class MilvusDB(VectorStoreBase):
"""
return self.client.get_collection_stats(collection_name=self.collection_name)
def list(self, filters: dict = None, limit: int = 100) -> list:
def list(self, filters: dict = None, top_k: int = 100) -> list:
"""
List all vectors in a collection.
Args:
filters (Dict, optional): Filters to apply to the list.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors.
"""
query_filter = self._create_filter(filters) if filters else None
result = self.client.query(collection_name=self.collection_name, filter=query_filter, limit=limit)
result = self.client.query(collection_name=self.collection_name, filter=query_filter, limit=top_k)
memories = []
for data in result:
obj = OutputData(id=data.get("id"), score=None, payload=data.get("metadata"))
+7 -7
View File
@@ -113,14 +113,14 @@ class MongoDB(VectorStoreBase):
except PyMongoError as e:
logger.error(f"Error inserting data: {e}")
def search(self, query: str, vectors: List[float], limit=5, filters: Optional[Dict] = None) -> List[OutputData]:
def search(self, query: str, vectors: List[float], top_k=5, filters: Optional[Dict] = None) -> List[OutputData]:
"""
Search for similar vectors using the vector search index.
Args:
query (str): Query string
vectors (List[float]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search.
Returns:
@@ -139,8 +139,8 @@ class MongoDB(VectorStoreBase):
{
"$vectorSearch": {
"index": self.index_name,
"limit": limit,
"numCandidates": min(limit * 20, 10000),
"limit": top_k,
"numCandidates": min(top_k * 20, 10000),
"queryVector": vectors,
"path": "embedding",
}
@@ -271,13 +271,13 @@ class MongoDB(VectorStoreBase):
logger.error(f"Error getting collection info: {e}")
return {}
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[OutputData]:
def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> List[OutputData]:
"""
List vectors in the collection.
Args:
filters (Dict, optional): Filters to apply to the list.
limit (int, optional): Number of vectors to return.
top_k (int, optional): Number of vectors to return.
Returns:
List[OutputData]: List of vectors.
@@ -292,7 +292,7 @@ class MongoDB(VectorStoreBase):
if filter_conditions:
query = {"$and": filter_conditions}
cursor = self.collection.find(query).limit(limit)
cursor = self.collection.find(query).limit(top_k)
results = [OutputData(id=str(doc["_id"]), score=None, payload=doc.get("payload")) for doc in cursor]
logger.info(f"Retrieved {len(results)} documents from collection '{self.collection_name}'.")
return results
+6 -6
View File
@@ -133,7 +133,7 @@ class NeptuneAnalyticsVector(VectorStoreBase):
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors using embedding similarity.
@@ -144,7 +144,7 @@ class NeptuneAnalyticsVector(VectorStoreBase):
Args:
query (str): Search query text (unused in vector search).
vectors (List[float]): Query embedding vector.
limit (int, optional): Maximum number of results to return. Defaults to 5.
top_k (int, optional): Maximum number of results to return. Defaults to 5.
filters (Optional[Dict]): Optional filters to apply to search results.
Returns:
@@ -159,7 +159,7 @@ class NeptuneAnalyticsVector(VectorStoreBase):
query_string = f"""
CALL neptune.algo.vectors.topKByEmbeddingWithFiltering({{
topK: {limit},
topK: {top_k},
embedding: {vectors}
{filter_clause}
}}
@@ -309,7 +309,7 @@ class NeptuneAnalyticsVector(VectorStoreBase):
pass
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[OutputData]:
def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> List[OutputData]:
"""
List all vectors in the collection with optional filtering.
@@ -317,7 +317,7 @@ class NeptuneAnalyticsVector(VectorStoreBase):
Args:
filters (Optional[Dict]): Optional filters to apply based on metadata.
limit (int, optional): Maximum number of vectors to return. Defaults to 100.
top_k (int, optional): Maximum number of vectors to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors with their metadata.
@@ -325,7 +325,7 @@ class NeptuneAnalyticsVector(VectorStoreBase):
where_clause = self._get_where_clause(filters) if filters else ""
para = {
"limit": limit,
"limit": top_k,
}
query_string = f"""
MATCH (n :{self.collection_name})
+7 -7
View File
@@ -158,7 +158,7 @@ class OpenSearchDB(VectorStoreBase):
return results
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""Search for similar vectors using OpenSearch k-NN search with optional filters."""
@@ -167,13 +167,13 @@ class OpenSearchDB(VectorStoreBase):
"knn": {
"vector_field": {
"vector": vectors,
"k": limit * 2,
"k": top_k * 2,
}
}
}
# Start building the full query
query_body = {"size": limit * 2, "query": None}
query_body = {"size": top_k * 2, "query": None}
# Prepare filter conditions if applicable
filter_clauses = []
@@ -196,7 +196,7 @@ class OpenSearchDB(VectorStoreBase):
hits = response["hits"]["hits"]
results = [
OutputData(id=hit["_source"].get("id"), score=hit["_score"], payload=hit["_source"].get("payload", {}))
for hit in hits[:limit] # Ensure we don't exceed limit
for hit in hits[:top_k] # Ensure we don't exceed top_k
]
return results
except Exception as e:
@@ -284,7 +284,7 @@ class OpenSearchDB(VectorStoreBase):
"""Get information about a collection (index)."""
return self.client.indices.get(index=name)
def list(self, filters: Optional[Dict] = None, limit: Optional[int] = None) -> List[OutputData]:
def list(self, filters: Optional[Dict] = None, top_k: Optional[int] = None) -> List[OutputData]:
try:
"""List all memories with optional filters."""
query: Dict = {"query": {"match_all": {}}}
@@ -299,8 +299,8 @@ class OpenSearchDB(VectorStoreBase):
if filter_clauses:
query["query"] = {"bool": {"filter": filter_clauses}}
if limit:
query["size"] = limit
if top_k:
query["size"] = top_k
response = self.client.search(index=self.collection_name, body=query)
hits = response["hits"]["hits"]
+6 -6
View File
@@ -203,7 +203,7 @@ class PGVector(VectorStoreBase):
self,
query: str,
vectors: list[float],
limit: Optional[int] = 5,
top_k: Optional[int] = 5,
filters: Optional[dict] = None,
) -> List[OutputData]:
"""
@@ -212,7 +212,7 @@ class PGVector(VectorStoreBase):
Args:
query (str): Query.
vectors (List[float]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -237,7 +237,7 @@ class PGVector(VectorStoreBase):
ORDER BY distance
LIMIT %s
""",
(vectors, *filter_params, limit),
(vectors, *filter_params, top_k),
)
results = cur.fetchall()
@@ -350,14 +350,14 @@ class PGVector(VectorStoreBase):
def list(
self,
filters: Optional[dict] = None,
limit: Optional[int] = 100
top_k: Optional[int] = 100
) -> List[OutputData]:
"""
List all vectors in a collection.
Args:
filters (Dict, optional): Filters to apply to the list.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors.
@@ -380,7 +380,7 @@ class PGVector(VectorStoreBase):
"""
with self._get_cursor() as cur:
cur.execute(query, (*filter_params, limit))
cur.execute(query, (*filter_params, top_k))
results = cur.fetchall()
return [[OutputData(id=str(r[0]), score=None, payload=r[2]) for r in results]]
+6 -6
View File
@@ -204,7 +204,7 @@ class PineconeDB(VectorStoreBase):
return pinecone_filter
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
@@ -212,7 +212,7 @@ class PineconeDB(VectorStoreBase):
Args:
query (str): Query.
vectors (list): List of vectors to search.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -222,7 +222,7 @@ class PineconeDB(VectorStoreBase):
query_params = {
"vector": vectors,
"top_k": limit,
"top_k": top_k,
"include_metadata": True,
"include_values": False,
}
@@ -320,13 +320,13 @@ class PineconeDB(VectorStoreBase):
"""
return self.client.describe_index(self.collection_name)
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[OutputData]:
def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> List[OutputData]:
"""
List vectors in an index with optional filtering.
Args:
filters (dict, optional): Filters to apply to the list. Defaults to None.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
dict: List of vectors with their metadata.
@@ -340,7 +340,7 @@ class PineconeDB(VectorStoreBase):
query_params = {
"vector": zero_vector,
"top_k": limit,
"top_k": top_k,
"include_metadata": True,
"include_values": True,
}
+6 -6
View File
@@ -308,14 +308,14 @@ class Qdrant(VectorStoreBase):
must_not=must_not or None,
)
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> list:
def search(self, query: str, vectors: list, top_k: int = 5, filters: dict = None) -> list:
"""
Search for similar vectors.
Args:
query (str): Query.
vectors (list): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -326,7 +326,7 @@ class Qdrant(VectorStoreBase):
collection_name=self.collection_name,
query=vectors,
query_filter=query_filter,
limit=limit,
limit=top_k,
)
return hits.points
@@ -404,13 +404,13 @@ class Qdrant(VectorStoreBase):
"""
return self.client.get_collection(collection_name=self.collection_name)
def list(self, filters: dict = None, limit: int = 100) -> list:
def list(self, filters: dict = None, top_k: int = 100) -> list:
"""
List all vectors in a collection.
Args:
filters (dict, optional): Filters to apply to the list. Defaults to None.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
list: List of vectors.
@@ -419,7 +419,7 @@ class Qdrant(VectorStoreBase):
result = self.client.scroll(
collection_name=self.collection_name,
scroll_filter=query_filter,
limit=limit,
limit=top_k,
with_payload=True,
with_vectors=False,
)
+5 -5
View File
@@ -141,7 +141,7 @@ class RedisDB(VectorStoreBase):
data.append(entry)
self.index.load(data, id_field="memory_id")
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None):
def search(self, query: str, vectors: list, top_k: int = 5, filters: dict = None):
conditions = [Tag(key) == value for key, value in filters.items() if value is not None]
filter = reduce(lambda x, y: x & y, conditions)
@@ -150,7 +150,7 @@ class RedisDB(VectorStoreBase):
vector_field_name="embedding",
return_fields=["memory_id", "hash", "agent_id", "run_id", "user_id", "memory", "metadata", "created_at"],
filter_expression=filter,
num_results=limit,
num_results=top_k,
)
results = self.index.query(v)
@@ -254,15 +254,15 @@ class RedisDB(VectorStoreBase):
# Recreate the index with the same parameters
self.create_col(collection_name, self.embedding_model_dims)
def list(self, filters: dict = None, limit: int = None) -> list:
def list(self, filters: dict = None, top_k: int = None) -> list:
"""
List all recent created memories from the vector store.
"""
conditions = [Tag(key) == value for key, value in filters.items() if value is not None]
filter = reduce(lambda x, y: x & y, conditions)
query = Query(str(filter)).sort_by("created_at", asc=False)
if limit is not None:
query = Query(str(filter)).sort_by("created_at", asc=False).paging(0, limit)
if top_k is not None:
query = Query(str(filter)).sort_by("created_at", asc=False).paging(0, top_k)
results = self.index.search(query)
return [
+5 -5
View File
@@ -99,12 +99,12 @@ class S3Vectors(VectorStoreBase):
vectors=vectors_to_put,
)
def search(self, query, vectors, limit=5, filters=None):
def search(self, query, vectors, top_k=5, filters=None):
params = {
"vectorBucketName": self.vector_bucket_name,
"indexName": self.collection_name,
"queryVector": {"float32": vectors},
"topK": limit,
"topK": top_k,
"returnMetadata": True,
"returnDistance": True,
}
@@ -149,7 +149,7 @@ class S3Vectors(VectorStoreBase):
response = self.client.get_index(vectorBucketName=self.vector_bucket_name, indexName=self.collection_name)
return response.get("index", {})
def list(self, filters=None, limit=None):
def list(self, filters=None, top_k=None):
# Note: list_vectors does not support metadata filtering.
if filters:
logger.warning("S3 Vectors `list` does not support metadata filtering. Ignoring filters.")
@@ -160,8 +160,8 @@ class S3Vectors(VectorStoreBase):
"returnData": False,
"returnMetadata": True,
}
if limit:
params["maxResults"] = limit
if top_k:
params["maxResults"] = top_k
paginator = self.client.get_paginator("list_vectors")
pages = paginator.paginate(**params)
+6 -6
View File
@@ -116,7 +116,7 @@ class Supabase(VectorStoreBase):
self.collection.upsert(records)
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
@@ -124,7 +124,7 @@ class Supabase(VectorStoreBase):
Args:
query (str): Query.
vectors (List[float]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -132,7 +132,7 @@ class Supabase(VectorStoreBase):
"""
filters = self._preprocess_filters(filters)
results = self.collection.query(
data=vectors, limit=limit, filters=filters, include_metadata=True, include_value=True
data=vectors, limit=top_k, filters=filters, include_metadata=True, include_value=True
)
return [OutputData(id=str(result[0]), score=float(result[1]), payload=result[2]) for result in results]
@@ -209,13 +209,13 @@ class Supabase(VectorStoreBase):
"index": {"method": info.index_method, "metric": info.distance_metric},
}
def list(self, filters: Optional[dict] = None, limit: int = 100) -> List[OutputData]:
def list(self, filters: Optional[dict] = None, top_k: int = 100) -> List[OutputData]:
"""
List vectors in the collection.
Args:
filters (Dict, optional): Filters to apply
limit (int, optional): Maximum number of results to return. Defaults to 100.
top_k (int, optional): Maximum number of results to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors
@@ -223,7 +223,7 @@ class Supabase(VectorStoreBase):
filters = self._preprocess_filters(filters)
query = [0] * self.embedding_model_dims
ids = self.collection.query(
data=query, limit=limit, filters=filters, include_metadata=True, include_value=False
data=query, limit=top_k, filters=filters, include_metadata=True, include_value=False
)
ids = [id[0] for id in ids]
records = self.collection.fetch(ids=ids)
+6 -6
View File
@@ -157,7 +157,7 @@ class TurbopufferDB(VectorStoreBase):
return ("And", tuple(conditions))
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
@@ -165,7 +165,7 @@ class TurbopufferDB(VectorStoreBase):
Args:
query (str): Query text (unused in vector search, kept for interface consistency).
vectors (list): Query vector to search with.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -173,7 +173,7 @@ class TurbopufferDB(VectorStoreBase):
"""
query_params = {
"rank_by": ("vector", "ANN", vectors),
"top_k": limit,
"top_k": top_k,
"include_attributes": True,
}
@@ -290,20 +290,20 @@ class TurbopufferDB(VectorStoreBase):
except Exception:
return {"name": self.collection_name}
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> list:
def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> list:
"""
List vectors in the namespace with optional filtering.
Args:
filters (dict, optional): Filters to apply. Defaults to None.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
list: Wrapped list of OutputData objects ([[results]]).
"""
query_params = {
"rank_by": ("vector", "ANN", [0.0] * self.embedding_model_dims),
"top_k": limit,
"top_k": top_k,
"include_attributes": True,
}
+7 -7
View File
@@ -98,7 +98,7 @@ class UpstashVector(VectorStoreBase):
self,
query: str,
vectors: List[list],
limit: int = 5,
top_k: int = 5,
filters: Optional[Dict] = None,
) -> List[OutputData]:
"""
@@ -106,7 +106,7 @@ class UpstashVector(VectorStoreBase):
Args:
query (list): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search.
Returns:
@@ -120,7 +120,7 @@ class UpstashVector(VectorStoreBase):
if self.enable_embeddings:
response = self.client.query(
data=query,
top_k=limit,
top_k=top_k,
filter=filters_str or "",
include_metadata=True,
namespace=self.collection_name,
@@ -129,7 +129,7 @@ class UpstashVector(VectorStoreBase):
queries = [
{
"vector": v,
"top_k": limit,
"top_k": top_k,
"filter": filters_str or "",
"include_metadata": True,
"namespace": self.collection_name,
@@ -205,12 +205,12 @@ class UpstashVector(VectorStoreBase):
return None
return OutputData(id=vector.id, score=None, payload=vector.metadata)
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[List[OutputData]]:
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.
limit (int, optional): Number of results to return. Defaults to 100.
top_k (int, optional): Number of results to return. Defaults to 100.
Returns:
List[OutputData]: Search results.
"""
@@ -233,7 +233,7 @@ class UpstashVector(VectorStoreBase):
)
with query:
while True:
if len(results) >= limit:
if len(results) >= top_k:
break
res = query.fetch_next(100)
if not res:
+8 -8
View File
@@ -410,14 +410,14 @@ class ValkeyDB(VectorStoreBase):
return memory_results
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None, ef_runtime: int = None):
def search(self, query: str, vectors: list, top_k: int = 5, filters: dict = None, ef_runtime: int = None):
"""
Search for similar vectors in the index.
Args:
query (str): The search query.
vectors (list): The vector to search for.
limit (int, optional): Maximum number of results to return. Defaults to 5.
top_k (int, optional): Maximum number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None.
ef_runtime (int, optional): HNSW ef_runtime parameter for this query. Only used with HNSW index. Defaults to None.
@@ -429,10 +429,10 @@ class ValkeyDB(VectorStoreBase):
# Build the KNN part with optional EF_RUNTIME for HNSW
if self.index_type == "hnsw" and ef_runtime is not None:
knn_part = f"[KNN {limit} @embedding $vec_param EF_RUNTIME {ef_runtime} AS vector_score]"
knn_part = f"[KNN {top_k} @embedding $vec_param EF_RUNTIME {ef_runtime} AS vector_score]"
else:
# For FLAT indexes or when ef_runtime is None, use basic KNN
knn_part = f"[KNN {limit} @embedding $vec_param AS vector_score]"
knn_part = f"[KNN {top_k} @embedding $vec_param AS vector_score]"
# Build the complete query
q = self._build_search_query(knn_part, filters)
@@ -768,7 +768,7 @@ class ValkeyDB(VectorStoreBase):
return q
def list(self, filters: dict = None, limit: int = None) -> list:
def list(self, filters: dict = None, top_k: int = None) -> list:
"""
List all recent created memories from the vector store.
@@ -778,7 +778,7 @@ class ValkeyDB(VectorStoreBase):
Values are used as-is without validation - wildcards, special characters,
lists, etc. are passed through literally to Valkey search.
Multiple filters are combined with AND logic.
limit (int, optional): Maximum number of results to return. Defaults to 1000
top_k (int, optional): Maximum number of results to return. Defaults to 1000
if not specified.
Returns:
@@ -789,10 +789,10 @@ class ValkeyDB(VectorStoreBase):
# Since Valkey search requires vector format, use a dummy vector search
# that returns all documents by using a zero vector and large K
dummy_vector = [0.0] * self.embedding_model_dims
search_limit = limit if limit is not None else 1000 # Large default
search_limit = top_k if top_k is not None else 1000 # Large default
# Use the existing search method which handles filters properly
search_results = self.search("", dummy_vector, limit=search_limit, filters=filters)
search_results = self.search("", dummy_vector, top_k=search_limit, filters=filters)
# Convert search results to list format (match Redis format)
class MemoryResult:
+14 -12
View File
@@ -5,7 +5,9 @@ from typing import Any, Dict, List, Optional, Tuple
import google.api_core.exceptions
from google.cloud import aiplatform, aiplatform_v1
from google.cloud.aiplatform.matching_engine.matching_engine_index_endpoint import Namespace
from google.cloud.aiplatform.matching_engine.matching_engine_index_endpoint import (
Namespace,
)
from google.oauth2 import service_account
from pydantic import BaseModel
@@ -206,20 +208,20 @@ class GoogleMatchingEngine(VectorStoreBase):
raise
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
Args:
query (str): Query.
vectors (List[float]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Optional[Dict], optional): Filters to apply to the search. Defaults to None.
Returns:
List[OutputData]: Search results (unwrapped)
"""
logger.debug("Starting search")
logger.debug("Limit: %d, Filters: %s", limit, filters)
logger.debug("Limit: %d, Filters: %s", top_k, filters)
try:
filter_namespaces = []
@@ -241,7 +243,7 @@ class GoogleMatchingEngine(VectorStoreBase):
response = self.index_endpoint.find_neighbors(
deployed_index_id=self.deployment_index_id,
queries=[vectors],
num_neighbors=limit,
num_neighbors=top_k,
filter=filter_namespaces if filter_namespaces else None,
return_full_datapoint=True,
)
@@ -453,12 +455,12 @@ class GoogleMatchingEngine(VectorStoreBase):
"region": self.region,
}
def list(self, filters: Optional[Dict] = None, limit: Optional[int] = None) -> List[List[OutputData]]:
def list(self, filters: Optional[Dict] = None, top_k: Optional[int] = None) -> List[List[OutputData]]:
"""List vectors matching the given filters.
Args:
filters: Optional filters to apply
limit: Optional maximum number of results to return
top_k: Optional maximum number of results to return
Returns:
List[List[OutputData]]: List of matching vectors wrapped in an extra array
@@ -466,17 +468,17 @@ class GoogleMatchingEngine(VectorStoreBase):
"""
logger.debug("Starting list operation")
logger.debug("Filters: %s", filters)
logger.debug("Limit: %s", limit)
logger.debug("Limit: %s", top_k)
try:
# Use a zero vector for the search
dimension = 768 # This should be configurable based on the model
zero_vector = [0.0] * dimension
# Use a large limit if none specified
search_limit = limit if limit is not None else 10000
# Use a large top_k if none specified
search_limit = top_k if top_k is not None else 10000
results = self.search(query=zero_vector, limit=search_limit, filters=filters)
results = self.search(query=zero_vector, top_k=search_limit, filters=filters)
logger.debug("Found %d results", len(results))
return [results] # Wrap in extra array to match interface
@@ -609,7 +611,7 @@ class GoogleMatchingEngine(VectorStoreBase):
logger.debug("Filter: %s", filter)
embedding = self.embedder.embed_query(query)
results = self.search(query=embedding, limit=k, filters=filter)
results = self.search(query=embedding, top_k=k, filters=filter)
docs_and_scores = [
(Document(page_content=result.payload.get("text", ""), metadata=result.payload), result.score)
+4 -4
View File
@@ -179,7 +179,7 @@ class Weaviate(VectorStoreBase):
batch.add_object(collection=self.collection_name, properties=data_object, uuid=object_id, vector=vector)
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
@@ -194,7 +194,7 @@ class Weaviate(VectorStoreBase):
response = collection.query.hybrid(
query="",
vector=vectors,
limit=limit,
limit=top_k,
filters=combined_filter,
return_properties=["hash", "created_at", "updated_at", "user_id", "agent_id", "run_id", "data", "category"],
return_metadata=MetadataQuery(score=True),
@@ -313,7 +313,7 @@ class Weaviate(VectorStoreBase):
return schema
return None
def list(self, filters=None, limit=100) -> List[OutputData]:
def list(self, filters=None, top_k=100) -> List[OutputData]:
"""
List all vectors in a collection.
"""
@@ -325,7 +325,7 @@ class Weaviate(VectorStoreBase):
filter_conditions.append(Filter.by_property(key).equal(value))
combined_filter = Filter.all_of(filter_conditions) if filter_conditions else None
response = collection.query.fetch_objects(
limit=limit,
limit=top_k,
filters=combined_filter,
return_properties=["hash", "created_at", "updated_at", "user_id", "agent_id", "run_id", "data", "category"],
)