diff --git a/src/rai_core/rai/messages/artifacts.py b/src/rai_core/rai/messages/artifacts.py index dbc34537a..e8d8c0401 100644 --- a/src/rai_core/rai/messages/artifacts.py +++ b/src/rai_core/rai/messages/artifacts.py @@ -26,18 +26,19 @@ def store_artifacts( tool_call_id: str, artifacts: List[Any], db_path="artifact_database.pkl" ): # TODO(boczekbartek): refactor - db_path = Path(db_path) - if not db_path.is_file(): - artifact_database = {} - with open("artifact_database.pkl", "wb") as file: + path = Path(db_path) + if not path.is_file(): + path.parent.mkdir(parents=True, exist_ok=True) + artifact_database: dict = {} + with path.open("wb") as file: pickle.dump(artifact_database, file) - with open("artifact_database.pkl", "rb") as file: + with path.open("rb") as file: artifact_database = pickle.load(file) if tool_call_id not in artifact_database: artifact_database[tool_call_id] = artifacts else: artifact_database[tool_call_id].extend(artifacts) - with open("artifact_database.pkl", "wb") as file: + with path.open("wb") as file: pickle.dump(artifact_database, file) @@ -45,11 +46,11 @@ def get_stored_artifacts( tool_call_id: str, db_path="artifact_database.pkl" ) -> List[Any]: # TODO(boczekbartek): refactor - db_path = Path(db_path) - if not db_path.is_file(): + db_path_p = Path(db_path) + if not db_path_p.is_file(): return [] - with db_path.open("rb") as db: + with db_path_p.open("rb") as db: artifact_database = pickle.load(db) if tool_call_id in artifact_database: return artifact_database[tool_call_id] diff --git a/tests/messages/test_artifacts_db_path.py b/tests/messages/test_artifacts_db_path.py new file mode 100644 index 000000000..285ad91b9 --- /dev/null +++ b/tests/messages/test_artifacts_db_path.py @@ -0,0 +1,33 @@ +# Copyright (C) 2025 Robotec.AI +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +from rai.messages.artifacts import get_stored_artifacts, store_artifacts + + +def test_store_artifacts_honors_db_path(tmp_path: Path): + db = tmp_path / "custom_artifacts.pkl" + store_artifacts("tc1", ["a"], db_path=str(db)) + assert db.is_file() + assert get_stored_artifacts("tc1", db_path=str(db)) == ["a"] + store_artifacts("tc1", ["b"], db_path=str(db)) + assert get_stored_artifacts("tc1", db_path=str(db)) == ["a", "b"] + assert get_stored_artifacts("missing", db_path=str(db)) == [] + + +def test_store_artifacts_creates_nested_path(tmp_path: Path): + db = tmp_path / "nested" / "dir" / "db.pkl" + store_artifacts("x", [1], db_path=str(db)) + assert get_stored_artifacts("x", db_path=str(db)) == [1]