Checkout version logic and datetime utils

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())

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],

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()

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