Repository: Quansight/ragna
Pull request #350 (2024-03-11 13:31:36)

improve error handling for builtin assistants
I spend almost an hour today trying to find a bug that wasn't there, because we didn't raise a proper error when our API call failed.


Patch:
diff --git a/environment-dev.yml b/environment-dev.yml
index cb4bdd4..ca1daaf 100644
--- a/environment-dev.yml
+++ b/environment-dev.yml
@@ -9,6 +9,7 @@ dependencies:
       - python-dotenv
       - pytest >=6
       - pytest-mock
+      - pytest-asyncio
       - mypy ==1.6.1
       - pre-commit
       - types-aiofiles
diff --git a/pyproject.toml b/pyproject.toml
index 2a5c5bf..574813c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -97,16 +97,17 @@ ignore = ["E501"]
 
 [tool.pytest.ini_options]
 minversion = "6.0"
-addopts = "-ra --tb=short"
+addopts = "-ra --tb=short --asyncio-mode=auto"
 testpaths = [
     "tests",
 ]
 filterwarnings = [
     "error",
+    "ignore::ResourceWarning",
     # httpx 0.27.0 deprecated some functionality that the test client of starlette /
     # FastApi use. This should be resolved by the next release of these libraries.
     # See https://github.com/encode/starlette/issues/2524
-    "ignore:The 'app' shortcut is now deprecated:DeprecationWarning"
+    "ignore:The 'app' shortcut is now deprecated:DeprecationWarning",
 ]
 xfail_strict = true
 
diff --git a/ragna/assistants/_ai21labs.py b/ragna/assistants/_ai21labs.py
index 289840b..d6bf994 100644
--- a/ragna/assistants/_ai21labs.py
+++ b/ragna/assistants/_ai21labs.py
@@ -1,6 +1,6 @@
 from typing import AsyncIterator, cast
 
-from ragna.core import RagnaException, Source
+from ragna.core import Source
 
 from ._api import ApiAssistant
 
@@ -53,11 +53,7 @@ class Ai21LabsAssistant(ApiAssistant):
                 "system": self._make_system_content(sources),
             },
         )
-
-        if response.is_error:
-            raise RagnaException(
-                status_code=response.status_code, response=response.json()
-            )
+        await self._assert_api_call_is_success(response)
 
         yield cast(str, response.json()["outputs"][0]["text"])
 
diff --git a/ragna/assistants/_anthropic.py b/ragna/assistants/_anthropic.py
index 652bab7..02bc2f1 100644
--- a/ragna/assistants/_anthropic.py
+++ b/ragna/assistants/_anthropic.py
@@ -53,6 +53,8 @@ class AnthropicApiAssistant(ApiAssistant):
                 "stream": True,
             },
         ) as event_source:
+            await self._assert_api_call_is_success(event_source.response)
+
             async for sse in event_source.aiter_sse():
                 data = json.loads(sse.data)
                 if data["type"] != "completion":
diff --git a/ragna/assistants/_api.py b/ragna/assistants/_api.py
index cc03650..1a5bd3a 100644
--- a/ragna/assistants/_api.py
+++ b/ragna/assistants/_api.py
@@ -1,11 +1,14 @@
 import abc
+import contextlib
+import json
 import os
 from typing import AsyncIterator
 
 import httpx
+from httpx import Response
 
 import ragna
-from ragna.core import Assistant, EnvVarRequirement, Requirement, Source
+from ragna.core import Assistant, EnvVarRequirement, RagnaException, Requirement, Source
 
 
 class ApiAssistant(Assistant):
@@ -35,3 +38,19 @@ class ApiAssistant(Assistant):
         self, prompt: str, sources: list[Source], *, max_new_tokens: int
     ) -> AsyncIterator[str]:
         ...
+
+    async def _assert_api_call_is_success(self, response: Response) -> None:
+        if response.is_success:
+            return
+
+        content = await response.aread()
+        with contextlib.suppress(Exception):
+            content = json.loads(content)
+
+        raise RagnaException(
+            "API call failed",
+            request_method=response.request.method,
+            request_url=str(response.request.url),
+            response_status_code=response.status_code,
+            response_content=content,
+        )
diff --git a/ragna/assistants/_cohere.py b/ragna/assistants/_cohere.py
index 7e5c928..48516fd 100644
--- a/ragna/assistants/_cohere.py
+++ b/ragna/assistants/_cohere.py
@@ -54,12 +54,15 @@ class CohereApiAssistant(ApiAssistant):
                 "documents": self._make_source_documents(sources),
             },
         ) as response:
-            if response.is_error:
-                raise RagnaException(status_code=response.status_code)
+            await self._assert_api_call_is_success(response)
+
             async for chunk in response.aiter_lines():
                 event = json.loads(chunk)
                 if event["event_type"] == "stream-end":
-                    break
+                    if event["event_type"] == "COMPLETE":
+                        break
+
+                    raise RagnaException(event["error_message"])
                 if "text" in event:
                     yield cast(str, event["text"])
 
diff --git a/ragna/assistants/_google.py b/ragna/assistants/_google.py
index d27df04..4a91416 100644
--- a/ragna/assistants/_google.py
+++ b/ragna/assistants/_google.py
@@ -87,6 +87,8 @@ class GoogleApiAssistant(ApiAssistant):
                 },
             },
         ) as response:
+            await self._assert_api_call_is_success(response)
+
             async for chunk in ijson.items(
                 AsyncIteratorReader(response.aiter_bytes(1024)),
                 "item.candidates.item.content.parts.item.text",
diff --git a/ragna/assistants/_mosaicml.py b/ragna/assistants/_mosaicml.py
index 5234119..bfeb037 100644
--- a/ragna/assistants/_mosaicml.py
+++ b/ragna/assistants/_mosaicml.py
@@ -1,6 +1,6 @@
 from typing import AsyncIterator, cast
 
-from ragna.core import RagnaException, Source
+from ragna.core import Source
 
 from ._api import ApiAssistant
 
@@ -43,10 +43,8 @@ class MosaicmlApiAssistant(ApiAssistant):
                 "parameters": {"temperature": 0.0, "max_new_tokens": max_new_tokens},
             },
         )
-        if response.is_error:
-            raise RagnaException(
-                status_code=response.status_code, response=response.json()
-            )
+        await self._assert_api_call_is_success(response)
+
         yield cast(str, response.json()["outputs"][0]).replace(instruction, "").strip()
 
 
diff --git a/ragna/assistants/_openai.py b/ragna/assistants/_openai.py
index cfd4d5b..4f66b3d 100644
--- a/ragna/assistants/_openai.py
+++ b/ragna/assistants/_openai.py
@@ -60,6 +60,8 @@ class OpenaiApiAssistant(ApiAssistant):
                 "stream": True,
             },
         ) as event_source:
+            await self._assert_api_call_is_success(event_source.response)
+
             async for sse in event_source.aiter_sse():
                 data = json.loads(sse.data)
                 choice = data["choices"][0]
diff --git a/ragna/deploy/_ui/modal_configuration.py b/ragna/deploy/_ui/modal_configuration.py
index 7567c57..6821fa5 100644
--- a/ragna/deploy/_ui/modal_configuration.py
+++ b/ragna/deploy/_ui/modal_configuration.py
@@ -36,7 +36,7 @@ class ChatConfig(param.Parameterized):
     max_context_tokens = param.Integer(
         step=500,
         bounds=(1, 8000),
-        default=4000,
+        default=2000,
         doc=(
             "Maximum number of context tokens and in turn the number of document chunks "
             "pulled out of the vector database."


Test patch:
diff --git a/tests/assistants/__init__.py b/tests/assistants/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/assistants/test_api.py b/tests/assistants/test_api.py
new file mode 100644
index 0000000..9796145
--- /dev/null
+++ b/tests/assistants/test_api.py
@@ -0,0 +1,28 @@
+import os
+
+import pytest
+
+from ragna import assistants
+from ragna._compat import anext
+from ragna.assistants._api import ApiAssistant
+from ragna.core import RagnaException
+from tests.utils import skip_on_windows
+
+API_ASSISTANTS = [
+    assistant
+    for assistant in assistants.__dict__.values()
+    if isinstance(assistant, type)
+    and issubclass(assistant, ApiAssistant)
+    and assistant is not ApiAssistant
+]
+
+
+@skip_on_windows
+@pytest.mark.parametrize("assistant", API_ASSISTANTS)
+async def test_api_call_error_smoke(mocker, assistant):
+    mocker.patch.dict(os.environ, {assistant._API_KEY_ENV_VAR: "SENTINEL"})
+
+    chunks = assistant().answer(prompt="?", sources=[])
+
+    with pytest.raises(RagnaException, match="API call failed"):
+        await anext(chunks)
diff --git a/tests/utils.py b/tests/utils.py
new file mode 100644
index 0000000..356da1d
--- /dev/null
+++ b/tests/utils.py
@@ -0,0 +1,7 @@
+import platform
+
+import pytest
+
+skip_on_windows = pytest.mark.skipif(
+    platform.system() == "Windows", reason="Test is broken skipped on Windows"
+)