Repository: Quansight/ragna (python)

Issue:
use metadata filter for chats
### Feature description

See #422 for work that has to happen before this. See #256 for a discussion.

- `SourceStorage.retrieve` needs to take a `MetadataFilter` object as first input argument rather than a list of documents
- The builtin source storages need to be refactored in a way to only have a single index. The metadata of each document should be stored alongside the document in the DB. During retrieval, the metadata filter needs to be translated into the dialect of the DB
- The first input argument of `Rag().chat()` needs to be renamed to `input` and should allow a list of `Document`s as well as a `MetadataFilter`:
  - Passing a list of documents should retain the current behavior by storing them in the source storage and creating a metadata filter to select only these documents in the retrieval stage.
  - Passing a metadata filter should skip the storing and just set the prepared flag on the chat
  The cases of passing both types or none as explored in https://github.com/Quansight/ragna/issues/256#issuecomment-1954057313 will be saved for later. For now only a single input is fine
- For the REST API, instead of overloading one input parameter as done for the Python API, we should have two mutually exclusive `documents` and `metadata_filter` inputs.

### Value and/or benefit

_No response_

### Anything else?

_No response_

Pull request:
Add metadata filter handling for builtin source storages
Closes #423. 

Patch:
diff --git a/docs/examples/gallery_streaming.py b/docs/examples/gallery_streaming.py
index 09a39ab..e67ba38 100644
--- a/docs/examples/gallery_streaming.py
+++ b/docs/examples/gallery_streaming.py
@@ -69,7 +69,7 @@ with open(document_path, "w") as file:
     file.write(ragna_docs.SAMPLE_CONTENT)
 
 chat = Rag().chat(
-    documents=[document_path],
+    input=[document_path],
     source_storage=source_storages.RagnaDemoSourceStorage,
     assistant=DemoStreamingAssistant,
 )
diff --git a/docs/tutorials/gallery_custom_components.py b/docs/tutorials/gallery_custom_components.py
index 8a411f8..4e7674d 100644
--- a/docs/tutorials/gallery_custom_components.py
+++ b/docs/tutorials/gallery_custom_components.py
@@ -121,7 +121,7 @@ with open(document_path, "w") as file:
 from ragna import Rag
 
 chat = Rag().chat(
-    documents=[document_path],
+    input=[document_path],
     source_storage=TutorialSourceStorage,
     assistant=TutorialAssistant,
 )
@@ -275,7 +275,7 @@ class ElaborateTutorialAssistant(Assistant):
 # creating a chat.
 
 chat = Rag().chat(
-    documents=[document_path],
+    input=[document_path],
     source_storage=TutorialSourceStorage,
     assistant=ElaborateTutorialAssistant,
     my_required_parameter=3,
@@ -291,7 +291,7 @@ print(await chat.answer("Hello!"))
 
 try:
     Rag().chat(
-        documents=[document_path],
+        input=[document_path],
         source_storage=TutorialSourceStorage,
         assistant=ElaborateTutorialAssistant,
     )
@@ -302,7 +302,7 @@ except Exception as exc:
 
 try:
     Rag().chat(
-        documents=[document_path],
+        input=[document_path],
         source_storage=TutorialSourceStorage,
         assistant=ElaborateTutorialAssistant,
         my_required_parameter="bar",
@@ -404,7 +404,7 @@ class AsyncAssistant(Assistant):
 
 
 chat = Rag().chat(
-    documents=[document_path],
+    input=[document_path],
     source_storage=TutorialSourceStorage,
     assistant=AsyncAssistant,
 )
diff --git a/docs/tutorials/gallery_python_api.py b/docs/tutorials/gallery_python_api.py
index 23667a7..b7d17dc 100644
--- a/docs/tutorials/gallery_python_api.py
+++ b/docs/tutorials/gallery_python_api.py
@@ -110,7 +110,7 @@ from ragna.assistants import RagnaDemoAssistant
 from ragna import Rag
 
 chat = Rag().chat(
-    documents=[document_path],
+    input=[document_path],
     source_storage=RagnaDemoSourceStorage,
     assistant=RagnaDemoAssistant,
 )
diff --git a/ragna/assistants/_demo.py b/ragna/assistants/_demo.py
index 997ca5a..f9bf644 100644
--- a/ragna/assistants/_demo.py
+++ b/ragna/assistants/_demo.py
@@ -42,7 +42,7 @@ class RagnaDemoAssistant(Assistant):
         prompt, sources = (message := messages[-1]).content, message.sources
         sources_display = []
         for source in sources:
-            source_display = f"- {source.document.name}"
+            source_display = f"- {source.document_name}"
             if source.location:
                 source_display += f", {source.location}"
             source_display += f": {textwrap.shorten(source.content, width=100)}"
diff --git a/ragna/core/_components.py b/ragna/core/_components.py
index bff4979..e592698 100644
--- a/ragna/core/_components.py
+++ b/ragna/core/_components.py
@@ -4,6 +4,7 @@ import abc
 import enum
 import functools
 import inspect
+import uuid
 from typing import (
     AsyncIterable,
     AsyncIterator,
@@ -18,6 +19,7 @@ import pydantic
 import pydantic.utils
 
 from ._document import Document
+from ._metadata_filter import MetadataFilter
 from ._utils import RequirementsMixin, merge_models
 
 
@@ -75,6 +77,7 @@ class Component(RequirementsMixin):
             annotations = get_type_hints(method)
             # Skip over the protocol parameters in order for the model below to only
             # comprise concrete parameters.
+
             for _ in range(num_protocol_params):
                 next(params)
 
@@ -112,14 +115,16 @@ class Source(pydantic.BaseModel):
         num_tokens: Number of tokens of the content.
     """
 
-    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
-
     id: str
-    document: Document
+    document_id: uuid.UUID
+    document_name: str
     location: str
     content: str
     num_tokens: int
 
+    def __hash__(self) -> int:
+        return hash(self.id)
+
 
 class SourceStorage(Component, abc.ABC):
     __ragna_protocol_methods__ = ["store", "retrieve"]
@@ -134,7 +139,7 @@ class SourceStorage(Component, abc.ABC):
         ...
 
     @abc.abstractmethod
-    def retrieve(self, documents: list[Document], prompt: str) -> list[Source]:
+    def retrieve(self, metadata_filter: MetadataFilter, prompt: str) -> list[Source]:
         """Retrieve sources for a given prompt.
 
         Args:
diff --git a/ragna/core/_rag.py b/ragna/core/_rag.py
index 15154ea..16b629d 100644
--- a/ragna/core/_rag.py
+++ b/ragna/core/_rag.py
@@ -27,6 +27,7 @@ from starlette.concurrency import iterate_in_threadpool, run_in_threadpool
 
 from ._components import Assistant, Component, Message, MessageRole, SourceStorage
 from ._document import Document, LocalDocument
+from ._metadata_filter import MetadataFilter
 from ._utils import RagnaException, default_user, merge_models
 
 T = TypeVar("T")
@@ -80,8 +81,8 @@ class Rag(Generic[C]):
 
     def chat(
         self,
+        input: Any,
         *,
-        documents: Iterable[Any],
         source_storage: Union[Type[SourceStorage], SourceStorage],
         assistant: Union[Type[Assistant], Assistant],
         **params: Any,
@@ -89,7 +90,7 @@ class Rag(Generic[C]):
         """Create a new [ragna.core.Chat][].
 
         Args:
-            documents: Documents to use. If any item is not a [ragna.core.Document][],
+            input: Documents to use. If any item is not a [ragna.core.Document][],
                 [ragna.core.LocalDocument.from_path][] is invoked on it.
             source_storage: Source storage to use.
             assistant: Assistant to use.
@@ -97,7 +98,7 @@ class Rag(Generic[C]):
         """
         return Chat(
             self,
-            documents=documents,
+            input=input,
             source_storage=source_storage,
             assistant=assistant,
             **params,
@@ -149,14 +150,15 @@ class Chat:
         self,
         rag: Rag,
         *,
-        documents: Iterable[Any],
+        input: Any,
         source_storage: Union[Type[SourceStorage], SourceStorage],
         assistant: Union[Type[Assistant], Assistant],
         **params: Any,
     ) -> None:
         self._rag = rag
 
-        self.documents = self._parse_documents(documents)
+        self.documents, self.metadata_filter, self._prepared = self._parse_input(input)
+
         self.source_storage = cast(
             SourceStorage, self._rag._load_component(source_storage)
         )
@@ -168,7 +170,6 @@ class Chat:
         self.params = params
         self._unpacked_params = self._unpack_chat_params(params)
 
-        self._prepared = False
         self._messages: list[Message] = []
 
     async def prepare(self) -> Message:
@@ -184,7 +185,7 @@ class Chat:
             ragna.core.RagnaException: If chat is already
                 [`prepare`][ragna.core.Chat.prepare]d.
         """
-        if self._prepared:
+        if self._prepared or self.documents is None:
             raise RagnaException(
                 "Chat is already prepared",
                 chat=self,
@@ -220,7 +221,9 @@ class Chat:
                 detail=RagnaException.EVENT,
             )
 
-        sources = await self._run(self.source_storage.retrieve, self.documents, prompt)
+        sources = await self._run(
+            self.source_storage.retrieve, self.metadata_filter, prompt
+        )
 
         question = Message(content=prompt, role=MessageRole.USER, sources=sources)
         self._messages.append(question)
@@ -237,9 +240,14 @@ class Chat:
 
         return answer
 
-    def _parse_documents(self, documents: Iterable[Any]) -> list[Document]:
-        documents_ = []
-        for document in documents:
+    def _parse_input(
+        self, input: Iterable[Any]
+    ) -> tuple[Optional[list[Document]], MetadataFilter, bool]:
+        if isinstance(input, MetadataFilter):
+            return None, input, True
+
+        documents = []
+        for document in input:
             if not isinstance(document, Document):
                 document = LocalDocument.from_path(document)
 
@@ -250,8 +258,15 @@ class Chat:
                     http_status_code=404,
                 )
 
-            documents_.append(document)
-        return documents_
+            documents.append(document)
+
+        metadata_filter = MetadataFilter.or_(
+            [
+                MetadataFilter.eq("document_id", str(document.id))
+                for document in documents
+            ]
+        )
+        return documents, metadata_filter, False
 
     def _unpack_chat_params(
         self, params: dict[str, Any]
diff --git a/ragna/deploy/_api/core.py b/ragna/deploy/_api/core.py
index 78ee615..3540f01 100644
--- a/ragna/deploy/_api/core.py
+++ b/ragna/deploy/_api/core.py
@@ -214,7 +214,7 @@ def app(*, config: Config, ignore_unavailable_components: bool) -> FastAPI:
         session: database.Session, *, user: str, chat: schemas.Chat
     ) -> ragna.core.Chat:
         core_chat = rag.chat(
-            documents=[
+            input=[
                 config.document(
                     id=document.id,
                     name=document.name,
diff --git a/ragna/source_storages/_chroma.py b/ragna/source_storages/_chroma.py
index 8082d4c..ad4651c 100644
--- a/ragna/source_storages/_chroma.py
+++ b/ragna/source_storages/_chroma.py
@@ -1,13 +1,16 @@
+from __future__ import annotations
+
 import uuid
+from typing import TYPE_CHECKING, Any, cast
 
 import ragna
-from ragna.core import (
-    Document,
-    Source,
-)
+from ragna.core import Document, MetadataFilter, MetadataOperator, Source
 
 from ._vector_database import VectorDatabaseSourceStorage
 
+if TYPE_CHECKING:
+    import chromadb
+
 
 class Chroma(VectorDatabaseSourceStorage):
     """[Chroma vector database](https://www.trychroma.com/)
@@ -33,17 +36,19 @@ class Chroma(VectorDatabaseSourceStorage):
             )
         )
 
+    def _get_collection(self) -> chromadb.Collection:
+        return self._client.get_or_create_collection(
+            self._embedding_id, embedding_function=self._embedding_function
+        )
+
     def store(
         self,
         documents: list[Document],
         *,
-        chat_id: uuid.UUID,
         chunk_size: int = 500,
         chunk_overlap: int = 250,
     ) -> None:
-        collection = self._client.create_collection(
-            str(chat_id), embedding_function=self._embedding_function
-        )
+        collection = self._get_collection()
 
         ids = []
         texts = []
@@ -59,6 +64,8 @@ class Chroma(VectorDatabaseSourceStorage):
                 metadatas.append(
                     {
                         "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,
                     }
@@ -70,22 +77,56 @@ class Chroma(VectorDatabaseSourceStorage):
             metadatas=metadatas,  # type: ignore[arg-type]
         )
 
+    # https://docs.trychroma.com/usage-guide#using-where-filters
+    _METADATA_OPERATOR_MAP = {
+        MetadataOperator.AND: "$and",
+        MetadataOperator.OR: "$or",
+        MetadataOperator.EQ: "$eq",
+        MetadataOperator.NE: "$ne",
+        MetadataOperator.LT: "$lt",
+        MetadataOperator.LE: "$lte",
+        MetadataOperator.GT: "$gt",
+        MetadataOperator.GE: "$gte",
+        MetadataOperator.IN: "$in",
+        MetadataOperator.NOT_IN: "$nin",
+    }
+
+    def _translate_metadata_filter(
+        self, metadata_filter: MetadataFilter
+    ) -> dict[str, Any]:
+        if metadata_filter.operator is MetadataOperator.RAW:
+            return cast(dict[str, Any], metadata_filter.value)
+        elif metadata_filter.operator in {MetadataOperator.AND, MetadataOperator.OR}:
+            return {
+                self._METADATA_OPERATOR_MAP[metadata_filter.operator]: [
+                    self._translate_metadata_filter(child)
+                    for child in metadata_filter.value
+                ]
+            }
+        else:
+            return {
+                metadata_filter.key: {
+                    self._METADATA_OPERATOR_MAP[
+                        metadata_filter.operator
+                    ]: metadata_filter.value
+                }
+            }
+
     def retrieve(
         self,
-        documents: list[Document],
+        metadata_filter: MetadataFilter,
         prompt: str,
         *,
         chat_id: uuid.UUID,
         chunk_size: int = 500,
         num_tokens: int = 1024,
     ) -> list[Source]:
-        collection = self._client.get_collection(
-            str(chat_id), embedding_function=self._embedding_function
-        )
+        collection = self._get_collection()
 
         include = ["distances", "metadatas", "documents"]
         result = collection.query(
             query_texts=prompt,
+            where=self._translate_metadata_filter(metadata_filter),
             n_results=min(
                 # We cannot retrieve source by a maximum number of tokens. Thus, we
                 # estimate how many sources we have to query. We overestimate by a
@@ -121,12 +162,13 @@ class Chroma(VectorDatabaseSourceStorage):
         #  2. Whatever threshold we use is very much dependent on the encoding method
         #  Thus, we likely need to have a callable parameter for this class
 
-        document_map = {str(document.id): document for document in documents}
         return self._take_sources_up_to_max_tokens(
             (
                 Source(
                     id=result["ids"],
-                    document=document_map[result["metadatas"]["document_id"]],
+                    # FIXME: We no longer have access to the document here
+                    # maybe reflect the same in the demo component
+                    document_id=result["metadatas"]["document_id"],
                     location=result["metadatas"]["page_numbers"],
                     content=result["documents"],
                     num_tokens=result["metadatas"]["num_tokens"],
diff --git a/ragna/source_storages/_demo.py b/ragna/source_storages/_demo.py
index 09095cc..6f80359 100644
--- a/ragna/source_storages/_demo.py
+++ b/ragna/source_storages/_demo.py
@@ -1,7 +1,16 @@
+import functools
 import textwrap
 import uuid
+from typing import Any
 
-from ragna.core import Document, Source, SourceStorage
+from ragna.core import (
+    Document,
+    MetadataFilter,
+    MetadataOperator,
+    RagnaException,
+    Source,
+    SourceStorage,
+)
 
 
 class RagnaDemoSourceStorage(SourceStorage):
@@ -20,23 +29,86 @@ class RagnaDemoSourceStorage(SourceStorage):
         return "Ragna/DemoSourceStorage"
 
     def __init__(self) -> None:
-        self._storage: dict[uuid.UUID, list[Source]] = {}
+        self._storage: list[dict[str, Any]] = []
 
-    def store(self, documents: list[Document], *, chat_id: uuid.UUID) -> None:
-        self._storage[chat_id] = [
+    def store(self, documents: list[Document]) -> None:
+        self._storage.extend(
+            [
+                dict(
+                    document_id=document.id,
+                    document_name=document.name,
+                    **document.metadata,
+                    __id__=str(uuid.uuid4()),
+                    __location__=(
+                        f"page {page.number}"
+                        if (page := next(document.extract_pages())).number
+                        else ""
+                    ),
+                    __content__=(content := textwrap.shorten(page.text, width=100)),
+                    __num_tokens__=len(content.split()),
+                )
+                for document in documents
+            ]
+        )
+
+    _METADATA_OPERATOR_MAP = {
+        MetadataOperator.EQ: lambda a, b: a == b,
+        MetadataOperator.NE: lambda a, b: a != b,
+        MetadataOperator.LT: lambda a, b: a < b,
+        MetadataOperator.LE: lambda a, b: a <= b,
+        MetadataOperator.GT: lambda a, b: a > b,
+        MetadataOperator.GE: lambda a, b: a >= b,
+        MetadataOperator.IN: lambda a, b: a in b,
+        MetadataOperator.NOT_IN: lambda a, b: a not in b,
+    }
+
+    def _apply_filter(
+        self, metadata_filter: MetadataFilter
+    ) -> list[tuple[int, dict[str, Any]]]:
+        if metadata_filter.operator is MetadataOperator.RAW:
+            raise RagnaException
+        elif metadata_filter.operator in {MetadataOperator.AND, MetadataOperator.OR}:
+            idcs_groups = []
+            rows_map = {}
+            for child in metadata_filter.value:
+                idcs_group = set()
+                for idx, row in self._apply_filter(child):
+                    idcs_group.add(idx)
+                    if idx not in rows_map:
+                        rows_map[idx] = row
+                idcs_groups.append(idcs_group)
+            idcs = functools.reduce(
+                (
+                    set.intersection
+                    if metadata_filter.operator is MetadataOperator.AND
+                    else set.union
+                ),
+                idcs_groups,
+            )
+            return [(idx, rows_map[idx]) for idx in sorted(idcs)]
+        else:
+            rows_with_idx = []
+            for idx, row in enumerate(self._storage):
+                value = row.get(metadata_filter.key)
+                if value is None:
+                    continue
+
+                if self._METADATA_OPERATOR_MAP[metadata_filter.operator](
+                    value, metadata_filter.value
+                ):
+                    rows_with_idx.append((idx, row))
+
+            return rows_with_idx
+
+    def retrieve(self, metadata_filter: MetadataFilter, prompt: str) -> list[Source]:
+        return [
             Source(
-                id=str(uuid.uuid4()),
-                document=document,
-                location=f"page {page.number}"
-                if (page := next(document.extract_pages())).number
-                else "",
-                content=(content := textwrap.shorten(page.text, width=100)),
-                num_tokens=len(content.split()),
+                id=row["__id__"],
+                document_id=row["document_id"],
+                document_name=row["document_name"],
+                location=row["__location__"],
+                content=row["__content__"],
+                num_tokens=row["__num_tokens__"],
             )
-            for document in documents
+            for _, row in self._apply_filter(metadata_filter)
         ]
-
-    def retrieve(
-        self, documents: list[Document], prompt: str, *, chat_id: uuid.UUID
-    ) -> list[Source]:
-        return self._storage[chat_id]
diff --git a/ragna/source_storages/_lancedb.py b/ragna/source_storages/_lancedb.py
index bd5003f..cba3801 100644
--- a/ragna/source_storages/_lancedb.py
+++ b/ragna/source_storages/_lancedb.py
@@ -1,7 +1,15 @@
 import uuid
+from typing import cast
 
 import ragna
-from ragna.core import Document, PackageRequirement, Requirement, Source
+from ragna.core import (
+    Document,
+    MetadataFilter,
+    MetadataOperator,
+    PackageRequirement,
+    Requirement,
+    Source,
+)
 
 from ._vector_database import VectorDatabaseSourceStorage
 
@@ -35,6 +43,7 @@ class LanceDB(VectorDatabaseSourceStorage):
         import pyarrow as pa
 
         self._db = lancedb.connect(ragna.local_root() / "lancedb")
+        # Create schema at runtime!
         self._schema = pa.schema(
             [
                 pa.field("id", pa.string()),
@@ -55,11 +64,10 @@ class LanceDB(VectorDatabaseSourceStorage):
         self,
         documents: list[Document],
         *,
-        chat_id: uuid.UUID,
         chunk_size: int = 500,
         chunk_overlap: int = 250,
     ) -> None:
-        table = self._db.create_table(name=str(chat_id), schema=self._schema)
+        table = self._db.create_table(name=self._embedding_id, schema=self._schema)
 
         for document in documents:
             for chunk in self._chunk_pages(
@@ -72,6 +80,7 @@ class LanceDB(VectorDatabaseSourceStorage):
                         {
                             "id": str(uuid.uuid4()),
                             "document_id": str(document.id),
+                            "document_name": str(document.name),
                             "page_numbers": self._page_numbers_to_str(
                                 chunk.page_numbers
                             ),
@@ -84,6 +93,41 @@ class LanceDB(VectorDatabaseSourceStorage):
                     ]
                 )
 
+    # https://lancedb.github.io/lancedb/sql/
+    _METADATA_OPERATOR_MAP = {
+        MetadataOperator.AND: "AND",
+        MetadataOperator.OR: "OR",
+        MetadataOperator.EQ: "=",
+        MetadataOperator.LT: "<",
+        MetadataOperator.LE: "<=",
+        MetadataOperator.GT: ">",
+        MetadataOperator.GE: ">=",
+        MetadataOperator.IN: "IN",
+    }
+
+    def _translate_metadata_filter(self, metadata_filter: MetadataFilter) -> str:
+        if metadata_filter.operator is MetadataOperator.RAW:
+            return cast(str, metadata_filter.value)
+        elif metadata_filter.operator in {
+            MetadataOperator.AND,
+            MetadataOperator.OR,
+        }:
+            return f" {self._METADATA_OPERATOR_MAP[metadata_filter.operator]} ".join(
+                f"({self._translate_metadata_filter(child)})"
+                for child in metadata_filter.value
+            )
+        elif metadata_filter.operator is MetadataOperator.NE:
+            return f"NOT ({self._translate_metadata_filter(MetadataFilter.eq(metadata_filter.key, metadata_filter.value))})"
+        elif metadata_filter.operator is MetadataOperator.NOT_IN:
+            return f"NOT ({self._translate_metadata_filter(MetadataFilter.in_(metadata_filter.key, metadata_filter.value))})"
+        else:
+            value = (
+                tuple(metadata_filter.value)
+                if metadata_filter.operator is MetadataOperator.IN
+                else metadata_filter.value
+            )
+            return f"{metadata_filter.key} {self._METADATA_OPERATOR_MAP[metadata_filter.operator]} {value!r}"
+
     def retrieve(
         self,
         documents: list[Document],
diff --git a/ragna/source_storages/_vector_database.py b/ragna/source_storages/_vector_database.py
index b44c16d..55dd5bd 100644
--- a/ragna/source_storages/_vector_database.py
+++ b/ragna/source_storages/_vector_database.py
@@ -1,4 +1,5 @@
 import dataclasses
+import hashlib
 import itertools
 from collections import deque
 from typing import (
@@ -64,14 +65,14 @@ class VectorDatabaseSourceStorage(SourceStorage):
         ]
 
     def __init__(self) -> None:
-        import chromadb.api
         import chromadb.utils.embedding_functions
         import tiktoken
 
-        self._embedding_function = cast(
-            chromadb.api.types.EmbeddingFunction,
-            chromadb.utils.embedding_functions.DefaultEmbeddingFunction(),
-        )
+        self._embedding_function = chromadb.utils.embedding_functions.ONNXMiniLM_L6_V2()
+        self._embedding_name = self._embedding_function.MODEL_NAME
+        self._embedding_id = hashlib.md5(
+            self._embedding_name.encode(), usedforsecurity=False
+        ).hexdigest()
         # https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2#all-minilm-l6-v2
         self._embedding_dimensions = 384
         self._tokenizer = tiktoken.get_encoding("cl100k_base")


Test patch:
diff --git a/tests/core/test_e2e.py b/tests/core/test_e2e.py
index 18a7a16..a7ad414 100644
--- a/tests/core/test_e2e.py
+++ b/tests/core/test_e2e.py
@@ -17,7 +17,7 @@ def test_e2e(tmp_local_root):
 
     async def main(*, documents, source_storage, assistant):
         async with Rag().chat(
-            documents=documents,
+            input=documents,
             source_storage=source_storage,
             assistant=assistant,
         ) as chat:
diff --git a/tests/core/test_rag.py b/tests/core/test_rag.py
index 34de28f..b1766a6 100644
--- a/tests/core/test_rag.py
+++ b/tests/core/test_rag.py
@@ -21,7 +21,7 @@ class TestChat:
         **params,
     ):
         return Rag().chat(
-            documents=documents,
+            input=documents,
             source_storage=source_storage,
             assistant=assistant,
             **params,