Normalize timezone-aware bounds and stored TinyDB timestamps

This commit is contained in:
Andras Schmelczer 2026-10-27 19:56:37 +00:00
commit d06e982b46
2 changed files with 27 additions and 2 deletions

View file

@ -6,6 +6,7 @@ import pandas as pd
from filelock import FileLock
from tinydb import TinyDB
from ..helper.normalise_datetime import normalise_datetime
from ..views import Filter, SortBy, Trace
from .tracing_database_driver import TracingDatabaseDriver
@ -52,11 +53,15 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
has_feedback: Optional[bool] = None,
sort_by: Sequence[SortBy] = [],
) -> Tuple[List[Trace], int]:
since = normalise_datetime(since) if since is not None else None
until = normalise_datetime(until) if until is not None else None
def does_match(d: Dict[str, Any]) -> bool:
created = normalise_datetime(datetime.fromisoformat(d["created"]))
return (
not set(conjunctive_tags) - set(d["tags"])
and (since is None or datetime.fromisoformat(d["created"]) >= since)
and (until is None or datetime.fromisoformat(d["created"]) <= until)
and (since is None or created >= since)
and (until is None or created <= until)
and (
has_feedback is None or has_feedback == (d["feedback"] is not None)
)

View file

@ -1,5 +1,6 @@
import multiprocessing as mp
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from unittest.mock import patch
@ -7,6 +8,7 @@ from unittest.mock import patch
from tinydb.table import Table
from great_ai import ParallelTinyDbDriver, Trace
from great_ai.context import get_context
def write_traces(path: str, index: int, barrier: Any) -> None:
@ -54,3 +56,21 @@ def test_spawned_writers_share_database_lock(tmp_path: Path) -> None:
process.close()
with patch.object(ParallelTinyDbDriver, "path_to_db", path):
assert len(ParallelTinyDbDriver().query()[0]) == 6
def test_tinydb_date_queries_normalize_timezones(isolated_context: None) -> None:
database = get_context().tracing_database
for i, created in enumerate(("2026-01-01T12:00:00", "2026-01-01T14:00:00+02:00")):
database.save(
Trace(
trace_id=str(i),
created=created,
original_execution_time_ms=0,
models=[],
tags=[],
logged_values={},
)
)
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