Addresses comment about centralising raise_kibana_error

Author: mafaldasv-dev

src/elastic_evals/agent_builder/client.py

@@ -29,8 +29,8 @@
     UpdateAgentRequest,
     UpdateToolRequest,
 )
-from elastic_evals.api.response import parse_error_body
-from elastic_evals.api.retry import is_retryable_status_code, retry_kibana_api_call
+from elastic_evals.api.response import raise_kibana_error
+from elastic_evals.api.retry import retry_kibana_api_call
 from elastic_evals.utils.logging import log
 
 logger = log.getChild(__name__)
@@ -315,19 +315,12 @@ async def call_converse(
         )
 
     def _raise_error(self, response: httpx.Response, *, context: str) -> NoReturn:
-        status_code = response.status_code
-        body, body_text = parse_error_body(response)
-
-        message = f"Agent Builder request failed ({context}) with {status_code}"
-        if body_text:
-            message = f"{message}: {body_text}"
-
-        logger.error(message)
-        raise AgentBuilderError(
-            message=message,
-            status_code=status_code,
-            body=body,
-            retryable=is_retryable_status_code(status_code),
+        raise_kibana_error(
+            response,
+            error_cls=AgentBuilderError,
+            context="Agent Builder request",
+            operation=context,
+            logger=logger,
         )
 
     def _parse_success_body(self, response: httpx.Response, *, context: str) -> Any:

src/elastic_evals/api/datasets_client.py

@@ -7,7 +7,6 @@
 from __future__ import annotations
 
 import uuid
-from typing import NoReturn
 
 import httpx
 
@@ -23,8 +22,8 @@
 )
 from elastic_evals.api.errors import DatasetSyncError
 from elastic_evals.api.headers import build_kibana_headers
-from elastic_evals.api.response import parse_error_body
-from elastic_evals.api.retry import is_retryable_status_code, retry_kibana_api_call
+from elastic_evals.api.response import raise_kibana_error
+from elastic_evals.api.retry import retry_kibana_api_call
 
 
 def compute_dataset_id(name: str) -> str:
@@ -65,7 +64,11 @@ async def upsert(
         if response.status_code == 200:
             return UpsertDatasetResponse.model_validate(response.json())
 
-        self._raise_dataset_sync_error(response)
+        raise_kibana_error(
+            response,
+            error_cls=DatasetSyncError,
+            context="Kibana dataset sync request",
+        )
 
     @retry_kibana_api_call
     async def get(self, dataset_id: str) -> GetDatasetResponse:
@@ -85,19 +88,8 @@ async def get(self, dataset_id: str) -> GetDatasetResponse:
                 retryable=False,
             )
 
-        self._raise_dataset_sync_error(response)
-
-    def _raise_dataset_sync_error(self, response: httpx.Response) -> NoReturn:
-        status_code = response.status_code
-        body, body_text = parse_error_body(response)
-
-        message = f"Kibana dataset sync request failed with {status_code}"
-        if body_text:
-            message = f"{message}: {body_text}"
-
-        raise DatasetSyncError(
-            message=message,
-            status_code=status_code,
-            body=body,
-            retryable=is_retryable_status_code(status_code),
+        raise_kibana_error(
+            response,
+            error_cls=DatasetSyncError,
+            context="Kibana dataset sync request",
         )

src/elastic_evals/api/evaluators_client.py

@@ -27,8 +27,8 @@
     ValidateEvaluatorsResponse,
 )
 from elastic_evals.api.headers import build_kibana_headers
-from elastic_evals.api.response import parse_error_body
-from elastic_evals.api.retry import is_retryable_status_code, retry_kibana_api_call
+from elastic_evals.api.response import raise_kibana_error
+from elastic_evals.api.retry import retry_kibana_api_call
 
 
 class KibanaEvaluatorsClient:
@@ -95,15 +95,8 @@ async def _request(
         if response.status_code == 200:
             return response
 
-        status_code = response.status_code
-        body, body_text = parse_error_body(response)
-        message = f"Kibana evaluators request failed with {status_code}"
-        if body_text:
-            message = f"{message}: {body_text}"
-
-        raise KibanaEvaluatorsError(
-            message=message,
-            status_code=status_code,
-            body=body,
-            retryable=is_retryable_status_code(status_code),
+        raise_kibana_error(
+            response,
+            error_cls=KibanaEvaluatorsError,
+            context="Kibana evaluators request",
         )

src/elastic_evals/api/response.py

@@ -7,10 +7,14 @@
 from __future__ import annotations
 
 import json
-from typing import Any
+import logging
+from typing import Any, NoReturn
 
 import httpx
 
+from elastic_evals.api.errors import KibanaAPIError
+from elastic_evals.api.retry import is_retryable_status_code
+
 
 def parse_error_body(response: httpx.Response) -> tuple[Any, str]:
     """Parse an error body while preserving non-JSON response text."""
@@ -19,3 +23,29 @@ def parse_error_body(response: httpx.Response) -> tuple[Any, str]:
         return body, json.dumps(body, ensure_ascii=True)
     except (ValueError, json.JSONDecodeError):
         return response.text, response.text
+
+
+def raise_kibana_error(
+    response: httpx.Response,
+    *,
+    error_cls: type[KibanaAPIError],
+    context: str,
+    operation: str | None = None,
+    logger: logging.Logger | None = None,
+) -> NoReturn:
+    """Raise a typed Kibana API error from an unsuccessful response."""
+    body, body_text = parse_error_body(response)
+    message = f"{context} failed"
+    if operation is not None:
+        message = f"{message} ({operation})"
+    message = f"{message} with {response.status_code}"
+    if body_text:
+        message = f"{message}: {body_text}"
+    if logger is not None:
+        logger.error(message)
+    raise error_cls(
+        message=message,
+        status_code=response.status_code,
+        body=body,
+        retryable=is_retryable_status_code(response.status_code),
+    )

src/elastic_evals/api/scores_client.py

@@ -11,8 +11,8 @@
 from elastic_evals.api.constants import EVALS_SCORES_URL
 from elastic_evals.api.errors import IngestScoresError
 from elastic_evals.api.headers import build_kibana_headers
-from elastic_evals.api.response import parse_error_body
-from elastic_evals.api.retry import is_retryable_status_code, retry_kibana_api_call
+from elastic_evals.api.response import raise_kibana_error
+from elastic_evals.api.retry import retry_kibana_api_call
 from elastic_evals.api.scores_models import IngestScoresRequest, IngestScoresResponse
 from elastic_evals.utils.logging import log
 
@@ -54,16 +54,8 @@ async def ingest_scores(self, payload: IngestScoresRequest) -> IngestScoresRespo
                     )
             return parsed
 
-        status_code = response.status_code
-        body, body_text = parse_error_body(response)
-
-        message = f"Kibana score ingest request failed with {status_code}"
-        if body_text:
-            message = f"{message}: {body_text}"
-
-        raise IngestScoresError(
-            message=message,
-            status_code=status_code,
-            body=body,
-            retryable=is_retryable_status_code(status_code),
+        raise_kibana_error(
+            response,
+            error_cls=IngestScoresError,
+            context="Kibana score ingest request",
         )

tests/api/test_response.py

@@ -0,0 +1,58 @@
+# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+# or more contributor license agreements. Licensed under the Elastic License 2.0;
+# you may not use this file except in compliance with the Elastic License 2.0.
+
+import logging
+
+import httpx
+import pytest
+
+from elastic_evals.api.errors import KibanaEvaluatorsError
+from elastic_evals.api.response import raise_kibana_error
+
+
+def test_raise_kibana_error_parses_json_and_logs(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    logger = logging.getLogger(__name__)
+
+    with pytest.raises(KibanaEvaluatorsError) as exc_info:
+        raise_kibana_error(
+            httpx.Response(503, json={"message": "temporarily unavailable"}),
+            error_cls=KibanaEvaluatorsError,
+            context="Kibana evaluators request",
+            logger=logger,
+        )
+
+    assert exc_info.value.message == (
+        'Kibana evaluators request failed with 503: {"message": "temporarily unavailable"}'
+    )
+    assert exc_info.value.status_code == 503
+    assert exc_info.value.body == {"message": "temporarily unavailable"}
+    assert exc_info.value.retryable is True
+    assert exc_info.value.message in caplog.text
+
+
+def test_raise_kibana_error_preserves_text_body() -> None:
+    with pytest.raises(KibanaEvaluatorsError) as exc_info:
+        raise_kibana_error(
+            httpx.Response(400, text="invalid request"),
+            error_cls=KibanaEvaluatorsError,
+            context="Kibana evaluators request",
+        )
+
+    assert exc_info.value.message == "Kibana evaluators request failed with 400: invalid request"
+    assert exc_info.value.body == "invalid request"
+    assert exc_info.value.retryable is False
+
+
+def test_raise_kibana_error_includes_operation() -> None:
+    with pytest.raises(KibanaEvaluatorsError) as exc_info:
+        raise_kibana_error(
+            httpx.Response(400, text="invalid request"),
+            error_cls=KibanaEvaluatorsError,
+            context="Kibana evaluators request",
+            operation="evaluate",
+        )
+
+    assert exc_info.value.message == "Kibana evaluators request failed (evaluate) with 400: invalid request"