diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 554b86bd..ed1b3dbb 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -1,9 +1,10 @@ import asyncio import json import logging -from datetime import timedelta +from datetime import datetime, timedelta from importlib import metadata from pathlib import Path +from typing import Any from uuid import uuid4 import lancedb @@ -59,10 +60,13 @@ class Store: skip_validation: bool = False, create: bool = False, read_only: bool = False, + before: datetime | None = None, ): self.db_path: Path = db_path self._config = config - self._read_only = read_only + self._before = before + # Time-travel mode is always read-only + self._read_only = read_only or (before is not None) self.embedder = get_embedder(config=self._config) self._vacuum_lock = asyncio.Lock() @@ -89,9 +93,13 @@ class Store: # Initialize tables (creates them if they don't exist) self._init_tables() + # Checkout tables to historical state if before is specified + if before is not None: + self._checkout_tables_before(before) + # Run upgrades only on existing databases, set version for new ones # Skip upgrades in read-only mode (they would fail anyway) - if not read_only: + if not self._read_only: if is_new_db: self._set_initial_version() else: @@ -418,3 +426,83 @@ class Store: def _connection(self): """Compatibility property for repositories expecting _connection.""" return self + + def _checkout_tables_before(self, before: datetime) -> None: + """Checkout all tables to their state at or before the given datetime. + + Args: + before: The datetime to checkout to + + Raises: + ValueError: If no version exists before the given datetime + """ + # LanceDB stores timestamps as naive datetimes in local time. + # Convert 'before' to naive local time for comparison. + if before.tzinfo is not None: + # Convert to local time and make naive + before_local = before.astimezone().replace(tzinfo=None) + else: + # Already naive, assume local time + before_local = before + + tables = [ + ("documents", self.documents_table), + ("chunks", self.chunks_table), + ("settings", self.settings_table), + ] + + for table_name, table in tables: + versions = table.list_versions() + # Find the latest version at or before the target datetime + # Versions are sorted by version number, not timestamp, so we need to check all + best_version = None + best_timestamp = None + + for v in versions: + # LanceDB version timestamps are naive datetime objects in local time + v_timestamp = v["timestamp"] + # Make sure it's naive for comparison + if v_timestamp.tzinfo is not None: + v_timestamp = v_timestamp.replace(tzinfo=None) + + if v_timestamp <= before_local: + if best_timestamp is None or v_timestamp > best_timestamp: + best_version = v["version"] + best_timestamp = v_timestamp + + if best_version is None: + # Find the earliest version to report in error message + if versions: + earliest = min(versions, key=lambda v: v["timestamp"]) + earliest_ts = earliest["timestamp"] + raise ValueError( + f"No data exists before {before}. " + f"Database was created on {earliest_ts}" + ) + else: + raise ValueError( + f"No data exists before {before}. Table has no versions." + ) + + # Checkout to the found version + table.checkout(best_version) + + def list_table_versions(self, table_name: str) -> list[dict[str, Any]]: + """List version history for a table. + + Args: + table_name: Name of the table ("documents", "chunks", or "settings") + + Returns: + List of version info dicts with "version" and "timestamp" keys + """ + table_map = { + "documents": self.documents_table, + "chunks": self.chunks_table, + "settings": self.settings_table, + } + table = table_map.get(table_name) + if table is None: + raise ValueError(f"Unknown table: {table_name}") + + return list(table.list_versions()) diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index 496579d1..2b8b3fd8 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -1,8 +1,10 @@ import sys +from datetime import UTC, datetime from importlib import metadata from pathlib import Path from typing import TYPE_CHECKING, Any +from dateutil import parser as dateutil_parser from packaging.version import Version, parse if TYPE_CHECKING: @@ -12,6 +14,55 @@ if TYPE_CHECKING: from haiku.rag.graph.research.models import Citation +def parse_datetime(s: str) -> datetime: + """Parse a datetime string into a datetime object. + + Supports: + - ISO 8601 format: "2025-01-15T14:30:00", "2025-01-15T14:30:00Z", "2025-01-15T14:30:00+00:00" + - Date only: "2025-01-15" (interpreted as 00:00:00) + - Various other formats via dateutil + + Args: + s: String to parse + + Returns: + Parsed datetime object + + Raises: + ValueError: If the string cannot be parsed + """ + try: + return dateutil_parser.parse(s) + except (ValueError, TypeError) as e: + raise ValueError( + f"Could not parse datetime: {s}. " + "Use ISO 8601 format (e.g., 2025-01-15T14:30:00) or date (e.g., 2025-01-15)" + ) from e + + +def to_utc(dt: datetime) -> datetime: + """Convert a datetime to UTC. + + - Naive datetimes are assumed to be local time and converted to UTC + - Datetimes with timezone info are converted to UTC + - UTC datetimes are returned as-is + + Args: + dt: Datetime to convert + + Returns: + Datetime in UTC timezone + """ + if dt.tzinfo is None: + # Naive datetime - assume local time + local_dt = dt.astimezone() # Adds local timezone + return local_dt.astimezone(UTC) + elif dt.tzinfo == UTC: + return dt + else: + return dt.astimezone(UTC) + + def apply_common_settings( settings: Any | None, settings_class: type[Any], diff --git a/tests/store/test_time_travel.py b/tests/store/test_time_travel.py new file mode 100644 index 00000000..bc7474f6 --- /dev/null +++ b/tests/store/test_time_travel.py @@ -0,0 +1,112 @@ +import asyncio +from datetime import UTC, datetime, timedelta + +import pytest + +from haiku.rag.store import ReadOnlyError, Store +from haiku.rag.store.models import Document +from haiku.rag.store.repositories.document import DocumentRepository + + +class TestStoreTimeTravel: + def test_store_with_before_is_read_only(self, temp_db_path): + """Store with before parameter is automatically read-only.""" + # Create a store first + store = Store(temp_db_path, create=True) + store.close() + + # Open with before - should be read-only + before = datetime.now(UTC) + timedelta(hours=1) + store = Store(temp_db_path, before=before) + assert store.is_read_only is True + store.close() + + def test_store_before_raises_on_write(self, temp_db_path): + """Store with before parameter raises on write operations.""" + store = Store(temp_db_path, create=True) + store.close() + + before = datetime.now(UTC) + timedelta(hours=1) + store = Store(temp_db_path, before=before) + + with pytest.raises(ReadOnlyError): + store._assert_writable() + store.close() + + @pytest.mark.asyncio + async def test_store_before_checks_out_historical_state(self, temp_db_path): + """Store with before parameter checks out tables to historical state.""" + # Create store and add a document + store = Store(temp_db_path, create=True) + repo = DocumentRepository(store) + await repo.create(Document(content="First document")) + + # Get the version timestamp after first document + versions_after_first = store.list_table_versions("documents") + # Find the latest version timestamp + latest_version = max(versions_after_first, key=lambda v: v["version"]) + time_after_first = latest_version["timestamp"] + + # Wait a bit to ensure the next write gets a distinct timestamp + await asyncio.sleep(0.5) + + # Add second document + await repo.create(Document(content="Second document")) + + # Verify we have more versions now + versions_after_second = store.list_table_versions("documents") + assert len(versions_after_second) > len(versions_after_first) + + store.close() + + # Open at historical state (using the timestamp from after first write) + store = Store(temp_db_path, before=time_after_first) + repo = DocumentRepository(store) + + # Should only see first document + docs = await repo.list_all() + assert len(docs) == 1 + assert docs[0].content == "First document" + store.close() + + # Open at current state + store = Store(temp_db_path) + repo = DocumentRepository(store) + + # Should see both documents + docs = await repo.list_all() + assert len(docs) == 2 + store.close() + + def test_store_before_no_version_raises(self, temp_db_path): + """Store with before datetime before any version raises ValueError.""" + store = Store(temp_db_path, create=True) + store.close() + + # Try to open before the database was created + before = datetime(2000, 1, 1, tzinfo=UTC) + with pytest.raises(ValueError) as exc_info: + Store(temp_db_path, before=before) + assert "No data exists before" in str(exc_info.value) + + def test_current_table_versions_returns_versions(self, temp_db_path): + """current_table_versions returns dict of table versions.""" + store = Store(temp_db_path, create=True) + versions = store.current_table_versions() + + assert "documents" in versions + assert "chunks" in versions + assert "settings" in versions + assert all(isinstance(v, int) for v in versions.values()) + store.close() + + def test_list_table_versions_returns_history(self, temp_db_path): + """list_table_versions returns version history for a table.""" + store = Store(temp_db_path, create=True) + versions = store.list_table_versions("documents") + + assert len(versions) >= 1 + for v in versions: + assert "version" in v + assert "timestamp" in v + store.close() diff --git a/tests/utils/test_datetime.py b/tests/utils/test_datetime.py new file mode 100644 index 00000000..6d2f406b --- /dev/null +++ b/tests/utils/test_datetime.py @@ -0,0 +1,85 @@ +from datetime import UTC, datetime, timezone + +import pytest + +from haiku.rag.utils import parse_datetime, to_utc + + +class TestParseDateTime: + def test_parse_iso8601_with_timezone(self): + """Parse ISO 8601 datetime with timezone.""" + result = parse_datetime("2025-01-15T14:30:00+00:00") + assert result.year == 2025 + assert result.month == 1 + assert result.day == 15 + assert result.hour == 14 + assert result.minute == 30 + assert result.second == 0 + assert result.tzinfo is not None + + def test_parse_iso8601_without_timezone(self): + """Parse ISO 8601 datetime without timezone (naive).""" + result = parse_datetime("2025-01-15T14:30:00") + assert result.year == 2025 + assert result.month == 1 + assert result.day == 15 + assert result.hour == 14 + assert result.minute == 30 + + def test_parse_date_only(self): + """Parse date-only string as start of day.""" + result = parse_datetime("2025-01-15") + assert result.year == 2025 + assert result.month == 1 + assert result.day == 15 + assert result.hour == 0 + assert result.minute == 0 + assert result.second == 0 + + def test_parse_various_formats(self): + """Parse various datetime formats.""" + # ISO with Z suffix + result = parse_datetime("2025-01-15T14:30:00Z") + assert result.year == 2025 + assert result.month == 1 + assert result.day == 15 + + # With milliseconds + result = parse_datetime("2025-01-15T14:30:00.123") + assert result.microsecond == 123000 + + def test_parse_invalid_raises_value_error(self): + """Invalid datetime string raises ValueError.""" + with pytest.raises(ValueError) as exc_info: + parse_datetime("not-a-datetime") + assert "Could not parse datetime" in str(exc_info.value) + + +class TestToUtc: + def test_naive_datetime_assumes_local_and_converts(self): + """Naive datetime is assumed local and converted to UTC.""" + naive = datetime(2025, 1, 15, 14, 30, 0) + result = to_utc(naive) + assert result.tzinfo == UTC + + def test_utc_datetime_unchanged(self): + """UTC datetime is returned as-is.""" + utc_dt = datetime(2025, 1, 15, 14, 30, 0, tzinfo=UTC) + result = to_utc(utc_dt) + assert result == utc_dt + assert result.tzinfo == UTC + + def test_other_timezone_converts_to_utc(self): + """Datetime with other timezone is converted to UTC.""" + from datetime import timedelta + + # Create a datetime at UTC+5 + tz_plus5 = timezone(timedelta(hours=5)) + dt_plus5 = datetime(2025, 1, 15, 19, 30, 0, tzinfo=tz_plus5) + + result = to_utc(dt_plus5) + + # 19:30 UTC+5 = 14:30 UTC + assert result.tzinfo == UTC + assert result.hour == 14 + assert result.minute == 30