diff --git a/great_ai/persistence/mongodb_driver.py b/great_ai/persistence/mongodb_driver.py index c0df0a0..6035e08 100644 --- a/great_ai/persistence/mongodb_driver.py +++ b/great_ai/persistence/mongodb_driver.py @@ -1,3 +1,4 @@ +import re from datetime import datetime from typing import Any, Dict, List, Optional, Sequence, Tuple @@ -102,6 +103,19 @@ class MongoDbDriver(TracingDatabaseDriver): return operator_mapping["="] 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( self, *, @@ -121,9 +135,7 @@ class MongoDbDriver(TracingDatabaseDriver): and_query: List[Dict[str, Any]] = [] and_query.extend({"tags": tag} for tag in conjunctive_tags) - and_query.extend( - {f.property: {self._get_operator(f): f.value}} for f in conjunctive_filters - ) + and_query.extend(self._get_filter(f) for f in conjunctive_filters) if not and_query: and_query.append({}) diff --git a/great_ai/persistence/parallel_tinydb_driver.py b/great_ai/persistence/parallel_tinydb_driver.py index 1b5f2af..4250600 100644 --- a/great_ai/persistence/parallel_tinydb_driver.py +++ b/great_ai/persistence/parallel_tinydb_driver.py @@ -84,10 +84,14 @@ class ParallelTinyDbDriver(TracingDatabaseDriver): getattr(df[f.property], operator_mapping[f.operator])(f.value) ] elif operator == "contains": + pattern = f.value_as_text.casefold() df = df.loc[ - df[f.property].str.contains( - str(int(f.value)) if isinstance(f.value, float) else f.value, - case=False, + df[f.property].map( + lambda value: ( + pattern in value.casefold() + if isinstance(value, str) + else not isinstance(f.value, str) and value == f.value + ) ) ] diff --git a/great_ai/views/filter.py b/great_ai/views/filter.py index 49a6f66..dcd8a4c 100644 --- a/great_ai/views/filter.py +++ b/great_ai/views/filter.py @@ -1,3 +1,4 @@ +import builtins from typing import Union from pydantic import BaseModel @@ -9,3 +10,9 @@ class Filter(BaseModel): property: str operator: Operator 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) diff --git a/tests/test_database_regressions.py b/tests/test_database_regressions.py index e2c84da..d223b09 100644 --- a/tests/test_database_regressions.py +++ b/tests/test_database_regressions.py @@ -5,10 +5,10 @@ from pathlib import Path from typing import Any from unittest.mock import patch -from tinydb.table import Table - from great_ai import ParallelTinyDbDriver, Trace 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: @@ -74,3 +74,45 @@ def test_tinydb_date_queries_normalize_timezones(isolated_context: None) -> None noon = datetime(2026, 1, 1, 12, tzinfo=timezone.utc) assert database.query(since=noon, until=noon)[1] == 2 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 + ) diff --git a/tests/test_mongodb_queries.py b/tests/test_mongodb_queries.py index fd2e460..e4d057e 100644 --- a/tests/test_mongodb_queries.py +++ b/tests/test_mongodb_queries.py @@ -1,10 +1,10 @@ from datetime import datetime, timedelta, timezone from typing import Any, Iterator, Tuple -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest - from great_ai import MongoDbDriver +from great_ai.views import Filter @pytest.fixture @@ -44,3 +44,24 @@ def test_empty_batch_is_a_noop(mongo: Tuple[MongoDbDriver, Any]) -> None: driver, collection = mongo assert driver.save_batch([]) == [] 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"}}, + ] + }