Repository: Quansight/ragna
Pull request #349 (2024-03-11 10:03:29)

add test for a streaming assistant
Previously, we indeed tested `stream=True`, but we used the demo assistant to do it which does not support streaming. Meaning, we only ever got a single chunk from the API.

Test for this was proposed in https://github.com/Quansight/ragna/pull/317#pullrequestreview-1886007856.


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/deploy/api/test_e2e.py b/tests/deploy/api/test_e2e.py
index 278c776..e19cb2c 100644
--- a/tests/deploy/api/test_e2e.py
+++ b/tests/deploy/api/test_e2e.py
@@ -4,25 +4,45 @@ import httpx_sse
 import pytest
 from fastapi.testclient import TestClient
 
+from ragna.assistants import RagnaDemoAssistant
 from ragna.deploy import Config
 from ragna.deploy._api import app
 
 from .utils import authenticate
 
 
-@pytest.mark.parametrize("stream_answer", [True, False])
-def test_e2e(tmp_local_root, stream_answer):
-    config = Config(local_root=tmp_local_root)
-    check_api(config, stream_answer=stream_answer)
+class TestAssistant(RagnaDemoAssistant):
+    @property
+    def max_input_size(self) -> int:
+        return 0
+
+    def answer(self, prompt, sources, *, multiple_answer_chunks: bool):
+        content = next(super().answer(prompt, sources))
+
+        if multiple_answer_chunks:
+            for chunk in content.split(" "):
+                yield f"{chunk} "
+        else:
+            yield content
 
 
-def check_api(config, *, stream_answer):
+@pytest.mark.parametrize("multiple_answer_chunks", [True, False])
+@pytest.mark.parametrize("stream_answer", [True, False])
+def test_e2e(tmp_local_root, multiple_answer_chunks, stream_answer):
+    config = Config(local_root=tmp_local_root, assistants=[TestAssistant])
+
     document_root = config.local_root / "documents"
     document_root.mkdir()
     document_path = document_root / "test.txt"
     with open(document_path, "w") as file:
         file.write("!\n")
 
+    # Reset starlette_sse AppStatus for each run
+    # See https://github.com/sysid/sse-starlette/issues/59
+    from sse_starlette.sse import AppStatus
+
+    AppStatus.should_exit_event = None
+
     with TestClient(app(config=config, ignore_unavailable_components=False)) as client:
         authenticate(client)
 
@@ -66,7 +86,7 @@ def check_api(config, *, stream_answer):
             "name": "test-chat",
             "source_storage": source_storage,
             "assistant": assistant,
-            "params": {},
+            "params": {"multiple_answer_chunks": multiple_answer_chunks},
             "documents": [document],
         }
         chat = client.post("/chats", json=chat_metadata).raise_for_status().json()
@@ -96,10 +116,10 @@ def check_api(config, *, stream_answer):
                 json={"prompt": prompt, "stream": True},
             ) as event_source:
                 for sse in event_source.iter_sse():
-                    chunk = json.loads(sse.data)
-                    chunks.append(chunk["content"])
-            message = chunk
-            message["content"] = "".join(chunks)
+                    chunks.append(json.loads(sse.data))
+            message = chunks[0]
+            assert all(chunk["sources"] is None for chunk in chunks[1:])
+            message["content"] = "".join(chunk["content"] for chunk in chunks)
         else:
             message = (
                 client.post(f"/chats/{chat['id']}/answer", json={"prompt": prompt})
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"
+)