From 1f84f07d3ce174dfa069d754511006d13d7aa289 Mon Sep 17 00:00:00 2001 From: Jakub Walaszczyk Date: Tue, 26 May 2026 10:52:45 +0200 Subject: [PATCH 1/3] Enhances vector store registration error handling Signed-off-by: Jakub Walaszczyk Assisted-by: Claude Code --- ai4rag/core/experiment/exception_handler.py | 16 +++++- ai4rag/core/experiment/experiment.py | 18 +++--- scripts/format.sh | 4 +- .../core/experiment/test_exception_handler.py | 55 +++++++++++++++---- 4 files changed, 71 insertions(+), 22 deletions(-) diff --git a/ai4rag/core/experiment/exception_handler.py b/ai4rag/core/experiment/exception_handler.py index c1bf54a..d7c4360 100644 --- a/ai4rag/core/experiment/exception_handler.py +++ b/ai4rag/core/experiment/exception_handler.py @@ -39,6 +39,20 @@ def __repr__(self) -> str: ) +class VectorStoreInitializationError(AI4RAGError): + """Exception representing error during vector store creation or retrieval.""" + + def __init__(self, exception, embedding_model_id): + super().__init__(exception) + self.embedding_model_id = embedding_model_id + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}: Unable to initialize vector store for embedding model " + f"'{self.embedding_model_id}' due to: {repr(self.exception)}" + ) + + class GenerationError(AI4RAGError): """Exception representing error during retrieval or inference.""" @@ -115,4 +129,4 @@ def get_final_error_msg(self) -> str: error_content = next((er for er in self.errors if most_common_error_type_name in er.__class__.__name__)) - return f"{error_content}. " f"To find more details please see generated logs file." + return error_content diff --git a/ai4rag/core/experiment/experiment.py b/ai4rag/core/experiment/experiment.py index 8856c77..0721a8e 100644 --- a/ai4rag/core/experiment/experiment.py +++ b/ai4rag/core/experiment/experiment.py @@ -17,6 +17,7 @@ AssetSaveError, ExperimentExceptionHandler, IndexingError, + VectorStoreInitializationError, ) from ai4rag.core.experiment.mps import ModelsPreSelector from ai4rag.core.experiment.results import EvaluationResult, ExperimentResults @@ -366,13 +367,16 @@ def run_single_evaluation(self, rag_params: RAGParamsType) -> float: reuse_collection_name = self._get_reusable_collection_name(indexing_params=indexing_params) - vector_store = get_vector_store( - vs_type=self.vector_store_type, - embedding_model=embedding_model, - reuse_collection_name=reuse_collection_name, - client=self.client, - ogx_vector_io_provider_id=self.ogx_vector_io_provider_id, - ) + try: + vector_store = get_vector_store( + vs_type=self.vector_store_type, + embedding_model=embedding_model, + reuse_collection_name=reuse_collection_name, + client=self.client, + ogx_vector_io_provider_id=self.ogx_vector_io_provider_id, + ) + except Exception as exc: + raise VectorStoreInitializationError(exc, embedding_model_id=embedding_model.model_id) from exc collection_name = vector_store.collection_name diff --git a/scripts/format.sh b/scripts/format.sh index 2c84d76..71bbc2f 100755 --- a/scripts/format.sh +++ b/scripts/format.sh @@ -12,11 +12,11 @@ echo "Running formatters on: ${TARGETS[*]}" echo "" echo "==> isort" -uv run isort "${TARGETS[@]}" +uv run --extra code_check isort "${TARGETS[@]}" echo "" echo "==> black" -uv run black "${TARGETS[@]}" +uv run --extra code_check black "${TARGETS[@]}" echo "" echo "==> copyright_check" diff --git a/tests/unit/ai4rag/core/experiment/test_exception_handler.py b/tests/unit/ai4rag/core/experiment/test_exception_handler.py index 71ae950..7390d44 100644 --- a/tests/unit/ai4rag/core/experiment/test_exception_handler.py +++ b/tests/unit/ai4rag/core/experiment/test_exception_handler.py @@ -9,6 +9,7 @@ ExperimentExceptionHandler, GenerationError, IndexingError, + VectorStoreInitializationError, ) from ai4rag.utils.event_handler import LogLevel @@ -164,6 +165,33 @@ def test_asset_save_error_is_ai4rag_error(self): assert issubclass(AssetSaveError, AI4RAGError) +class TestVectorStoreInitializationError: + """Test suite for VectorStoreInitializationError exception class.""" + + def test_creation(self): + """Test creating VectorStoreInitializationError.""" + base_exception = ConnectionError("OGX unavailable") + error = VectorStoreInitializationError(base_exception, "embedding-model-1") + + assert error.exception == base_exception + assert error.embedding_model_id == "embedding-model-1" + + def test_repr(self): + """Test VectorStoreInitializationError __repr__.""" + base_exception = ConnectionError("OGX unavailable") + error = VectorStoreInitializationError(base_exception, "embedding-model-1") + + repr_str = repr(error) + assert "VectorStoreInitializationError" in repr_str + assert "Unable to initialize vector store" in repr_str + assert "embedding-model-1" in repr_str + assert "ConnectionError" in repr_str + + def test_is_ai4rag_error(self): + """Test that VectorStoreInitializationError inherits from AI4RAGError.""" + assert issubclass(VectorStoreInitializationError, AI4RAGError) + + class TestExperimentExceptionsHandlerInitialization: """Test suite for ExperimentExceptionsHandler initialization.""" @@ -284,11 +312,10 @@ def test_get_final_error_msg_single_error(self, mocker): # Reset logger mock after handle_exception mock_logger.reset_mock() - msg = handler.get_final_error_msg() + result = handler.get_final_error_msg() - assert isinstance(msg, str) - assert "IndexingError" in msg - assert "please see generated logs file" in msg.lower() + assert isinstance(result, IndexingError) + assert result is error mock_logger.error.assert_called_once() def test_get_final_error_msg_multiple_errors_same_type(self, mocker): @@ -307,9 +334,10 @@ def test_get_final_error_msg_multiple_errors_same_type(self, mocker): mock_logger.reset_mock() - msg = handler.get_final_error_msg() + result = handler.get_final_error_msg() - assert "IndexingError" in msg + assert isinstance(result, IndexingError) + assert result is errors[0] mock_logger.error.assert_called_once() def test_get_final_error_msg_multiple_error_types(self, mocker): @@ -328,10 +356,11 @@ def test_get_final_error_msg_multiple_error_types(self, mocker): mock_logger.reset_mock() - msg = handler.get_final_error_msg() + result = handler.get_final_error_msg() # Most common error type is IndexingError (2 occurrences) - assert "IndexingError" in msg + assert isinstance(result, IndexingError) + assert result is errors[0] mock_logger.error.assert_called_once() def test_get_final_error_msg_logs_all_errors(self, mocker): @@ -385,10 +414,11 @@ def test_full_workflow_with_event_handler(self, mocker): # Get final error message mock_logger.reset_mock() - msg = handler.get_final_error_msg() + result = handler.get_final_error_msg() # Most common error is IndexingError - assert "IndexingError" in msg + assert isinstance(result, IndexingError) + assert result is errors[0] mock_logger.error.assert_called_once() def test_full_workflow_without_event_handler(self, mocker): @@ -409,6 +439,7 @@ def test_full_workflow_without_event_handler(self, mocker): assert len(handler.errors) == 2 mock_logger.reset_mock() - msg = handler.get_final_error_msg() + result = handler.get_final_error_msg() - assert "GenerationError" in msg + assert isinstance(result, GenerationError) + assert result is errors[0] From 9859d47e0f2a4fd45248a64bc7750834d84e6579 Mon Sep 17 00:00:00 2001 From: Jakub Walaszczyk Date: Tue, 26 May 2026 10:58:42 +0200 Subject: [PATCH 2/3] Update error message Signed-off-by: Jakub Walaszczyk --- ai4rag/core/experiment/exception_handler.py | 7 ++++--- ai4rag/core/experiment/experiment.py | 6 +++++- .../ai4rag/core/experiment/test_exception_handler.py | 12 ++++++++++-- 3 files changed, 19 insertions(+), 6 deletions(-) diff --git a/ai4rag/core/experiment/exception_handler.py b/ai4rag/core/experiment/exception_handler.py index d7c4360..2020315 100644 --- a/ai4rag/core/experiment/exception_handler.py +++ b/ai4rag/core/experiment/exception_handler.py @@ -42,14 +42,15 @@ def __repr__(self) -> str: class VectorStoreInitializationError(AI4RAGError): """Exception representing error during vector store creation or retrieval.""" - def __init__(self, exception, embedding_model_id): + def __init__(self, exception, embedding_model_id, vector_store_provider_id): super().__init__(exception) self.embedding_model_id = embedding_model_id + self.vector_store_provider_id = vector_store_provider_id def __repr__(self) -> str: return ( - f"{self.__class__.__name__}: Unable to initialize vector store for embedding model " - f"'{self.embedding_model_id}' due to: {repr(self.exception)}" + f"{self.__class__.__name__}: Unable to initialize vector store for {self.vector_store_provider_id} " + f"embedding model '{self.embedding_model_id}' due to: {repr(self.exception)}" ) diff --git a/ai4rag/core/experiment/experiment.py b/ai4rag/core/experiment/experiment.py index 0721a8e..3a3e262 100644 --- a/ai4rag/core/experiment/experiment.py +++ b/ai4rag/core/experiment/experiment.py @@ -376,7 +376,11 @@ def run_single_evaluation(self, rag_params: RAGParamsType) -> float: ogx_vector_io_provider_id=self.ogx_vector_io_provider_id, ) except Exception as exc: - raise VectorStoreInitializationError(exc, embedding_model_id=embedding_model.model_id) from exc + raise VectorStoreInitializationError( + exc, + embedding_model_id=embedding_model.model_id, + vector_store_provider_id=self.ogx_vector_io_provider_id or "local_chroma" + ) from exc collection_name = vector_store.collection_name diff --git a/tests/unit/ai4rag/core/experiment/test_exception_handler.py b/tests/unit/ai4rag/core/experiment/test_exception_handler.py index 7390d44..6cf6139 100644 --- a/tests/unit/ai4rag/core/experiment/test_exception_handler.py +++ b/tests/unit/ai4rag/core/experiment/test_exception_handler.py @@ -171,7 +171,11 @@ class TestVectorStoreInitializationError: def test_creation(self): """Test creating VectorStoreInitializationError.""" base_exception = ConnectionError("OGX unavailable") - error = VectorStoreInitializationError(base_exception, "embedding-model-1") + error = VectorStoreInitializationError( + base_exception, + embedding_model_id="embedding-model-1", + vector_store_provider_id="vs-provider-id", + ) assert error.exception == base_exception assert error.embedding_model_id == "embedding-model-1" @@ -179,7 +183,11 @@ def test_creation(self): def test_repr(self): """Test VectorStoreInitializationError __repr__.""" base_exception = ConnectionError("OGX unavailable") - error = VectorStoreInitializationError(base_exception, "embedding-model-1") + error = VectorStoreInitializationError( + base_exception, + embedding_model_id="embedding-model-1", + vector_store_provider_id="vs-provider-id", + ) repr_str = repr(error) assert "VectorStoreInitializationError" in repr_str From e82b228be834d347d03245854c3c6917c28fae14 Mon Sep 17 00:00:00 2001 From: Jakub Walaszczyk Date: Tue, 26 May 2026 11:38:37 +0200 Subject: [PATCH 3/3] Black ai4rag Signed-off-by: Jakub Walaszczyk --- ai4rag/core/experiment/experiment.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ai4rag/core/experiment/experiment.py b/ai4rag/core/experiment/experiment.py index 3a3e262..a07b967 100644 --- a/ai4rag/core/experiment/experiment.py +++ b/ai4rag/core/experiment/experiment.py @@ -379,7 +379,7 @@ def run_single_evaluation(self, rag_params: RAGParamsType) -> float: raise VectorStoreInitializationError( exc, embedding_model_id=embedding_model.model_id, - vector_store_provider_id=self.ogx_vector_io_provider_id or "local_chroma" + vector_store_provider_id=self.ogx_vector_io_provider_id or "local_chroma", ) from exc collection_name = vector_store.collection_name