Repository: Quansight/ragna
Pull request #491 (2024-08-14 21:47:13)

488 corpus cli
PoC that addresses #488.

Based on a [gist](https://gist.github.com/pmeier/cc1576839d3a25389080dc214640c397) produced by @pmeier, as well as comments from #488.

An example usage from the current branch is:
```
ragna corpus ingest $(find MY_DOCUMENT_DIRECTORY -type f) --metadata-fields metadata.json --corpus-name MY_CORPUS_NAME --user UNAME --report-failures 2>> failures.txt

cat metadata.json
# {"MY_DOCUMENT_DIRECTORY/MY_FILE1": {"test_field": true}, "MY_DOCUMENT_DIRECTORY/MY_FILE2": {"test_field": false}}
```  


Patch:
diff --git a/ragna/core/_components.py b/ragna/core/_components.py
index 502024f..cf36f96 100644
--- a/ragna/core/_components.py
+++ b/ragna/core/_components.py
@@ -6,6 +6,7 @@ import functools
 import inspect
 import uuid
 from typing import (
+    Any,
     AsyncIterable,
     AsyncIterator,
     Iterator,
@@ -169,6 +170,24 @@ class SourceStorage(Component, abc.ABC):
             http_detail=RagnaException.MESSAGE,
         )
 
+    def list_metadata(
+        self, corpus_name: Optional[str] = None
+    ) -> dict[str, dict[str, tuple[str, list[Any]]]]:
+        """List available metadata for corpuses.
+
+        Args:
+            corpus_name: Only return metadata for this corpus.
+
+        Returns:
+            List of available metadata.
+        """
+        raise RagnaException(
+            "list_metadata 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 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/deploy/_cli/core.py b/ragna/deploy/_cli/core.py
index cfafb74..91b2e75 100644
--- a/ragna/deploy/_cli/core.py
+++ b/ragna/deploy/_cli/core.py
@@ -16,6 +16,7 @@ from ragna.deploy._api import app as api_app
 from ragna.deploy._ui import app as ui_app
 
 from .config import ConfigOption, check_config, init_config
+from .corpus import app as corpus_app
 
 app = typer.Typer(
     name="ragna",
@@ -24,6 +25,7 @@ app = typer.Typer(
     add_completion=False,
     pretty_exceptions_enable=False,
 )
+app.add_typer(corpus_app)
 
 
 def version_callback(value: bool) -> None:
diff --git a/ragna/deploy/_cli/corpus.py b/ragna/deploy/_cli/corpus.py
new file mode 100644
index 0000000..f4244ec
--- /dev/null
+++ b/ragna/deploy/_cli/corpus.py
@@ -0,0 +1,219 @@
+import json
+import sys
+from pathlib import Path
+from typing import Annotated, Optional
+
+import rich
+import typer
+from rich.console import Console
+from rich.panel import Panel
+from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
+
+from ragna.core._utils import default_user
+from ragna.deploy._api import database, orm
+
+from .config import ConfigOption
+
+app = typer.Typer(
+    name="corpus",
+    help="(Experimental) Interact with a corpus of documents.",
+    invoke_without_command=True,
+    no_args_is_help=True,
+)
+
+
+@app.callback()
+def experimental_warning() -> None:
+    lines = [
+        (
+            "[bold]ragna corpus[/bold] and all subcommands are in an experimental "
+            "state and subject to change in the future."
+        ),
+        (
+            "If you have feedback or want to suggest a feature, "
+            "please open an issue at "
+            "https://github.com/Quansight/ragna/issues/new/choose."
+        ),
+    ]
+    rich.print(
+        Panel("\n".join(lines), title=":rotating_light: Warning :rotating_light:")
+    )
+
+
+@app.command(help="Ingest documents into a given corpus.")
+def ingest(
+    documents: list[Path],
+    metadata_fields: Annotated[
+        Optional[Path],
+        typer.Option(
+            help="JSON file that contains mappings from document name "
+            "to metadata fields associated with a document.",
+        ),
+    ] = None,
+    corpus_name: Annotated[
+        str,
+        typer.Option(help="Name of the corpus to ingest the documents into."),
+    ] = "default",
+    config: ConfigOption = "./ragna.toml",  # type: ignore[assignment]
+    user: Annotated[
+        Optional[str],
+        typer.Option(help="User to link the documents to in the ragna database."),
+    ] = None,
+    report_failures: Annotated[
+        bool,
+        typer.Option(help="Output to STDERR the documents that failed to be ingested."),
+    ] = False,
+    ignore_log: Annotated[
+        bool, typer.Option(help="Ignore the log file and re-ingest all documents.")
+    ] = False,
+) -> None:
+    try:
+        document_factory = getattr(config.document, "from_path")
+    except AttributeError:
+        raise typer.BadParameter(
+            f"{config.document.__name__} does not support creating documents from a"
+            f"path. Please implement a `from_path` method."
+        )
+
+    try:
+        make_session = database.get_sessionmaker(config.api.database_url)
+    except Exception:
+        raise typer.BadParameter(
+            f"Could not connect to the database: {config.api.database_url}"
+        )
+
+    if metadata_fields:
+        try:
+            with open(metadata_fields) as file:
+                metadata = json.load(file)
+        except Exception:
+            raise typer.BadParameter(
+                f"Could not read the metadata fields file: {metadata_fields}"
+            )
+    else:
+        metadata = {}
+
+    if user is None:
+        user = default_user()
+    with make_session() as session:  # type: ignore[attr-defined]
+        user_id = database._get_user_id(session, user)
+
+    # Log (JSONL) for recording which files previously added to vector database.
+    # Each entry has keys for 'user', 'corpus_name', 'source_storage' and 'document'.
+    ingestion_log: dict[str, set[str]] = {}
+    if not ignore_log:
+        ingestion_log_file = Path.cwd() / ".ragna_ingestion_log.jsonl"
+        if ingestion_log_file.exists():
+            with open(ingestion_log_file) as stream:
+                for line in stream:
+                    entry = json.loads(line)
+                    if entry["corpus_name"] == corpus_name and entry["user"] == user:
+                        ingestion_log.setdefault(entry["source_storage"], set()).add(
+                            entry["document"]
+                        )
+
+    with Progress(
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        "[progress.percentage]{task.percentage:>3.1f}%",
+        TimeRemainingColumn(),
+    ) as progress:
+        overall_task = progress.add_task(
+            "[cyan]Adding document embeddings to source storages...",
+            total=len(config.source_storages),
+        )
+
+        for source_storage in config.source_storages:
+            BATCH_SIZE = 10
+            number_of_batches = len(documents) // BATCH_SIZE
+            source_storage_task = progress.add_task(
+                f"[green]Adding document embeddings to {source_storage.display_name()}...",
+                total=number_of_batches,
+            )
+
+            for batch_number in range(0, len(documents), BATCH_SIZE):
+                documents_not_ingested = []
+                document_instances = []
+                orm_documents = []
+
+                if source_storage.display_name() in ingestion_log:
+                    batch_doc_set = set(
+                        [
+                            str(doc)
+                            for doc in documents[
+                                batch_number : batch_number + BATCH_SIZE
+                            ]
+                        ]
+                    )
+                    if batch_doc_set.issubset(
+                        ingestion_log[source_storage.display_name()]
+                    ):
+                        progress.advance(source_storage_task)
+                        continue
+
+                for document in documents[batch_number : batch_number + BATCH_SIZE]:
+                    try:
+                        doc_instance = document_factory(
+                            document,
+                            metadata=(
+                                metadata[str(document)]
+                                if str(document) in metadata
+                                else None
+                            ),
+                        )
+                        document_instances.append(doc_instance)
+                        orm_documents.append(
+                            orm.Document(
+                                id=doc_instance.id,
+                                user_id=user_id,
+                                name=doc_instance.name,
+                                metadata_=doc_instance.metadata,
+                            )
+                        )
+                    except Exception:
+                        documents_not_ingested.append(document)
+
+                if not orm_documents:
+                    continue
+
+                try:
+                    session = make_session()
+                    session.add_all(orm_documents)
+                    source_storage().store(corpus_name, document_instances)
+                    session.commit()
+                except Exception:
+                    documents_not_ingested.extend(
+                        documents[batch_number : batch_number + BATCH_SIZE]
+                    )
+                    session.rollback()
+                finally:
+                    session.close()
+
+                if not ignore_log:
+                    with open(ingestion_log_file, "a") as stream:
+                        for document in documents[
+                            batch_number : batch_number + BATCH_SIZE
+                        ]:
+                            stream.write(
+                                json.dumps(
+                                    {
+                                        "user": user,
+                                        "corpus_name": corpus_name,
+                                        "source_storage": source_storage.display_name(),
+                                        "document": str(document),
+                                    }
+                                )
+                                + "\n"
+                            )
+
+                if report_failures:
+                    Console(file=sys.stderr).print(
+                        f"{source_storage.__name__} failed to embed:\n{documents_not_ingested}",
+                    )
+
+                progress.advance(source_storage_task)
+
+            progress.update(source_storage_task, completed=number_of_batches)
+            progress.advance(overall_task)
+
+        progress.update(overall_task, completed=len(config.source_storages))
diff --git a/ragna/source_storages/_chroma.py b/ragna/source_storages/_chroma.py
index 32985e4..70f8a90 100644
--- a/ragna/source_storages/_chroma.py
+++ b/ragna/source_storages/_chroma.py
@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 import uuid
+from collections import defaultdict
 from typing import TYPE_CHECKING, Any, Optional, cast
 
 import ragna
 from ragna.core import Document, MetadataFilter, MetadataOperator, Source
 
+from ._utils import raise_no_corpuses_available, raise_non_existing_corpus
 from ._vector_database import VectorDatabaseSourceStorage
 
 if TYPE_CHECKING:
@@ -36,14 +38,60 @@ class Chroma(VectorDatabaseSourceStorage):
             )
         )
 
-    def _get_collection(self, corpus_name: str) -> chromadb.Collection:
-        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 _get_collection(
+        self, corpus_name: str, *, create: bool = False
+    ) -> chromadb.Collection:
+        if create:
+            return self._client.get_or_create_collection(
+                corpus_name, embedding_function=self._embedding_function
+            )
+
+        collections = list(self._client.list_collections())
+        if not collections:
+            raise_no_corpuses_available(self)
+
+        try:
+            return next(
+                collection
+                for collection in collections
+                if collection.name == corpus_name
+            )
+        except StopIteration:
+            raise_non_existing_corpus(self, corpus_name)
+
+    def list_metadata(
+        self, corpus_name: Optional[str] = None
+    ) -> dict[str, dict[str, tuple[str, list[Any]]]]:
+        if corpus_name is None:
+            corpus_names = self.list_corpuses()
+        else:
+            corpus_names = [corpus_name]
+
+        metadata = {}
+        for corpus_name in corpus_names:
+            collection = self._get_collection(corpus_name)
+
+            corpus_metadata = defaultdict(set)
+            for row in cast(
+                dict[str, list[Any]],
+                collection.get(include=["metadatas"]),  # type: ignore[list-item]
+            )["metadatas"]:
+                for key, value in row.items():
+                    if (key.startswith("__") and key.endswith("__")) or value is None:
+                        continue
+
+                    corpus_metadata[key].add(value)
+
+            metadata[corpus_name] = {
+                key: ({type(value).__name__ for value in values}.pop(), sorted(values))
+                for key, values in corpus_metadata.items()
+            }
+
+        return metadata
+
     def store(
         self,
         corpus_name: str,
@@ -52,7 +100,7 @@ class Chroma(VectorDatabaseSourceStorage):
         chunk_size: int = 500,
         chunk_overlap: int = 250,
     ) -> None:
-        collection = self._get_collection(corpus_name=corpus_name)
+        collection = self._get_collection(corpus_name=corpus_name, create=True)
 
         ids = []
         texts = []
@@ -70,8 +118,10 @@ class Chroma(VectorDatabaseSourceStorage):
                         "document_id": str(document.id),
                         "document_name": document.name,
                         **document.metadata,
-                        "page_numbers": self._page_numbers_to_str(chunk.page_numbers),
-                        "num_tokens": chunk.num_tokens,
+                        "__page_numbers__": self._page_numbers_to_str(
+                            chunk.page_numbers
+                        ),
+                        "__num_tokens__": chunk.num_tokens,
                     }
                 )
 
@@ -124,7 +174,7 @@ class Chroma(VectorDatabaseSourceStorage):
     def retrieve(
         self,
         corpus_name: str,
-        metadata_filter: MetadataFilter,
+        metadata_filter: Optional[MetadataFilter],
         prompt: str,
         *,
         chunk_size: int = 500,
@@ -177,9 +227,9 @@ class Chroma(VectorDatabaseSourceStorage):
                     id=result["ids"],
                     document_name=result["metadatas"]["document_name"],
                     document_id=result["metadatas"]["document_id"],
-                    location=result["metadatas"]["page_numbers"],
+                    location=result["metadatas"]["__page_numbers__"],
                     content=result["documents"],
-                    num_tokens=result["metadatas"]["num_tokens"],
+                    num_tokens=result["metadatas"]["__num_tokens__"],
                 )
                 for result in results
             ),
diff --git a/ragna/source_storages/_demo.py b/ragna/source_storages/_demo.py
index 0d96f42..23f0435 100644
--- a/ragna/source_storages/_demo.py
+++ b/ragna/source_storages/_demo.py
@@ -1,10 +1,9 @@
 import functools
 import textwrap
 import uuid
+from collections import defaultdict
 from typing import Any, Callable, Optional, cast
 
-from fastapi import status
-
 from ragna.core import (
     Document,
     MetadataFilter,
@@ -14,6 +13,8 @@ from ragna.core import (
     SourceStorage,
 )
 
+from ._utils import raise_no_corpuses_available, raise_non_existing_corpus
+
 
 class RagnaDemoSourceStorage(SourceStorage):
     """Demo assistant without requirements.
@@ -36,8 +37,52 @@ class RagnaDemoSourceStorage(SourceStorage):
     def list_corpuses(self) -> list[str]:
         return list(self._storage.keys())
 
+    def _get_corpus(
+        self, corpus_name: str, *, create: bool = False
+    ) -> list[dict[str, Any]]:
+        if create:
+            return self._storage.setdefault(corpus_name, [])
+
+        if not self._storage:
+            raise_no_corpuses_available(self)
+
+        corpus = self._storage.get(corpus_name)
+        if corpus is None:
+            raise_non_existing_corpus(self, corpus_name)
+
+        return corpus
+
+    def list_metadata(
+        self, corpus_name: Optional[str] = None
+    ) -> dict[str, dict[str, tuple[str, list[Any]]]]:
+        if corpus_name is None:
+            corpus_names = self.list_corpuses()
+        else:
+            corpus_names = [corpus_name]
+
+        metadata = {}
+        for corpus_name in corpus_names:
+            corpus = self._get_corpus(corpus_name)
+            corpus_metadata = defaultdict(set)
+
+            for row in corpus:
+                for key, value in row.items():
+                    if (key.startswith("__") and key.endswith("__")) or value is None:
+                        continue
+
+                    corpus_metadata[key].add(value)
+
+            metadata[corpus_name] = {
+                key: ({type(value).__name__ for value in values}.pop(), sorted(values))
+                for key, values in corpus_metadata.items()
+            }
+
+        return metadata
+
     def store(self, corpus_name: str, documents: list[Document]) -> None:
-        corpus = self._storage.setdefault(corpus_name, [])
+        corpus = self._get_corpus(corpus_name, create=True)
+        # FIXME: handle updating metadata (either introducing new or filling with None)
+        #  and add a type check
         corpus.extend(
             [
                 dict(
@@ -114,14 +159,6 @@ class RagnaDemoSourceStorage(SourceStorage):
     def retrieve(
         self, corpus_name: str, metadata_filter: MetadataFilter, prompt: str
     ) -> list[Source]:
-        corpus = self._storage.get(corpus_name)
-        if corpus is None:
-            raise RagnaException(
-                "Corpus does not exist",
-                corpus_name=corpus_name,
-                http_status_code=status.HTTP_400_BAD_REQUEST,
-                http_detail=RagnaException.MESSAGE,
-            )
         return [
             Source(
                 id=row["__id__"],
@@ -131,5 +168,7 @@ class RagnaDemoSourceStorage(SourceStorage):
                 content=row["__content__"],
                 num_tokens=row["__num_tokens__"],
             )
-            for _, row in self._apply_filter(corpus, metadata_filter)
+            for _, row in self._apply_filter(
+                self._get_corpus(corpus_name), metadata_filter
+            )
         ]
diff --git a/ragna/source_storages/_lancedb.py b/ragna/source_storages/_lancedb.py
index 4f39ae3..e59c1df 100644
--- a/ragna/source_storages/_lancedb.py
+++ b/ragna/source_storages/_lancedb.py
@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import uuid
 from collections import defaultdict
-from typing import TYPE_CHECKING, Optional, cast
+from typing import TYPE_CHECKING, Any, Optional, cast
 
 import ragna
 from ragna.core import (
@@ -15,6 +15,7 @@ from ragna.core import (
     Source,
 )
 
+from ._utils import raise_no_corpuses_available, raise_non_existing_corpus
 from ._vector_database import VectorDatabaseSourceStorage
 
 if TYPE_CHECKING:
@@ -50,31 +51,80 @@ class LanceDB(VectorDatabaseSourceStorage):
 
         self._db = lancedb.connect(ragna.local_root() / "lancedb")
 
-    _VECTOR_COLUMN_NAME = "embedded_text"
+    def list_corpuses(self) -> list[str]:
+        return list(self._db.table_names())
 
-    def _get_table(self, corpus_name: str) -> lancedb.table.Table:
-        if corpus_name in self._db.table_names():
-            return self._db.open_table(corpus_name)
-        else:
+    _VECTOR_COLUMN_NAME = "__embedded_text__"
+
+    def _get_table(
+        self, corpus_name: str, *, create: bool = False
+    ) -> lancedb.table.Table:
+        table_names = list(self._db.table_names())
+        no_corpuses = not table_names
+        non_existing_corpus = corpus_name not in table_names
+
+        if (no_corpuses or non_existing_corpus) and create:
             import pyarrow as pa
 
             return self._db.create_table(
                 name=corpus_name,
                 schema=pa.schema(
                     [
-                        pa.field("id", pa.string()),
+                        pa.field("__id__", pa.string()),
                         pa.field("document_id", pa.string()),
                         pa.field("document_name", pa.string()),
-                        pa.field("page_numbers", pa.string()),
-                        pa.field("text", pa.string()),
+                        pa.field("__page_numbers__", pa.string()),
+                        pa.field("__text__", pa.string()),
                         pa.field(
                             self._VECTOR_COLUMN_NAME,
                             pa.list_(pa.float32(), self._embedding_dimensions),
                         ),
-                        pa.field("num_tokens", pa.int32()),
+                        pa.field("__num_tokens__", pa.int32()),
                     ]
                 ),
             )
+        elif no_corpuses:
+            raise_no_corpuses_available(self)
+        elif non_existing_corpus:
+            raise_non_existing_corpus(self, corpus_name)
+        else:
+            return self._db.open_table(corpus_name)
+
+    def list_metadata(
+        self, corpus_name: Optional[str] = None
+    ) -> dict[str, dict[str, tuple[str, list[Any]]]]:
+        if corpus_name is None:
+            corpus_names = self.list_corpuses()
+        else:
+            corpus_names = [corpus_name]
+
+        metadata = {}
+        for corpus_name in corpus_names:
+            table = self._get_table(corpus_name)
+            corpus_metadata = (
+                table.to_arrow()
+                .select(
+                    [
+                        key
+                        for key in table.schema.names
+                        if not (key.startswith("__") and key.endswith("__"))
+                    ]
+                )
+                .to_pydict()
+            )
+            corpus_metadata = {
+                key: set(values) for key, values in corpus_metadata.items()
+            }
+
+            metadata[corpus_name] = {
+                key: (
+                    {type(value).__name__ for value in values}.pop(),
+                    sorted(values),
+                )
+                for key, values in corpus_metadata.items()
+            }
+
+        return metadata
 
     _PYTHON_TO_LANCE_TYPE_MAP = {
         bool: "boolean",
@@ -91,7 +141,7 @@ class LanceDB(VectorDatabaseSourceStorage):
         chunk_size: int = 500,
         chunk_overlap: int = 250,
     ) -> None:
-        table = self._get_table(corpus_name)
+        table = self._get_table(corpus_name, create=True)
 
         document_field_types = defaultdict(set)
         for document in documents:
@@ -130,13 +180,13 @@ class LanceDB(VectorDatabaseSourceStorage):
                     # overridden by concrete values if present
                     **default_metadata,
                     **document.metadata,
-                    "id": str(uuid.uuid4()),
+                    "__id__": str(uuid.uuid4()),
                     "document_id": str(document.id),
                     "document_name": str(document.name),
-                    "page_numbers": self._page_numbers_to_str(chunk.page_numbers),
-                    "text": chunk.text,
+                    "__page_numbers__": self._page_numbers_to_str(chunk.page_numbers),
+                    "__text__": chunk.text,
                     self._VECTOR_COLUMN_NAME: self._embedding_function([chunk.text])[0],
-                    "num_tokens": chunk.num_tokens,
+                    "__num_tokens__": chunk.num_tokens,
                 }
                 for document in documents
                 for chunk in self._chunk_pages(
@@ -187,9 +237,6 @@ 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,
@@ -221,15 +268,15 @@ class LanceDB(VectorDatabaseSourceStorage):
         return self._take_sources_up_to_max_tokens(
             (
                 Source(
-                    id=result["id"],
+                    id=result["__id__"],
                     document_id=result["document_id"],
                     document_name=result["document_name"],
                     # For some reason adding an empty string during store() results
                     # in this field being None. Thus, we need to parse it back here.
                     # TODO: See if there is a configuration option for this
-                    location=result["page_numbers"] or "",
-                    content=result["text"],
-                    num_tokens=result["num_tokens"],
+                    location=result["__page_numbers__"] or "",
+                    content=result["__text__"],
+                    num_tokens=result["__num_tokens__"],
                 )
                 for result in results.to_pylist()
             ),
diff --git a/ragna/source_storages/_utils.py b/ragna/source_storages/_utils.py
new file mode 100644
index 0000000..2c76753
--- /dev/null
+++ b/ragna/source_storages/_utils.py
@@ -0,0 +1,28 @@
+from typing import NoReturn
+
+from fastapi import status
+
+from ragna.core import RagnaException, SourceStorage
+
+
+def raise_no_corpuses_available(source_storage: SourceStorage) -> NoReturn:
+    __tracebackhide__ = True
+    raise RagnaException(
+        "No corpuses available",
+        source_storage=str(source_storage),
+        http_status_code=status.HTTP_400_BAD_REQUEST,
+        http_detail=RagnaException.MESSAGE,
+    ) from None
+
+
+def raise_non_existing_corpus(
+    source_storage: SourceStorage, corpus_name: str
+) -> NoReturn:
+    __tracebackhide__ = True
+    raise RagnaException(
+        "Corpus does not exist",
+        source_storage=str(source_storage),
+        corpus_name=corpus_name,
+        http_status_code=status.HTTP_400_BAD_REQUEST,
+        http_detail=RagnaException.MESSAGE,
+    ) from None


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 28b3497..54617c3 100644
--- a/tests/source_storages/test_source_storages.py
+++ b/tests/source_storages/test_source_storages.py
@@ -1,7 +1,18 @@
+import random
+import string
+from collections import defaultdict
+
 import pytest
 
-from ragna.core import LocalDocument, MetadataFilter, PlainTextDocumentHandler
-from ragna.source_storages import Chroma, LanceDB
+from ragna.core import (
+    LocalDocument,
+    MetadataFilter,
+    PlainTextDocumentHandler,
+    RagnaException,
+)
+from ragna.source_storages import Chroma, LanceDB, RagnaDemoSourceStorage
+
+SOURCE_STORAGES = [Chroma, LanceDB, RagnaDemoSourceStorage]
 
 METADATAS = {
     0: {"key": "value"},
@@ -182,3 +193,90 @@ def test_corpus_names(tmp_local_root, source_storage_cls):
         num_tokens=num_tokens,
     )
     assert secret not in dummy_sources[0].content
+
+
+@pytest.mark.parametrize("cls", SOURCE_STORAGES)
+def test_no_corpuses_error(tmp_local_root, cls):
+    source_storage = cls()
+
+    raises_error = pytest.raises(RagnaException, match="No corpuses available")
+
+    with raises_error:
+        source_storage.list_metadata(corpus_name="")
+
+    with raises_error:
+        source_storage.retrieve(corpus_name="", metadata_filter=None, prompt="")
+
+
+@pytest.mark.parametrize("cls", SOURCE_STORAGES)
+def test_non_existing_corpus_error(tmp_local_root, cls):
+    document_root = tmp_local_root / "documents"
+    document_root.mkdir()
+
+    document_path = document_root / "document.txt"
+    with open(document_path, "w") as file:
+        file.write("!")
+    document = LocalDocument.from_path(document_path)
+
+    source_storage = cls()
+    source_storage.store(corpus_name="default", documents=[document])
+
+    unknown_corpus_name = "unknown"
+    raises_error = pytest.raises(RagnaException, match="Corpus does not exist")
+
+    with raises_error:
+        source_storage.list_metadata(corpus_name=unknown_corpus_name)
+
+    with raises_error:
+        source_storage.retrieve(
+            corpus_name=unknown_corpus_name, metadata_filter=None, prompt=""
+        )
+
+
+@pytest.mark.parametrize("cls", SOURCE_STORAGES)
+def test_list_metadata(tmp_local_root, cls):
+    document_root = tmp_local_root / "documents"
+    document_root.mkdir()
+
+    documents = []
+    for idx in range(5):
+        path = (document_root / f"document{idx}").with_suffix(
+            random.choice([".txt", ".md"])
+        )
+        with open(path, "w") as file:
+            file.write(
+                "".join(random.choices(string.ascii_letters, k=random.randint(1, 10)))
+            )
+        documents.append(LocalDocument.from_path(path))
+
+    corpuses = {
+        "corpus0": documents[: len(documents) // 2],
+        "corpus1": documents[len(documents) // 2 :],
+    }
+
+    source_storage = cls()
+    expected_metadata = {}
+    for corpus_name, documents in corpuses.items():
+        source_storage.store(corpus_name, documents)
+
+        corpus_metadata = defaultdict(set)
+        for document in documents:
+            corpus_metadata["document_id"].add(str(document.id))
+            corpus_metadata["document_name"].add(document.name)
+            for key, value in document.metadata.items():
+                corpus_metadata[key].add(value)
+
+        expected_metadata[corpus_name] = {
+            key: ({type(value).__name__ for value in values}.pop(), values)
+            for key, values in corpus_metadata.items()
+        }
+
+    # We don't impose an order of the available values so we turn them into a set
+    actual_metadata = {
+        corpus_name: {
+            key: (type_, set(values)) for key, (type_, values) in metadata.items()
+        }
+        for corpus_name, metadata in source_storage.list_metadata().items()
+    }
+
+    assert actual_metadata == expected_metadata