diff --git a/burr/core/persistence.py b/burr/core/persistence.py index f217f739d..c32bf8e96 100644 --- a/burr/core/persistence.py +++ b/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): diff --git a/burr/integrations/pydantic.py b/burr/integrations/pydantic.py index 5354537a2..dd1f95a58 100644 --- a/burr/integrations/pydantic.py +++ b/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 diff --git a/docs/concepts/serde.rst b/docs/concepts/serde.rst index c2ad47bda..20ad7a92f 100644 --- a/docs/concepts/serde.rst +++ b/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. diff --git a/examples/email-assistant/application.py b/examples/email-assistant/application.py index ec74bf911..79b23b5c0 100644 --- a/examples/email-assistant/application.py +++ b/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", diff --git a/tests/core/test_persistence.py b/tests/core/test_persistence.py index b990b7d6d..b362cd96d 100644 --- a/tests/core/test_persistence.py +++ b/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", [ diff --git a/tests/integrations/test_burr_pydantic.py b/tests/integrations/test_burr_pydantic.py index 6c25520f6..8fbb94c02 100644 --- a/tests/integrations/test_burr_pydantic.py +++ b/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,6 +204,7 @@ 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): @@ -207,6 +212,13 @@ def test__validate_and_extract_signature_types_error(fn, 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", [