Repository: Quansight/ragna
Pull request #495 (2024-08-20 20:39:05)

List corpuses
Addresses https://github.com/Quansight/ragna/issues/492



blakerosenthal: @pmeier Is this about right? I like the idea of raising `NotImplemented` on `source_storage.list_corpuses()` in case the user decides to ignore corpuses altogether.

Also right now the `/corpuses` endpoint returns the list of corpus names directly from the source_storages. I have two questions about this:
1. Should the api call distinguish which source storage the corpus name belongs to? E.g. `["LanceDB/<corpus_name>", "ChromaDB/<corpus_name>"]`? Maybe actually returning a JSON of `{"LanceDB": [<corpus_name>], "ChromaDB": [<corpus_name>]}` would be better?
2. For the built in source storages, the default corpus name is just the embedding ID. Maybe `list_corpuses()` should translate this to a prettier string like `"default"` for the sake of the UI.

Finally, should I add the actual UI changes as part of this issue, or will that be a separate issue?

pmeier: > I like the idea of raising `NotImplemented` on `source_storage.list_corpuses()` in case the user decides to ignore corpuses altogether.

On second thought, we should probably raise a descriptive `RagnaException` instead. We need to tell users what they have to do if they encounter this error and then exit gracefully.

> 1. Should the api call distinguish which source storage the corpus name belongs to? E.g. `["LanceDB/<corpus_name>", "ChromaDB/<corpus_name>"]`? Maybe actually returning a JSON of `{"LanceDB": [<corpus_name>], "ChromaDB": [<corpus_name>]}` would be better?

You are right, that needs to happen. I would go for the JSON approach where the key is the `.display_name()`.


> 2\. For the built in source storages, the default corpus name is just the embedding ID. Maybe `list_corpuses()` should translate this to a prettier string like `"default"` for the sake of the UI.

Yes, that needs to happen on the source storage itself. I'm also inclined to remove the feature of using the embedding ID as name for now. There is no use and it makes it more complicated. So feel free to just remove the embedding ID and just use `"default"`.
pmeier: > Finally, should I add the actual UI changes as part of this issue, or will that be a separate issue?

Up to you. When using two PRs we can merge this PR soon and don't have to wait for #484 to be merged.
blakerosenthal: > Up to you. When using two PRs we can merge this PR soon and don't have to wait for #484 to be merged.

Okay, let's keep them separate then

Patch:
diff --git a/ragna/core/_components.py b/ragna/core/_components.py
index 3b11845..502024f 100644
--- a/ragna/core/_components.py
+++ b/ragna/core/_components.py
@@ -17,10 +17,11 @@ from typing import (
 
 import pydantic
 import pydantic.utils
+from fastapi import status
 
 from ._document import Document
 from ._metadata_filter import MetadataFilter
-from ._utils import RequirementsMixin, merge_models
+from ._utils import RagnaException, RequirementsMixin, merge_models
 
 
 class Component(RequirementsMixin):
@@ -155,6 +156,19 @@ class SourceStorage(Component, abc.ABC):
         """
         ...
 
+    def list_corpuses(self) -> list[str]:
+        """List available corpuses.
+
+        Returns:
+            List of available corpuses.
+        """
+        raise RagnaException(
+            "list_corpuses is not implemented",
+            source_storage=self.__class__.display_name(),
+            http_status_code=status.HTTP_400_BAD_REQUEST,
+            http_detail=RagnaException.MESSAGE,
+        )
+
 
 class MessageRole(str, enum.Enum):
     """Message role
diff --git a/ragna/deploy/_api/core.py b/ragna/deploy/_api/core.py
index 5d40ef3..92f0ff2 100644
--- a/ragna/deploy/_api/core.py
+++ b/ragna/deploy/_api/core.py
@@ -1,6 +1,6 @@
 import contextlib
 import uuid
-from typing import Annotated, Any, AsyncIterator, Iterator, Type, Union, cast
+from typing import Annotated, Any, AsyncIterator, Iterator, Optional, Type, Union, cast
 
 import aiofiles
 from fastapi import (
@@ -145,6 +145,33 @@ def app(*, config: Config, ignore_unavailable_components: bool) -> FastAPI:
             ],
         )
 
+    @app.get("/corpuses")
+    async def get_corpuses(
+        _: UserDependency,
+        source_storage: Optional[str] = None,
+    ) -> dict[str, list[str]]:
+        if source_storage is not None:
+            component = components_map.get(source_storage)
+            if component is None or not isinstance(component, SourceStorage):
+                raise RagnaException(
+                    "Unknown source storage",
+                    display_name=source_storage,
+                    http_status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                    http_detail=RagnaException.MESSAGE,
+                )
+            source_storages = [component]
+        else:
+            source_storages = [
+                source_storage
+                for source_storage in components_map.values()
+                if isinstance(source_storage, SourceStorage)
+            ]
+
+        return {
+            source_storage.display_name(): source_storage.list_corpuses()
+            for source_storage in source_storages
+        }
+
     make_session = database.get_sessionmaker(config.api.database_url)
 
     @contextlib.contextmanager
diff --git a/ragna/source_storages/_chroma.py b/ragna/source_storages/_chroma.py
index b0ae180..32985e4 100644
--- a/ragna/source_storages/_chroma.py
+++ b/ragna/source_storages/_chroma.py
@@ -37,13 +37,13 @@ class Chroma(VectorDatabaseSourceStorage):
         )
 
     def _get_collection(self, corpus_name: str) -> chromadb.Collection:
-        if corpus_name == "default":
-            corpus_name = self._embedding_id
-
         return self._client.get_or_create_collection(
             corpus_name, embedding_function=self._embedding_function
         )
 
+    def list_corpuses(self) -> list[str]:
+        return [collection.name for collection in self._client.list_collections()]
+
     def store(
         self,
         corpus_name: str,
diff --git a/ragna/source_storages/_demo.py b/ragna/source_storages/_demo.py
index d8bc5a2..0d96f42 100644
--- a/ragna/source_storages/_demo.py
+++ b/ragna/source_storages/_demo.py
@@ -33,6 +33,9 @@ class RagnaDemoSourceStorage(SourceStorage):
     def __init__(self) -> None:
         self._storage: dict[str, list[dict[str, Any]]] = {}
 
+    def list_corpuses(self) -> list[str]:
+        return list(self._storage.keys())
+
     def store(self, corpus_name: str, documents: list[Document]) -> None:
         corpus = self._storage.setdefault(corpus_name, [])
         corpus.extend(
diff --git a/ragna/source_storages/_lancedb.py b/ragna/source_storages/_lancedb.py
index e1ecc6c..4f39ae3 100644
--- a/ragna/source_storages/_lancedb.py
+++ b/ragna/source_storages/_lancedb.py
@@ -53,9 +53,6 @@ class LanceDB(VectorDatabaseSourceStorage):
     _VECTOR_COLUMN_NAME = "embedded_text"
 
     def _get_table(self, corpus_name: str) -> lancedb.table.Table:
-        if corpus_name == "default":
-            corpus_name = self._embedding_id
-
         if corpus_name in self._db.table_names():
             return self._db.open_table(corpus_name)
         else:
@@ -190,6 +187,9 @@ class LanceDB(VectorDatabaseSourceStorage):
             )
             return f"{key} {operator} {value!r}"
 
+    def list_corpuses(self) -> list[str]:
+        return list(self._db.table_names())
+
     def retrieve(
         self,
         corpus_name: str,


Test patch:
diff --git a/tests/core/test_components.py b/tests/core/test_components.py
index 64fc398..324cbdf 100644
--- a/tests/core/test_components.py
+++ b/tests/core/test_components.py
@@ -3,7 +3,7 @@ import functools
 
 import pytest
 
-from ragna.core import Message
+from ragna.core import Message, RagnaException, SourceStorage
 
 
 def sync(async_test_fn):
@@ -81,3 +81,15 @@ class TestMessage:
         assert (await message.read()) == content
 
         assert message.content == content
+
+
+def test_method_not_implemented():
+    class TestSourceStorage(SourceStorage):
+        def store(self, **params):
+            pass
+
+        def retrieve(self, **params):
+            pass
+
+    with pytest.raises(RagnaException, match="not implemented"):
+        TestSourceStorage().list_corpuses()
diff --git a/tests/deploy/api/test_e2e.py b/tests/deploy/api/test_e2e.py
index f522e6f..74747c7 100644
--- a/tests/deploy/api/test_e2e.py
+++ b/tests/deploy/api/test_e2e.py
@@ -1,5 +1,6 @@
 import json
 
+import httpx
 import pytest
 from fastapi.testclient import TestClient
 
@@ -74,6 +75,9 @@ def test_e2e(tmp_local_root, multiple_answer_chunks, stream_answer, corpus_name)
         assert not chat["prepared"]
         assert chat["messages"] == []
 
+        corpuses = client.get("/corpuses").raise_for_status().json()
+        assert corpuses == {source_storage: []}
+
         assert client.get("/chats").raise_for_status().json() == [chat]
         assert client.get(f"/chats/{chat['id']}").raise_for_status().json() == chat
 
@@ -86,6 +90,21 @@ def test_e2e(tmp_local_root, multiple_answer_chunks, stream_answer, corpus_name)
         assert len(chat["messages"]) == 1
         assert chat["messages"][-1] == message
 
+        corpuses = client.get("/corpuses").raise_for_status().json()
+        assert corpuses == {source_storage: [corpus_name]}
+
+        corpuses = (
+            client.get("/corpuses", params={"source_storage": source_storage})
+            .raise_for_status()
+            .json()
+        )
+        assert corpuses == {source_storage: [corpus_name]}
+
+        with pytest.raises(httpx.HTTPStatusError, match="422 Unprocessable Entity"):
+            client.get(
+                "/corpuses", params={"source_storage": "unknown_source_storage"}
+            ).raise_for_status()
+
         prompt = "?"
         if stream_answer:
             with client.stream(
diff --git a/tests/source_storages/test_source_storages.py b/tests/source_storages/test_source_storages.py
index 91ed4e9..28b3497 100644
--- a/tests/source_storages/test_source_storages.py
+++ b/tests/source_storages/test_source_storages.py
@@ -155,9 +155,15 @@ def test_corpus_names(tmp_local_root, source_storage_cls):
 
     source_storage = source_storage_cls()
 
+    assert source_storage.list_corpuses() == []
     source_storage.store(secret_corpus_name, [secret_document])
     source_storage.store(dummy_corpus_name, [dummy_document])
 
+    assert set(source_storage.list_corpuses()) == {
+        dummy_corpus_name,
+        secret_corpus_name,
+    }
+
     prompt = "What is the secret number?"
     num_tokens = 4096