make Chat.prepare idempotent (#482)

Quansight · Aug 7, 2024 · c98b262 · c98b262
1 parent 3dc586f
commit c98b262
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 22 deletions.
diff --git a/ragna/core/_rag.py b/ragna/core/_rag.py
@@ -181,26 +181,18 @@ async def prepare(self) -> Message:
 
         Returns:
             Welcome message.
-
-        Raises:
-            ragna.core.RagnaException: If chat is already
-                [`prepare`][ragna.core.Chat.prepare]d.
         """
-        if self._prepared or self.documents is None:
-            raise RagnaException(
-                "Chat is already prepared",
-                chat=self,
-                http_status_code=400,
-                http_detail=RagnaException.EVENT,
-            )
-
-        await self._run(self.source_storage.store, self.documents)
-        self._prepared = True
-
         welcome = Message(
             content="How can I help you with the documents?",
             role=MessageRole.SYSTEM,
         )
+
+        if self._prepared:
+            return welcome
+
+        await self._run(self.source_storage.store, self.documents)
+        self._prepared = True
+
         self._messages.append(welcome)
         return welcome
 

diff --git a/tests/core/test_e2e.py b/tests/core/test_e2e.py
@@ -11,16 +11,12 @@
 @pytest.mark.parametrize("input_type", ["corpus", "metadata_filter", "documents"])
 def test_e2e(tmp_local_root, input_type):
     async def main(*, input, source_storage, assistant):
-        chat = Rag().chat(
+        async with Rag().chat(
             input=input,
             source_storage=source_storage,
             assistant=assistant,
-        )
-
-        if not (input is None or isinstance(input, MetadataFilter)):
-            await chat.prepare()
-
-        return await chat.answer("?")
+        ) as chat:
+            return await chat.answer("?")
 
     document_root = tmp_local_root / "documents"
     document_root.mkdir()