Repository: Quansight/ragna
Pull request #501 (2024-08-26 19:31:55)

add list_metadata functionality to web API
Addresses the API part of https://github.com/Quansight/ragna/issues/493.


Patch:
diff --git a/ragna/core/_components.py b/ragna/core/_components.py
index a720a2e..cf36f96 100644
--- a/ragna/core/_components.py
+++ b/ragna/core/_components.py
@@ -172,7 +172,7 @@ class SourceStorage(Component, abc.ABC):
 
     def list_metadata(
         self, corpus_name: Optional[str] = None
-    ) -> dict[str, dict[str, tuple[type, list[Any]]]]:
+    ) -> dict[str, dict[str, tuple[str, list[Any]]]]:
         """List available metadata for corpuses.
 
         Args:
diff --git a/ragna/deploy/_api/core.py b/ragna/deploy/_api/core.py
index 92f0ff2..8568823 100644
--- a/ragna/deploy/_api/core.py
+++ b/ragna/deploy/_api/core.py
@@ -145,11 +145,9 @@ 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]]:
+    def _get_source_storage_components(
+        source_storage: Optional[str],
+    ) -> list[SourceStorage]:
         if source_storage is not None:
             component = components_map.get(source_storage)
             if component is None or not isinstance(component, SourceStorage):
@@ -159,17 +157,32 @@ def app(*, config: Config, ignore_unavailable_components: bool) -> FastAPI:
                     http_status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                     http_detail=RagnaException.MESSAGE,
                 )
-            source_storages = [component]
+            return [component]
         else:
-            source_storages = [
+            return [
                 source_storage
                 for source_storage in components_map.values()
                 if isinstance(source_storage, SourceStorage)
             ]
 
+    @app.get("/corpuses")
+    async def get_corpuses(
+        _: UserDependency, source_storage: Optional[str] = None
+    ) -> dict[str, list[str]]:
         return {
             source_storage.display_name(): source_storage.list_corpuses()
-            for source_storage in source_storages
+            for source_storage in _get_source_storage_components(source_storage)
+        }
+
+    @app.get("/corpuses/metadata")
+    async def get_corpus_metadata(
+        _: UserDependency,
+        source_storage: Optional[str] = None,
+        corpus_name: Optional[str] = None,
+    ) -> dict[str, dict[str, dict[str, tuple[str, list[Any]]]]]:
+        return {
+            source_storage.display_name(): source_storage.list_metadata(corpus_name)
+            for source_storage in _get_source_storage_components(source_storage)
         }
 
     make_session = database.get_sessionmaker(config.api.database_url)
diff --git a/ragna/source_storages/_chroma.py b/ragna/source_storages/_chroma.py
index 9a9d95b..70f8a90 100644
--- a/ragna/source_storages/_chroma.py
+++ b/ragna/source_storages/_chroma.py
@@ -5,12 +5,7 @@ from collections import defaultdict
 from typing import TYPE_CHECKING, Any, Optional, cast
 
 import ragna
-from ragna.core import (
-    Document,
-    MetadataFilter,
-    MetadataOperator,
-    Source,
-)
+from ragna.core import Document, MetadataFilter, MetadataOperator, Source
 
 from ._utils import raise_no_corpuses_available, raise_non_existing_corpus
 from ._vector_database import VectorDatabaseSourceStorage
@@ -69,7 +64,7 @@ class Chroma(VectorDatabaseSourceStorage):
 
     def list_metadata(
         self, corpus_name: Optional[str] = None
-    ) -> dict[str, dict[str, tuple[type, list[Any]]]]:
+    ) -> dict[str, dict[str, tuple[str, list[Any]]]]:
         if corpus_name is None:
             corpus_names = self.list_corpuses()
         else:
@@ -91,7 +86,7 @@ class Chroma(VectorDatabaseSourceStorage):
                     corpus_metadata[key].add(value)
 
             metadata[corpus_name] = {
-                key: ({type(value) for value in values}.pop(), sorted(values))
+                key: ({type(value).__name__ for value in values}.pop(), sorted(values))
                 for key, values in corpus_metadata.items()
             }
 
diff --git a/ragna/source_storages/_demo.py b/ragna/source_storages/_demo.py
index 24d22dc..23f0435 100644
--- a/ragna/source_storages/_demo.py
+++ b/ragna/source_storages/_demo.py
@@ -54,7 +54,7 @@ class RagnaDemoSourceStorage(SourceStorage):
 
     def list_metadata(
         self, corpus_name: Optional[str] = None
-    ) -> dict[str, dict[str, tuple[type, list[Any]]]]:
+    ) -> dict[str, dict[str, tuple[str, list[Any]]]]:
         if corpus_name is None:
             corpus_names = self.list_corpuses()
         else:
@@ -73,7 +73,7 @@ class RagnaDemoSourceStorage(SourceStorage):
                     corpus_metadata[key].add(value)
 
             metadata[corpus_name] = {
-                key: ({type(value) for value in values}.pop(), sorted(values))
+                key: ({type(value).__name__ for value in values}.pop(), sorted(values))
                 for key, values in corpus_metadata.items()
             }
 
diff --git a/ragna/source_storages/_lancedb.py b/ragna/source_storages/_lancedb.py
index c09ca36..e59c1df 100644
--- a/ragna/source_storages/_lancedb.py
+++ b/ragna/source_storages/_lancedb.py
@@ -92,7 +92,7 @@ class LanceDB(VectorDatabaseSourceStorage):
 
     def list_metadata(
         self, corpus_name: Optional[str] = None
-    ) -> dict[str, dict[str, tuple[type, list[Any]]]]:
+    ) -> dict[str, dict[str, tuple[str, list[Any]]]]:
         if corpus_name is None:
             corpus_names = self.list_corpuses()
         else:
@@ -118,7 +118,7 @@ class LanceDB(VectorDatabaseSourceStorage):
 
             metadata[corpus_name] = {
                 key: (
-                    {type(value) for value in values}.pop(),
+                    {type(value).__name__ for value in values}.pop(),
                     sorted(values),
                 )
                 for key, values in corpus_metadata.items()


Test patch:
diff --git a/tests/deploy/api/test_e2e.py b/tests/deploy/api/test_e2e.py
index 74747c7..cf42b20 100644
--- a/tests/deploy/api/test_e2e.py
+++ b/tests/deploy/api/test_e2e.py
@@ -78,6 +78,9 @@ def test_e2e(tmp_local_root, multiple_answer_chunks, stream_answer, corpus_name)
         corpuses = client.get("/corpuses").raise_for_status().json()
         assert corpuses == {source_storage: []}
 
+        corpuses_metadata = client.get("/corpuses/metadata").raise_for_status().json()
+        assert corpuses_metadata == {source_storage: {}}
+
         assert client.get("/chats").raise_for_status().json() == [chat]
         assert client.get(f"/chats/{chat['id']}").raise_for_status().json() == chat
 
@@ -105,6 +108,33 @@ def test_e2e(tmp_local_root, multiple_answer_chunks, stream_answer, corpus_name)
                 "/corpuses", params={"source_storage": "unknown_source_storage"}
             ).raise_for_status()
 
+        corpuses_metadata = client.get("/corpuses/metadata").raise_for_status().json()
+        assert corpus_name in corpuses_metadata[source_storage]
+        metadata_keys = corpuses_metadata[source_storage][corpus_name].keys()
+        assert list(metadata_keys) == ["document_id", "document_name", "path"]
+        for key in metadata_keys:
+            assert corpuses_metadata[source_storage][corpus_name][key][0] == "str"
+
+        corpuses_metadata = (
+            client.get(
+                "/corpuses/metadata",
+                params={"source_storage": source_storage, corpus_name: corpus_name},
+            )
+            .raise_for_status()
+            .json()
+        )
+        assert corpus_name in corpuses_metadata[source_storage]
+        metadata_keys = corpuses_metadata[source_storage][corpus_name].keys()
+        assert list(metadata_keys) == ["document_id", "document_name", "path"]
+        for key in metadata_keys:
+            assert corpuses_metadata[source_storage][corpus_name][key][0] == "str"
+
+        with pytest.raises(httpx.HTTPStatusError, match="422 Unprocessable Entity"):
+            client.get(
+                "/corpuses/metadata",
+                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 92cee5f..54617c3 100644
--- a/tests/source_storages/test_source_storages.py
+++ b/tests/source_storages/test_source_storages.py
@@ -267,7 +267,7 @@ def test_list_metadata(tmp_local_root, cls):
                 corpus_metadata[key].add(value)
 
         expected_metadata[corpus_name] = {
-            key: ({type(value) for value in values}.pop(), values)
+            key: ({type(value).__name__ for value in values}.pop(), values)
             for key, values in corpus_metadata.items()
         }