Apply literal contains filters consistently to numeric and textual fields

This commit is contained in:
Andras Schmelczer 2026-10-27 19:57:42 +00:00
commit c5bf7d4a7d
5 changed files with 96 additions and 10 deletions

View file

@ -1,3 +1,4 @@
import re
from datetime import datetime from datetime import datetime
from typing import Any, Dict, List, Optional, Sequence, Tuple from typing import Any, Dict, List, Optional, Sequence, Tuple
@ -102,6 +103,19 @@ class MongoDbDriver(TracingDatabaseDriver):
return operator_mapping["="] return operator_mapping["="]
return operator_mapping[filter.operator] return operator_mapping[filter.operator]
def _get_filter(self, filter: Filter) -> Dict[str, Any]:
if filter.operator == "contains":
text_filter = {
filter.property: {
"$regex": re.escape(filter.value_as_text),
"$options": "i",
}
}
if isinstance(filter.value, str):
return text_filter
return {"$or": [{filter.property: {"$eq": filter.value}}, text_filter]}
return {filter.property: {self._get_operator(filter): filter.value}}
def query( def query(
self, self,
*, *,
@ -121,9 +135,7 @@ class MongoDbDriver(TracingDatabaseDriver):
and_query: List[Dict[str, Any]] = [] and_query: List[Dict[str, Any]] = []
and_query.extend({"tags": tag} for tag in conjunctive_tags) and_query.extend({"tags": tag} for tag in conjunctive_tags)
and_query.extend( and_query.extend(self._get_filter(f) for f in conjunctive_filters)
{f.property: {self._get_operator(f): f.value}} for f in conjunctive_filters
)
if not and_query: if not and_query:
and_query.append({}) and_query.append({})

View file

@ -84,10 +84,14 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
getattr(df[f.property], operator_mapping[f.operator])(f.value) getattr(df[f.property], operator_mapping[f.operator])(f.value)
] ]
elif operator == "contains": elif operator == "contains":
pattern = f.value_as_text.casefold()
df = df.loc[ df = df.loc[
df[f.property].str.contains( df[f.property].map(
str(int(f.value)) if isinstance(f.value, float) else f.value, lambda value: (
case=False, pattern in value.casefold()
if isinstance(value, str)
else not isinstance(f.value, str) and value == f.value
)
) )
] ]

View file

@ -1,3 +1,4 @@
import builtins
from typing import Union from typing import Union
from pydantic import BaseModel from pydantic import BaseModel
@ -9,3 +10,9 @@ class Filter(BaseModel):
property: str property: str
operator: Operator operator: Operator
value: Union[float, str] value: Union[float, str]
@builtins.property
def value_as_text(self) -> str:
if isinstance(self.value, float) and self.value.is_integer():
return str(int(self.value))
return str(self.value)

View file

@ -5,10 +5,10 @@ from pathlib import Path
from typing import Any from typing import Any
from unittest.mock import patch from unittest.mock import patch
from tinydb.table import Table
from great_ai import ParallelTinyDbDriver, Trace from great_ai import ParallelTinyDbDriver, Trace
from great_ai.context import get_context from great_ai.context import get_context
from great_ai.views import Filter
from tinydb.table import Table
def write_traces(path: str, index: int, barrier: Any) -> None: def write_traces(path: str, index: int, barrier: Any) -> None:
@ -74,3 +74,45 @@ def test_tinydb_date_queries_normalize_timezones(isolated_context: None) -> None
noon = datetime(2026, 1, 1, 12, tzinfo=timezone.utc) noon = datetime(2026, 1, 1, 12, tzinfo=timezone.utc)
assert database.query(since=noon, until=noon)[1] == 2 assert database.query(since=noon, until=noon)[1] == 2
assert database.query(since=noon.replace(hour=13))[1] == 0 assert database.query(since=noon.replace(hour=13))[1] == 0
def test_contains_handles_numbers_and_literal_punctuation(
isolated_context: None,
) -> None:
database = get_context().tracing_database
database.save(
Trace(
trace_id="contains",
created="2026-01-01",
original_execution_time_ms=2.5,
models=[],
tags=[],
logged_values={"text": "A[2.5"},
)
)
for property, value in (
("original_execution_time_ms", 2.5),
("text", "["),
("text", 2.5),
("text", "a["),
):
assert (
database.query(
conjunctive_filters=[
Filter(property=property, operator="contains", value=value)
]
)[1]
== 1
)
assert (
database.query(
conjunctive_filters=[
Filter(
property="original_execution_time_ms",
operator="contains",
value="[",
)
]
)[1]
== 0
)

View file

@ -1,10 +1,10 @@
from datetime import datetime, timedelta, timezone from datetime import datetime, timedelta, timezone
from typing import Any, Iterator, Tuple from typing import Any, Iterator, Tuple
from unittest.mock import MagicMock, patch from unittest.mock import patch
import pytest import pytest
from great_ai import MongoDbDriver from great_ai import MongoDbDriver
from great_ai.views import Filter
@pytest.fixture @pytest.fixture
@ -44,3 +44,24 @@ def test_empty_batch_is_a_noop(mongo: Tuple[MongoDbDriver, Any]) -> None:
driver, collection = mongo driver, collection = mongo
assert driver.save_batch([]) == [] assert driver.save_batch([]) == []
collection.insert_many.assert_not_called() collection.insert_many.assert_not_called()
def test_mongo_contains_uses_literal_case_insensitive_patterns(
mongo: Tuple[MongoDbDriver, Any],
) -> None:
driver, collection = mongo
driver.query(
conjunctive_filters=[Filter(property="text", operator="contains", value="a[")]
)
assert {
"text": {"$regex": r"a\[", "$options": "i"}
} in collection.find.call_args.kwargs["filter"]["$and"]
numeric = driver._get_filter(
Filter(property="score", operator="contains", value=2.5)
)
assert numeric == {
"$or": [
{"score": {"$eq": 2.5}},
{"score": {"$regex": r"2\.5", "$options": "i"}},
]
}