Replace begin/commit/rollback_transaction with 'async with storage_provider: ...'

.github/workflows/codeql.yml

@@ -63,6 +63,11 @@ jobs:
     # - name: Setup runtime (example)
     #   uses: actions/setup-example@v1
 
+      # Force Python 3.12 for this job (uv will pick this up)
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.12.x'
+
     # Initializes the CodeQL tools for scanning.
     - name: Initialize CodeQL
       uses: github/codeql-action/init@v3

TADA.md

@@ -4,7 +4,10 @@ Talk at PyBay is on Sat, Oct 18 in SF
 
 ## Software
 
-- Unify Podcast and Transcript (use shared message and metadata classes)?
+- We should change the commit state machine into a context manager
+  ('with storage_provider: ...') rather than requiring the user to do
+  the commit/rollback logic
+
 - Distinguish between release deps and build/dev deps?
 
 ### Specifically for VTT import (minor):

make.bat

@@ -50,7 +50,7 @@ goto end
 :build
 if not exist ".venv\" call make.bat venv
 echo Building package...
-.venv\Scripts\python -m build --wheel --installer uv
+uv build
 goto end
 
 :venv

pyproject.toml

@@ -59,7 +59,7 @@ packages = [
 
 [tool.setuptools.package-data]
 "typeagent.podcasts" = ["*.json"]
-"typeagent.emails" = ["*.json"]
+"typeagent.emails" = ["*.json", "*.txt"]
 
 [tool.pytest.ini_options]
 asyncio_default_fixture_loop_scope = "function"

test/test_add_messages_with_indexing.py

@@ -148,22 +148,17 @@ async def test_transaction_rollback_on_error():
 
         initial_count = await transcript.messages.size()
 
-        # Try to add batch with error (empty text_chunks should work, so this isn't a good test)
-        # Instead, let's just verify the transaction state machine works
+        # Verify the transaction context manager works
         try:
-            await storage.begin_transaction()
-            await storage.commit_transaction()
+            async with storage:
+                pass
         except Exception:
-            pytest.fail("Transaction state machine should work")
+            pytest.fail("Transaction context manager should work")
 
-        # Verify transaction state errors
-        with pytest.raises(RuntimeError, match="Transaction already active"):
-            await storage.begin_transaction()
-            await storage.begin_transaction()
-
-        await storage.rollback_transaction()
-
-        with pytest.raises(RuntimeError, match="Cannot commit"):
-            await storage.commit_transaction()
+        # Verify nested transactions fail
+        with pytest.raises(Exception):  # SQLite will raise an OperationalError
+            async with storage:
+                async with storage:  # This should fail
+                    pass
 
         await storage.close()

tools/query.py

@@ -482,7 +482,7 @@ def prsep():
 
 
 def make_arg_parser(description: str) -> argparse.ArgumentParser:
-    line_width = utils.cap(144, shutil.get_terminal_size().columns)
+    line_width = min(144, shutil.get_terminal_size().columns)
     parser = argparse.ArgumentParser(
         description=description,
         formatter_class=lambda *a, **b: argparse.HelpFormatter(

tools/test_email.py

@@ -75,8 +75,10 @@ async def main():
 
     if sys.argv[1:2]:
         base_path = Path(sys.argv[1])
-    else:
+    elif os.path.exists("/data"):
         base_path = Path("/data/testChat/knowpro/email/")
+    else:
+        base_path = Path(".")
 
     try:
         base_path.mkdir(parents=True, exist_ok=True)

typeagent/aitools/utils.py

@@ -17,8 +17,6 @@
 
 from pydantic_ai import Agent
 
-cap = min  # More readable name for capping a value at some limit.
-
 
 @contextmanager
 def timelog(label: str, verbose: bool = True):
@@ -51,7 +49,7 @@ def format_code(text: str, line_width=None) -> str:
     """
     if line_width is None:
         # Use the terminal width, but cap it to 200 characters.
-        line_width = cap(200, shutil.get_terminal_size().columns)
+        line_width = min(200, shutil.get_terminal_size().columns)
     formatted_text = black.format_str(
         text, mode=black.FileMode(line_length=line_width)
     ).rstrip()

typeagent/knowpro/conversation_base.py

@@ -136,9 +136,7 @@ async def add_messages_with_indexing(
         """
         storage = await self.settings.get_storage_provider()
 
-        await storage.begin_transaction()
-
-        try:
+        async with storage:
             start_points = IndexingStartPoints(
                 message_count=await self.messages.size(),
                 semref_count=await self.semantic_refs.size(),
@@ -162,14 +160,8 @@ async def add_messages_with_indexing(
                 - start_points.semref_count,
             )
 
-            await storage.commit_transaction()
-
             return result
 
-        except BaseException:
-            await storage.rollback_transaction()
-            raise
-
     async def _add_metadata_knowledge_incremental(
         self,
         start_from_message_ordinal: int,

typeagent/knowpro/interfaces.py

@@ -50,13 +50,6 @@ class DeletionInfo:
     reason: str | None = None
 
 
-class TransactionState(Enum):
-    """State machine for storage provider transactions."""
-
-    NONE = "none"
-    ACTIVE = "active"
-
-
 @dataclass
 class IndexingStartPoints:
     """Track collection sizes before adding new items."""
@@ -840,16 +833,17 @@ async def get_related_terms_index(self) -> ITermToRelatedTermsIndex: ...
     async def get_conversation_threads(self) -> IConversationThreads: ...
 
     # Transaction management
-    async def begin_transaction(self) -> None:
-        """Begin a transaction. Must not be in active transaction."""
+    async def __aenter__(self) -> Self:
+        """Enter transaction context. Calls begin_transaction()."""
         ...
 
-    async def commit_transaction(self) -> None:
-        """Commit active transaction. Must be in ACTIVE state."""
-        ...
-
-    async def rollback_transaction(self) -> None:
-        """Rollback active transaction. Must be in ACTIVE state."""
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: Any,
+    ) -> None:
+        """Exit transaction context. Commits on success, rolls back on exception."""
         ...
 
     async def close(self) -> None: ...

typeagent/storage/memory/provider.py

@@ -53,21 +53,17 @@ def __init__(
         thread_settings = message_text_settings.embedding_index_settings
         self._conversation_threads = ConversationThreads(thread_settings)
 
-    async def begin_transaction(self) -> None:
-        """Begin a transaction. No-op for in-memory storage."""
-        pass
-
-    async def commit_transaction(self) -> None:
-        """Commit active transaction. No-op for in-memory storage."""
-        pass
+    async def __aenter__(self) -> "MemoryStorageProvider[TMessage]":
+        """Enter transaction context. No-op for in-memory storage."""
+        return self
 
-    async def rollback_transaction(self) -> None:
-        """
-        Rollback active transaction. No-op for in-memory storage.
-
-        Note: In-memory storage cannot rollback changes as they are
-        immediately applied to the in-memory data structures.
-        """
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: object,
+    ) -> None:
+        """Exit transaction context. No-op for in-memory storage."""
         pass
 
     async def get_semantic_ref_index(self) -> ITermToSemanticRefIndex:

typeagent/storage/sqlite/provider.py

@@ -106,33 +106,22 @@ def __init__(
         # Connect message collection to message text index for automatic indexing
         self._message_collection.set_message_text_index(self._message_text_index)
 
-        # Transaction state tracking
-        self._transaction_state = interfaces.TransactionState.NONE
-
-    async def begin_transaction(self) -> None:
-        """Begin a transaction. Must not be in active transaction."""
-        if self._transaction_state == interfaces.TransactionState.ACTIVE:
-            raise RuntimeError("Transaction already active")
+    async def __aenter__(self) -> "SqliteStorageProvider[TMessage]":
+        """Enter transaction context."""
         self.db.execute("BEGIN IMMEDIATE")
-        self._transaction_state = interfaces.TransactionState.ACTIVE
+        return self
 
-    async def commit_transaction(self) -> None:
-        """Commit active transaction. Must be in ACTIVE state."""
-        if self._transaction_state != interfaces.TransactionState.ACTIVE:
-            raise RuntimeError(
-                f"Cannot commit: transaction state is {self._transaction_state.value}"
-            )
-        self.db.commit()
-        self._transaction_state = interfaces.TransactionState.NONE
-
-    async def rollback_transaction(self) -> None:
-        """Rollback active transaction. Must be in ACTIVE state."""
-        if self._transaction_state != interfaces.TransactionState.ACTIVE:
-            raise RuntimeError(
-                f"Cannot rollback: transaction state is {self._transaction_state.value}"
-            )
-        self.db.rollback()
-        self._transaction_state = interfaces.TransactionState.NONE
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: object,
+    ) -> None:
+        """Exit transaction context. Commits on success, rolls back on exception."""
+        if exc_type is None:
+            self.db.commit()
+        else:
+            self.db.rollback()
 
     async def close(self) -> None:
         """Close the database connection. COMMITS."""