fix(sdk): removing deprecating param from our sdk and docs changes with it (#4740)
This commit is contained in:
@@ -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"]))
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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
@@ -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]
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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})
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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]]
|
||||
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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 [
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"],
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user