refactor: add entity ID and search param validation, rename textLemmatized field, update tests (#4843)

This commit is contained in:
Kartik
2026-04-15 20:57:09 +05:30
committed by GitHub
parent 9692726db4
commit e6d6276bb9
10 changed files with 667 additions and 41 deletions
+150 -17
View File
@@ -110,6 +110,64 @@ def _reject_top_level_entity_params(kwargs: Dict[str, Any], method_name: str) ->
)
def _validate_and_trim_entity_id(value: Optional[str], name: str) -> Optional[str]:
"""
Validates and normalizes an entity ID.
- Trims leading/trailing whitespace
- Rejects empty or whitespace-only strings
- Rejects strings containing internal whitespace
Args:
value: The entity ID value to validate
name: The parameter name (for error messages)
Returns:
The trimmed entity ID, or None if input is None
Raises:
ValueError: If entity ID is invalid
"""
if value is None:
return None
trimmed = value.strip()
if trimmed == "":
raise ValueError(
f"Invalid {name}: cannot be empty or whitespace-only. Provide a valid identifier."
)
if any(c.isspace() for c in trimmed):
raise ValueError(
f"Invalid {name}: cannot contain whitespace. Provide a valid identifier without spaces."
)
return trimmed
def _validate_search_params(threshold: Optional[float] = None, top_k: Optional[int] = None) -> None:
"""
Validates search parameters.
Args:
threshold: Similarity threshold (must be between 0 and 1)
top_k: Number of results to return (must be non-negative integer)
Raises:
ValueError: If threshold or top_k are invalid
"""
if threshold is not None:
if not isinstance(threshold, (int, float)):
raise ValueError("threshold must be a valid number")
if threshold < 0 or threshold > 1:
raise ValueError(
f"Invalid threshold: {threshold}. Must be between 0 and 1 (inclusive)."
)
if top_k is not None:
if not isinstance(top_k, int) or isinstance(top_k, bool):
raise ValueError("top_k must be a valid integer")
if top_k < 0:
raise ValueError(
f"Invalid top_k: {top_k}. Must be a non-negative integer."
)
def _is_sensitive_field(field_name: str) -> bool:
"""Check if a field should be redacted for telemetry safety.
@@ -217,9 +275,14 @@ def _build_filters_and_metadata(
base_metadata_template = deepcopy(input_metadata) if input_metadata else {}
effective_query_filters = deepcopy(input_filters) if input_filters else {}
# ---------- add all provided session ids ----------
# ---------- validate and add all provided session ids ----------
session_ids_provided = []
# Validate and trim entity IDs
user_id = _validate_and_trim_entity_id(user_id, "user_id")
agent_id = _validate_and_trim_entity_id(agent_id, "agent_id")
run_id = _validate_and_trim_entity_id(run_id, "run_id")
if user_id:
base_metadata_template["user_id"] = user_id
effective_query_filters["user_id"] = user_id
@@ -876,7 +939,7 @@ class Memory(MemoryBase):
self,
*,
filters: Optional[Dict[str, Any]] = None,
top_k: int = 100,
top_k: int = 20,
**kwargs,
):
"""
@@ -886,20 +949,38 @@ class Memory(MemoryBase):
filters (dict): Filter dict containing entity IDs and optional metadata filters.
Must contain at least one of: user_id, agent_id, run_id.
Example: filters={"user_id": "u1", "agent_id": "a1"}
top_k (int, optional): The maximum number of memories to return. Defaults to 100.
top_k (int, optional): The maximum number of memories to return. Defaults to 20.
Returns:
dict: A dictionary containing a list of memories under the "results" key.
Example for v1.1+: `{"results": [{"id": "...", "memory": "...", ...}]}`
Raises:
ValueError: If filters doesn't contain at least one of user_id, agent_id, run_id.
ValueError: If filters doesn't contain at least one of user_id, agent_id, run_id,
or if top_k is invalid.
"""
# Reject top-level entity params - must use filters instead
_reject_top_level_entity_params(kwargs, "get_all")
# Validate top_k
_validate_search_params(top_k=top_k)
# Validate and trim entity IDs in filters
effective_filters = dict(filters) if filters else {}
if "user_id" in effective_filters:
effective_filters["user_id"] = _validate_and_trim_entity_id(
effective_filters["user_id"], "user_id"
)
if "agent_id" in effective_filters:
effective_filters["agent_id"] = _validate_and_trim_entity_id(
effective_filters["agent_id"], "agent_id"
)
if "run_id" in effective_filters:
effective_filters["run_id"] = _validate_and_trim_entity_id(
effective_filters["run_id"], "run_id"
)
# Validate filters contains at least one entity ID
effective_filters = filters or {}
if not any(key in effective_filters for key in ("user_id", "agent_id", "run_id")):
raise ValueError(
"filters must contain at least one of: user_id, agent_id, run_id. "
@@ -968,7 +1049,7 @@ class Memory(MemoryBase):
self,
query: str,
*,
top_k: int = 100,
top_k: int = 20,
filters: Optional[Dict[str, Any]] = None,
threshold: float = 0.1,
rerank: bool = False,
@@ -979,7 +1060,7 @@ class Memory(MemoryBase):
Args:
query (str): Query to search for.
top_k (int, optional): Maximum number of results to return. Defaults to 100.
top_k (int, optional): Maximum number of results to return. Defaults to 20.
filters (dict): Filter dict containing entity IDs and optional metadata filters.
Must contain at least one of: user_id, agent_id, run_id.
Example: filters={"user_id": "u1", "agent_id": "a1"}
@@ -1008,13 +1089,29 @@ class Memory(MemoryBase):
Example for v1.1+: `{"results": [{"id": "...", "memory": "...", "score": 0.8, ...}]}`
Raises:
ValueError: If filters doesn't contain at least one of user_id, agent_id, run_id.
ValueError: If filters doesn't contain at least one of user_id, agent_id, run_id,
or if threshold/top_k values are invalid.
"""
# Reject top-level entity params - must use filters instead
_reject_top_level_entity_params(kwargs, "search")
# Validate filters contains at least one entity ID
# Validate search parameters (before applying defaults)
_validate_search_params(threshold=threshold, top_k=top_k)
# Validate and trim entity IDs in filters
effective_filters = filters.copy() if filters else {}
if "user_id" in effective_filters:
effective_filters["user_id"] = _validate_and_trim_entity_id(
effective_filters["user_id"], "user_id"
)
if "agent_id" in effective_filters:
effective_filters["agent_id"] = _validate_and_trim_entity_id(
effective_filters["agent_id"], "agent_id"
)
if "run_id" in effective_filters:
effective_filters["run_id"] = _validate_and_trim_entity_id(
effective_filters["run_id"], "run_id"
)
if not any(key in effective_filters for key in ("user_id", "agent_id", "run_id")):
raise ValueError(
"filters must contain at least one of: user_id, agent_id, run_id. "
@@ -2126,7 +2223,7 @@ class AsyncMemory(MemoryBase):
self,
*,
filters: Optional[Dict[str, Any]] = None,
top_k: int = 100,
top_k: int = 20,
**kwargs,
):
"""
@@ -2136,20 +2233,38 @@ class AsyncMemory(MemoryBase):
filters (dict): Filter dict containing entity IDs and optional metadata filters.
Must contain at least one of: user_id, agent_id, run_id.
Example: filters={"user_id": "u1", "agent_id": "a1"}
top_k (int, optional): The maximum number of memories to return. Defaults to 100.
top_k (int, optional): The maximum number of memories to return. Defaults to 20.
Returns:
dict: A dictionary containing a list of memories under the "results" key.
Example for v1.1+: `{"results": [{"id": "...", "memory": "...", ...}]}`
Raises:
ValueError: If filters doesn't contain at least one of user_id, agent_id, run_id.
ValueError: If filters doesn't contain at least one of user_id, agent_id, run_id,
or if top_k is invalid.
"""
# Reject top-level entity params - must use filters instead
_reject_top_level_entity_params(kwargs, "get_all")
# Validate top_k
_validate_search_params(top_k=top_k)
# Validate and trim entity IDs in filters
effective_filters = dict(filters) if filters else {}
if "user_id" in effective_filters:
effective_filters["user_id"] = _validate_and_trim_entity_id(
effective_filters["user_id"], "user_id"
)
if "agent_id" in effective_filters:
effective_filters["agent_id"] = _validate_and_trim_entity_id(
effective_filters["agent_id"], "agent_id"
)
if "run_id" in effective_filters:
effective_filters["run_id"] = _validate_and_trim_entity_id(
effective_filters["run_id"], "run_id"
)
# Validate filters contains at least one entity ID
effective_filters = filters or {}
if not any(key in effective_filters for key in ("user_id", "agent_id", "run_id")):
raise ValueError(
"filters must contain at least one of: user_id, agent_id, run_id. "
@@ -2218,7 +2333,7 @@ class AsyncMemory(MemoryBase):
self,
query: str,
*,
top_k: int = 100,
top_k: int = 20,
filters: Optional[Dict[str, Any]] = None,
threshold: float = 0.1,
rerank: bool = False,
@@ -2229,7 +2344,7 @@ class AsyncMemory(MemoryBase):
Args:
query (str): Query to search for.
top_k (int, optional): Maximum number of results to return. Defaults to 100.
top_k (int, optional): Maximum number of results to return. Defaults to 20.
filters (dict): Filter dict containing entity IDs and optional metadata filters.
Must contain at least one of: user_id, agent_id, run_id.
Example: filters={"user_id": "u1", "agent_id": "a1"}
@@ -2258,13 +2373,31 @@ class AsyncMemory(MemoryBase):
Example for v1.1+: `{"results": [{"id": "...", "memory": "...", "score": 0.8, ...}]}`
Raises:
ValueError: If filters doesn't contain at least one of user_id, agent_id, run_id.
ValueError: If filters doesn't contain at least one of user_id, agent_id, run_id,
or if threshold/top_k values are invalid.
"""
# Reject top-level entity params - must use filters instead
_reject_top_level_entity_params(kwargs, "search")
# Validate filters contains at least one entity ID
# Validate search parameters (before applying defaults)
_validate_search_params(threshold=threshold, top_k=top_k)
# Validate and trim entity IDs in filters
effective_filters = filters.copy() if filters else {}
if "user_id" in effective_filters:
effective_filters["user_id"] = _validate_and_trim_entity_id(
effective_filters["user_id"], "user_id"
)
if "agent_id" in effective_filters:
effective_filters["agent_id"] = _validate_and_trim_entity_id(
effective_filters["agent_id"], "agent_id"
)
if "run_id" in effective_filters:
effective_filters["run_id"] = _validate_and_trim_entity_id(
effective_filters["run_id"], "run_id"
)
# Validate filters contains at least one entity ID
if not any(key in effective_filters for key in ("user_id", "agent_id", "run_id")):
raise ValueError(
"filters must contain at least one of: user_id, agent_id, run_id. "