[pull] main from apache:main

burr/core/persistence.py

@@ -33,6 +33,13 @@
 except ImportError:
     Self = None
 
+# Error message template for uninitialized SQLitePersister
+_UNINITIALIZED_PERSISTER_ERROR = (
+    "Uninitialized persister: table '{table_name}' does not exist. "
+    "Make sure to call .initialize() on the persister before passing it "
+    "to the ApplicationBuilder."
+)
+
 
 class PersistedStateData(TypedDict):
     partition_key: str
@@ -444,12 +451,19 @@ def list_app_ids(self, partition_key: Optional[str], **kwargs) -> list[str]:
         )
 
         cursor = self.connection.cursor()
-        cursor.execute(
-            f"SELECT DISTINCT app_id FROM {self.table_name} "
-            f"WHERE partition_key = ? "
-            f"ORDER BY created_at DESC",
-            (partition_key,),
-        )
+        try:
+            cursor.execute(
+                f"SELECT DISTINCT app_id FROM {self.table_name} "
+                f"WHERE partition_key = ? "
+                f"ORDER BY created_at DESC",
+                (partition_key,),
+            )
+        except sqlite3.OperationalError as e:
+            if "no such table" in str(e):
+                raise RuntimeError(
+                    _UNINITIALIZED_PERSISTER_ERROR.format(table_name=self.table_name)
+                ) from e
+            raise
         app_ids = [row[0] for row in cursor.fetchall()]
         return app_ids
 
@@ -475,27 +489,34 @@ def load(
         )
         logger.debug("Loading %s, %s, %s", partition_key, app_id, sequence_id)
         cursor = self.connection.cursor()
-        if app_id is None:
-            # get latest for all app_ids
-            cursor.execute(
-                f"SELECT position, state, sequence_id, app_id, created_at, status FROM {self.table_name} "
-                f"WHERE partition_key = ? "
-                f"ORDER BY CREATED_AT DESC LIMIT 1",
-                (partition_key,),
-            )
-        elif sequence_id is None:
-            cursor.execute(
-                f"SELECT position, state, sequence_id, app_id, created_at, status FROM {self.table_name} "
-                f"WHERE partition_key = ? AND app_id = ? "
-                f"ORDER BY sequence_id DESC LIMIT 1",
-                (partition_key, app_id),
-            )
-        else:
-            cursor.execute(
-                f"SELECT position, state, sequence_id, app_id, created_at, status FROM {self.table_name} "
-                f"WHERE partition_key = ? AND app_id = ? AND sequence_id = ?",
-                (partition_key, app_id, sequence_id),
-            )
+        try:
+            if app_id is None:
+                # get latest for all app_ids
+                cursor.execute(
+                    f"SELECT position, state, sequence_id, app_id, created_at, status FROM {self.table_name} "
+                    f"WHERE partition_key = ? "
+                    f"ORDER BY CREATED_AT DESC LIMIT 1",
+                    (partition_key,),
+                )
+            elif sequence_id is None:
+                cursor.execute(
+                    f"SELECT position, state, sequence_id, app_id, created_at, status FROM {self.table_name} "
+                    f"WHERE partition_key = ? AND app_id = ? "
+                    f"ORDER BY sequence_id DESC LIMIT 1",
+                    (partition_key, app_id),
+                )
+            else:
+                cursor.execute(
+                    f"SELECT position, state, sequence_id, app_id, created_at, status FROM {self.table_name} "
+                    f"WHERE partition_key = ? AND app_id = ? AND sequence_id = ?",
+                    (partition_key, app_id, sequence_id),
+                )
+        except sqlite3.OperationalError as e:
+            if "no such table" in str(e):
+                raise RuntimeError(
+                    _UNINITIALIZED_PERSISTER_ERROR.format(table_name=self.table_name)
+                ) from e
+            raise
         row = cursor.fetchone()
         if row is None:
             return None
@@ -551,11 +572,18 @@ def save(
         )
         cursor = self.connection.cursor()
         json_state = json.dumps(state.serialize(**self.serde_kwargs))
-        cursor.execute(
-            f"INSERT INTO {self.table_name} (partition_key, app_id, sequence_id, position, state, status) "
-            f"VALUES (?, ?, ?, ?, ?, ?)",
-            (partition_key, app_id, sequence_id, position, json_state, status),
-        )
+        try:
+            cursor.execute(
+                f"INSERT INTO {self.table_name} (partition_key, app_id, sequence_id, position, state, status) "
+                f"VALUES (?, ?, ?, ?, ?, ?)",
+                (partition_key, app_id, sequence_id, position, json_state, status),
+            )
+        except sqlite3.OperationalError as e:
+            if "no such table" in str(e):
+                raise RuntimeError(
+                    _UNINITIALIZED_PERSISTER_ERROR.format(table_name=self.table_name)
+                ) from e
+            raise
         self.connection.commit()
 
     def cleanup(self):

burr/integrations/pydantic.py

@@ -129,14 +129,17 @@ def _validate_and_extract_signature_types(
             "action must be the state object. Got signature: {sig}."
         )
     type_hints = typing.get_type_hints(fn)
+    state_model = type_hints.get("state")
 
-    if (state_model := type_hints["state"]) is inspect.Parameter.empty or not issubclass(
-        state_model, pydantic.BaseModel
+    if (
+        state_model is None
+        or state_model is inspect.Parameter.empty
+        or not issubclass(state_model, pydantic.BaseModel)
     ):
         raise ValueError(
             f"Function fn: {fn.__qualname__} is not a valid pydantic action. "
-            "a type annotation of a type extending: pydantic.BaseModel. Got parameter "
-            "state: {state_model.__qualname__}."
+            "The 'state' parameter must be annotated with a type extending pydantic.BaseModel. "
+            f"Got: {state_model}."
         )
     if (ret_hint := type_hints.get("return")) is None or not issubclass(
         ret_hint, pydantic.BaseModel

docs/concepts/serde.rst

@@ -19,6 +19,7 @@
 
 .. _serde:
 
+
 ================================
 Serialization / Deserialization
 ================================
@@ -130,3 +131,11 @@ Requirements for the serializer/deserializer functions:
 
     1. The serializer function needs to return a dictionary.
     2. Both function signatures needs to have a ``**kwargs`` parameter to allow for custom arguments to be passed in. We advise namespacing the kwargs provided to avoid conflicts with other serializers/deserializers.
+
+
+Limitations of State Serialization
+----------------------------------
+
+Currently, it is only possible to override the serialization and deserialization behavior of individual fields within a state.
+Overriding the serialization of the entire state object as a whole is not supported at this time.
+If you need custom handling, you must apply it at the field level.

examples/email-assistant/application.py

@@ -54,7 +54,7 @@ def determine_clarifications(state: State) -> State:
     client = _get_openai_client()
     # TODO -- use instructor to get a pydantic model
     result = client.chat.completions.create(
-        model="gpt-4",
+        model="gpt-4o-mini",
         messages=[
             {
                 "role": "system",
@@ -130,7 +130,7 @@ def formulate_draft(state: State) -> Tuple[dict, State]:
     prompt = " ".join(instructions)
 
     result = client.chat.completions.create(
-        model="gpt-4",
+        model="gpt-4o-mini",
         messages=[
             {
                 "role": "system",

tests/core/test_persistence.py

@@ -94,6 +94,38 @@ def test_sqlite_persistence_is_initialized_true_new_connection(tmp_path):
     p2.cleanup()
 
 
+def test_sqlite_persister_load_without_initialize_raises_runtime_error():
+    """Test that calling load() without initialize() raises a clear RuntimeError."""
+    persister = SQLLitePersister(db_path=":memory:", table_name="test_table")
+    try:
+        with pytest.raises(RuntimeError, match="Uninitialized persister"):
+            persister.load("partition_key", "app_id")
+    finally:
+        persister.cleanup()
+
+
+def test_sqlite_persister_save_without_initialize_raises_runtime_error():
+    """Test that calling save() without initialize() raises a clear RuntimeError."""
+    persister = SQLLitePersister(db_path=":memory:", table_name="test_table")
+    try:
+        with pytest.raises(RuntimeError, match="Uninitialized persister"):
+            persister.save(
+                "partition_key", "app_id", 1, "position", State({"key": "value"}), "completed"
+            )
+    finally:
+        persister.cleanup()
+
+
+def test_sqlite_persister_list_app_ids_without_initialize_raises_runtime_error():
+    """Test that calling list_app_ids() without initialize() raises a clear RuntimeError."""
+    persister = SQLLitePersister(db_path=":memory:", table_name="test_table")
+    try:
+        with pytest.raises(RuntimeError, match="Uninitialized persister"):
+            persister.list_app_ids("partition_key")
+    finally:
+        persister.cleanup()
+
+
 @pytest.mark.parametrize(
     "method_name,kwargs",
     [

tests/integrations/test_burr_pydantic.py

@@ -185,6 +185,10 @@ def _fn_with_no_return_type(state: OriginalModel):
     ...
 
 
+def _fn_with_untyped_state_arg(state) -> OriginalModel:
+    ...
+
+
 def _fn_correct_same_itype_otype(state: OriginalModel, input_1: int) -> OriginalModel:
     ...
 
@@ -200,13 +204,21 @@ def _fn_correct_diff_itype_otype(state: OriginalModel, input_1: int) -> NestedMo
         (_fn_with_incorrect_state_arg, ValueError),
         (_fn_with_incorrect_return_type, ValueError),
         (_fn_with_no_return_type, ValueError),
+        (_fn_with_untyped_state_arg, ValueError),
     ],
 )
 def test__validate_and_extract_signature_types_error(fn, expected_exception):
     with pytest.raises(expected_exception=expected_exception):
         _validate_and_extract_signature_types(fn)
 
 
+def test__validate_and_extract_signature_types_untyped_state_error_message():
+    """Test that the error message is informative when state is not type-annotated."""
+    with pytest.raises(ValueError) as excinfo:
+        _validate_and_extract_signature_types(_fn_with_untyped_state_arg)
+    assert "'state' parameter must be annotated" in str(excinfo.value)
+
+
 @pytest.mark.parametrize(
     "fn,expected",
     [