Repository: Quansight/ragna (python)

Issue:
Add flag to disable unavailable components for the API / UI
As is, our official docker image has a problem: we do not include any builtin assistants:

https://github.com/Quansight/ragna/blob/b3cc95ab5b4228b03af9937afc9dc842ccf25432/ragna-docker.toml#L5-L13

The reason for this is that if we include them by default, we would force the user to supply API keys for all of them or otherwise the API / UI would fail to start.

However, from a UX perspective this is quite bad. Imagine I have an `OPENAI_API_KEY` and I want to run ragna with docker. The only way I currently have short of building my own image with a new config, is to set the assistants through an environment variable and thus overwrite the config:

```bash
$ docker run \
    -p 31476:31476 -p 31477:31477 \
    -e OPENAI_API_KEY=$OPENAI_API_KEY \
    -e RAGNA_COMPONENTS_ASSISTANTS='["ragna.assistants.Gpt4"]' \
    ragna
```

(Before #331 is fixed, the env var needs to be `RAGNA_CORE_ASSISTANTS`)

I propose we introduce a `--ignore-unavailable-components` (happy for naming help here). Instead of failing if a component is not available

https://github.com/Quansight/ragna/blob/b3cc95ab5b4228b03af9937afc9dc842ccf25432/ragna/core/_rag.py#L53-L56

just ignores the component. Of course we should still fail if no component is available.

This way we can put all of our builtin components in the config and the user can "activate" them by supplying the corresponding API key. Meaning, they can run

```bash
$ docker run ragna check
```

to see what API keys they can supply and afterwards run something like

```bash
$ docker run \
    -p 31476:31476 -p 31477:31477 \
    -e OPENAI_API_KEY=$OPENAI_API_KEY \
    ragna
```

to start the UI with the OpenAI assistants without selecting them by hand. Note that this would also select the gpt-3.5 variant automatically, while in the workaround above only gpt-4 would be available.



Pull request description:
Add option to ignore unavailable components
Closes #332.

Patch:
diff --git a/Dockerfile b/Dockerfile
index be3a481..de26063 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -30,4 +30,4 @@ ENV LANCEDB_CONFIG_DIR=/var/ragna/lancedb.config
 COPY ragna-docker.toml ragna.toml
 
 ENTRYPOINT ["ragna"]
-CMD ["ui", "--start-api", "--no-open-browser"]
+CMD ["ui", "--start-api", "--ignore-unavailable-components", "--no-open-browser"]
diff --git a/ragna-docker.toml b/ragna-docker.toml
index ac576b0..1874fb3 100644
--- a/ragna-docker.toml
+++ b/ragna-docker.toml
@@ -7,7 +7,18 @@ source_storages = [
     "ragna.source_storages.LanceDB"
 ]
 assistants = [
-    "ragna.assistants.RagnaDemoAssistant"
+    "ragna.assistants.Jurassic2Ultra",
+    "ragna.assistants.Claude",
+    "ragna.assistants.ClaudeInstant",
+    "ragna.assistants.Command",
+    "ragna.assistants.CommandLight",
+    "ragna.assistants.RagnaDemoAssistant",
+    "ragna.assistants.GeminiPro",
+    "ragna.assistants.GeminiUltra",
+    "ragna.assistants.Mpt7bInstruct",
+    "ragna.assistants.Mpt30bInstruct",
+    "ragna.assistants.Gpt4",
+    "ragna.assistants.Gpt35Turbo16k",
 ]
 
 [api]
diff --git a/ragna/core/_rag.py b/ragna/core/_rag.py
index f63e603..6cdff12 100644
--- a/ragna/core/_rag.py
+++ b/ragna/core/_rag.py
@@ -44,7 +44,9 @@ class Rag(Generic[C]):
     def __init__(self) -> None:
         self._components: dict[Type[C], C] = {}
 
-    def _load_component(self, component: Union[Type[C], C]) -> C:
+    def _load_component(
+        self, component: Union[Type[C], C], *, ignore_unavailable: bool = False
+    ) -> Optional[C]:
         cls: Type[C]
         instance: Optional[C]
 
@@ -60,6 +62,9 @@ class Rag(Generic[C]):
         if cls not in self._components:
             if instance is None:
                 if not cls.is_available():
+                    if ignore_unavailable:
+                        return None
+
                     raise RagnaException(
                         "Component not available", name=cls.display_name()
                     )
@@ -148,8 +153,10 @@ class Chat:
         self._rag = rag
 
         self.documents = self._parse_documents(documents)
-        self.source_storage = self._rag._load_component(source_storage)
-        self.assistant = self._rag._load_component(assistant)
+        self.source_storage = cast(
+            SourceStorage, self._rag._load_component(source_storage)
+        )
+        self.assistant = cast(Assistant, self._rag._load_component(assistant))
 
         special_params = SpecialChatParams().model_dump()
         special_params.update(params)
diff --git a/ragna/deploy/_api/core.py b/ragna/deploy/_api/core.py
index c0fc4bd..2236c01 100644
--- a/ragna/deploy/_api/core.py
+++ b/ragna/deploy/_api/core.py
@@ -1,11 +1,19 @@
 import contextlib
-import itertools
 import uuid
 from typing import Annotated, Any, AsyncIterator, Iterator, Type, cast
 
 import aiofiles
 import sse_starlette
-from fastapi import Body, Depends, FastAPI, Form, HTTPException, Request, UploadFile
+from fastapi import (
+    Body,
+    Depends,
+    FastAPI,
+    Form,
+    HTTPException,
+    Request,
+    UploadFile,
+    status,
+)
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 
@@ -13,26 +21,48 @@ import ragna
 import ragna.core
 from ragna._compat import aiter, anext
 from ragna._utils import handle_localhost_origins
-from ragna.core import Component, Rag, RagnaException
+from ragna.core import Assistant, Component, Rag, RagnaException, SourceStorage
 from ragna.core._rag import SpecialChatParams
 from ragna.deploy import Config
 
 from . import database, schemas
 
 
-def app(config: Config) -> FastAPI:
+def app(*, config: Config, ignore_unavailable_components: bool) -> FastAPI:
     ragna.local_root(config.local_root)
 
     rag = Rag()  # type: ignore[var-annotated]
-    components_map: dict[str, Component] = {
-        component.display_name(): rag._load_component(component)
-        for component in itertools.chain(config.source_storages, config.assistants)
-    }
+    components_map: dict[str, Component] = {}
+    for components in [config.source_storages, config.assistants]:
+        components = cast(list[Type[Component]], components)
+        at_least_one = False
+        for component in components:
+            loaded_component = rag._load_component(
+                component, ignore_unavailable=ignore_unavailable_components
+            )
+            if loaded_component is None:
+                print(
+                    f"Ignoring {component.display_name()}, because it is not available."
+                )
+            else:
+                at_least_one = True
+                components_map[component.display_name()] = loaded_component
+
+        if not at_least_one:
+            raise RagnaException(
+                "No component available",
+                components=[component.display_name() for component in components],
+            )
 
     def get_component(display_name: str) -> Component:
         component = components_map.get(display_name)
         if component is None:
-            raise RagnaException("Unknown component", display_name=display_name)
+            raise RagnaException(
+                "Unknown component",
+                display_name=display_name,
+                http_status_code=status.HTTP_404_NOT_FOUND,
+                http_detail=RagnaException.MESSAGE,
+            )
 
         return component
 
@@ -97,11 +127,14 @@ def app(config: Config) -> FastAPI:
         return schemas.Components(
             documents=sorted(config.document.supported_suffixes()),
             source_storages=[
-                _get_component_json_schema(source_storage)
-                for source_storage in config.source_storages
+                _get_component_json_schema(type(source_storage))
+                for source_storage in components_map.values()
+                if isinstance(source_storage, SourceStorage)
             ],
             assistants=[
-                _get_component_json_schema(assistant) for assistant in config.assistants
+                _get_component_json_schema(type(assistant))
+                for assistant in components_map.values()
+                if isinstance(assistant, Assistant)
             ],
         )
 
diff --git a/ragna/deploy/_cli/core.py b/ragna/deploy/_cli/core.py
index b973678..fa493c2 100644
--- a/ragna/deploy/_cli/core.py
+++ b/ragna/deploy/_cli/core.py
@@ -78,8 +78,23 @@ def check(config: ConfigOption = "./ragna.toml") -> None:  # type: ignore[assign
 def api(
     *,
     config: ConfigOption = "./ragna.toml",  # type: ignore[assignment]
+    ignore_unavailable_components: Annotated[
+        bool,
+        typer.Option(
+            help=(
+                "Ignore components that are not available, "
+                "i.e. their requirements are not met. "
+            )
+        ),
+    ] = False,
 ) -> None:
-    uvicorn.run(api_app(config), host=config.api.hostname, port=config.api.port)
+    uvicorn.run(
+        api_app(
+            config=config, ignore_unavailable_components=ignore_unavailable_components
+        ),
+        host=config.api.hostname,
+        port=config.api.port,
+    )
 
 
 @app.command(help="Start the web UI.")
@@ -93,6 +108,16 @@ def ui(
             show_default="Start if the API is not served at the configured URL.",
         ),
     ] = None,
+    ignore_unavailable_components: Annotated[
+        bool,
+        typer.Option(
+            help=(
+                "Ignore components that are not available, "
+                "i.e. their requirements are not met. "
+                "This option as no effect if --no-start-api is used."
+            )
+        ),
+    ] = False,
     open_browser: Annotated[
         bool,
         typer.Option(help="Open the web UI in the browser when it is started."),
@@ -116,6 +141,7 @@ def ui(
                 "api",
                 "--config",
                 config.__ragna_cli_config_path__,  # type: ignore[attr-defined]
+                f"--{'' if ignore_unavailable_components else 'no-'}ignore-unavailable-components",
             ],
             stdout=sys.stdout,
             stderr=sys.stderr,
@@ -140,7 +166,7 @@ def ui(
                 )
                 raise typer.Exit(1)
 
-        ui_app(config, open_browser=open_browser).serve()  # type: ignore[no-untyped-call]
+        ui_app(config=config, open_browser=open_browser).serve()  # type: ignore[no-untyped-call]
     finally:
         if process is not None:
             process.kill()
diff --git a/ragna/deploy/_ui/app.py b/ragna/deploy/_ui/app.py
index 7590db8..c0bca73 100644
--- a/ragna/deploy/_ui/app.py
+++ b/ragna/deploy/_ui/app.py
@@ -133,7 +133,7 @@ class App(param.Parameterized):
         )
 
 
-def app(config: Config, *, open_browser: bool) -> App:
+def app(*, config: Config, open_browser: bool) -> App:
     return App(
         hostname=config.ui.hostname,
         port=config.ui.port,
diff --git a/scripts/docs/gen_files.py b/scripts/docs/gen_files.py
index b4cc78a..32a0eef 100644
--- a/scripts/docs/gen_files.py
+++ b/scripts/docs/gen_files.py
@@ -51,7 +51,7 @@ def cli_reference():
 
 
 def api_reference():
-    app = api_app(Config())
+    app = api_app(config=Config(), ignore_unavailable_components=False)
     openapi_json = fastapi.openapi.utils.get_openapi(
         title=app.title,
         version=app.version,


Test patch:
diff --git a/tests/deploy/api/test_components.py b/tests/deploy/api/test_components.py
new file mode 100644
index 0000000..505b0ca
--- /dev/null
+++ b/tests/deploy/api/test_components.py
@@ -0,0 +1,102 @@
+import pytest
+from fastapi import status
+from fastapi.testclient import TestClient
+
+from ragna import assistants
+from ragna.core import RagnaException
+from ragna.deploy import Config
+from ragna.deploy._api import app
+
+from .utils import authenticate
+
+
+@pytest.mark.parametrize("ignore_unavailable_components", [True, False])
+def test_ignore_unavailable_components(ignore_unavailable_components):
+    available_assistant = assistants.RagnaDemoAssistant
+    assert available_assistant.is_available()
+
+    unavailable_assistant = assistants.Gpt4
+    assert not unavailable_assistant.is_available()
+
+    config = Config(assistants=[available_assistant, unavailable_assistant])
+
+    if ignore_unavailable_components:
+        with TestClient(
+            app(
+                config=config,
+                ignore_unavailable_components=ignore_unavailable_components,
+            )
+        ) as client:
+            authenticate(client)
+
+            components = client.get("/components").raise_for_status().json()
+            assert [assistant["title"] for assistant in components["assistants"]] == [
+                available_assistant.display_name()
+            ]
+    else:
+        with pytest.raises(RagnaException, match="not available"):
+            app(
+                config=config,
+                ignore_unavailable_components=ignore_unavailable_components,
+            )
+
+
+def test_ignore_unavailable_components_at_least_one():
+    unavailable_assistant = assistants.Gpt4
+    assert not unavailable_assistant.is_available()
+
+    config = Config(assistants=[unavailable_assistant])
+
+    with pytest.raises(RagnaException, match="No component available"):
+        app(
+            config=config,
+            ignore_unavailable_components=True,
+        )
+
+
+def test_unknown_component(tmp_local_root):
+    config = Config(local_root=tmp_local_root)
+
+    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")
+
+    with TestClient(
+        app(config=Config(), ignore_unavailable_components=False)
+    ) as client:
+        authenticate(client)
+
+        document_upload = (
+            client.post("/document", json={"name": document_path.name})
+            .raise_for_status()
+            .json()
+        )
+        document = document_upload["document"]
+        assert document["name"] == document_path.name
+
+        parameters = document_upload["parameters"]
+        with open(document_path, "rb") as file:
+            client.request(
+                parameters["method"],
+                parameters["url"],
+                data=parameters["data"],
+                files={"file": file},
+            )
+
+        response = client.post(
+            "/chats",
+            json={
+                "name": "test-chat",
+                "source_storage": "unknown_source_storage",
+                "assistant": "unknown_assistant",
+                "params": {},
+                "documents": [document],
+            },
+        )
+
+        assert response.status_code == status.HTTP_404_NOT_FOUND
+
+        error = response.json()["error"]
+        assert "Unknown component" in error["message"]
diff --git a/tests/deploy/api/test_e2e.py b/tests/deploy/api/test_e2e.py
index b5c7e99..278c776 100644
--- a/tests/deploy/api/test_e2e.py
+++ b/tests/deploy/api/test_e2e.py
@@ -1,14 +1,14 @@
 import json
-import os
 
 import httpx_sse
 import pytest
 from fastapi.testclient import TestClient
 
-from ragna.core._utils import default_user
 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):
@@ -23,22 +23,8 @@ def check_api(config, *, stream_answer):
     with open(document_path, "w") as file:
         file.write("!\n")
 
-    with TestClient(app(config)) as client:
-        username = default_user()
-        token = (
-            client.post(
-                "/token",
-                data={
-                    "username": username,
-                    "password": os.environ.get(
-                        "RAGNA_DEMO_AUTHENTICATION_PASSWORD", username
-                    ),
-                },
-            )
-            .raise_for_status()
-            .json()
-        )
-        client.headers["Authorization"] = f"Bearer {token}"
+    with TestClient(app(config=config, ignore_unavailable_components=False)) as client:
+        authenticate(client)
 
         assert client.get("/chats").raise_for_status().json() == []
 
diff --git a/tests/deploy/api/utils.py b/tests/deploy/api/utils.py
new file mode 100644
index 0000000..abcf141
--- /dev/null
+++ b/tests/deploy/api/utils.py
@@ -0,0 +1,23 @@
+import os
+
+from fastapi.testclient import TestClient
+
+from ragna.core._utils import default_user
+
+
+def authenticate(client: TestClient) -> None:
+    username = default_user()
+    token = (
+        client.post(
+            "/token",
+            data={
+                "username": username,
+                "password": os.environ.get(
+                    "RAGNA_DEMO_AUTHENTICATION_PASSWORD", username
+                ),
+            },
+        )
+        .raise_for_status()
+        .json()
+    )
+    client.headers["Authorization"] = f"Bearer {token}"