Add test with uploaded document

Author: Julien Almarcha

app/endpoints/multiagents/__init__.py

@@ -137,7 +137,7 @@ async def multiagents(
         user_id=request.app.state.user.id,
     )
     initial_docs = [doc.chunk.content for doc in searches]
-    initial_refs = [doc.chunk.metadata.document_name for doc in searches]
+    initial_refs = [doc.chunk.metadata.get("document_name") for doc in searches]
 
     async def go_multiagents(body, initial_docs, initial_refs, n_retry, max_retry=5, window=5):
         docs_tmp = initial_docs[n_retry * window : (n_retry + 1) * window]
@@ -155,7 +155,7 @@ async def go_multiagents(body, initial_docs, initial_refs, n_retry, max_retry=5,
             return await go_multiagents(body, initial_docs, initial_refs, n_retry=n_retry + 1, max_retry=5, window=5)
         elif choice in [1, 2]:
             pass
-        elif choice == 4 or n_retry >= max_retry:  # else ?
+        elif choice == 4 or n_retry >= max_retry:
             searches = await context.documents.search(
                 session=session,
                 collection_ids=body.collections,
@@ -167,7 +167,9 @@ async def go_multiagents(body, initial_docs, initial_refs, n_retry, max_retry=5,
                 user_id=request.app.state.user.id,
             )
             docs_tmp = [doc.chunk.content for doc in searches]
-            refs_tmp = [doc.chunk.metadata.document_name for doc in searches]
+            refs_tmp = [doc.chunk.metadata.get("document_name") for doc in searches]
+        else:
+            raise ValueError(f"Unknown choice: {choice}")
         prompts = get_prompt_teller_multi(body.prompt, docs_tmp, choice)
         answers = await ask_in_parallel(model, prompts, user, body.max_tokens_intermediate)
         prompt = PROMPT_CONCAT.format(prompt=body.prompt, answers=answers)

app/tests/test_multiagents.py

@@ -1,8 +1,9 @@
+import os
 import pytest
 from uuid import uuid4
 from fastapi.testclient import TestClient
 from app.schemas.collections import CollectionVisibility
-from app.utils.variables import ENDPOINT__MULTIAGENTS, ENDPOINT__COLLECTIONS
+from app.utils.variables import ENDPOINT__FILES, ENDPOINT__MULTIAGENTS, ENDPOINT__COLLECTIONS
 
 
 @pytest.mark.usefixtures("client")
@@ -26,8 +27,6 @@ def test_multiagents_basic(self, client: TestClient):
             "k": 3,
             "rff_k": 1,
             "score_threshold": 0.5,
-            "writers_model": "writers_model_example",
-            "supervisor_model": "supervisor_model_example",
             "max_tokens": 50,
             "max_tokens_intermediate": 20,
             "model": "albert-small",
@@ -45,3 +44,49 @@ def test_multiagents_basic(self, client: TestClient):
         assert "n_retry" in data
         assert "sources_refs" in data
         assert "sources_content" in data
+
+    def test_multiagent_rag(self, client: TestClient):
+        """
+        Test the /multiagents endpoint after uploading a PDF file,
+        then calling the endpoint similarly to test_multiagents_basic
+        but using a different prompt.
+        """
+        # Create a private collection for the test
+        collection_name = f"test_collection_{str(uuid4())}"
+        params = {"name": collection_name, "visibility": CollectionVisibility.PRIVATE}
+        response = client.post_without_permissions(url=f"/v1{ENDPOINT__COLLECTIONS}", json=params)
+        assert response.status_code == 201, response.text
+        collection_id = response.json()["id"]
+
+        # Upload pdf.pdf into the new collection
+        file_path = "app/tests/assets/pdf.pdf"
+        with open(file_path, "rb") as file:
+            files = {"file": (os.path.basename(file_path), file, "application/pdf")}
+            data = {"request": '{"collection": "%s"}' % collection_id}
+            upload_response = client.post_without_permissions(url=f"/v1{ENDPOINT__FILES}", data=data, files=files)
+            file.close()
+        assert upload_response.status_code == 201, upload_response.text
+
+        # Now test the /multiagents endpoint with a different prompt
+        payload = {
+            "prompt": "Quel est le montant maximum des actes dont la signature de la première ministre peut être déléguée ?",
+            "collections": [collection_id],
+            "method": "semantic",
+            "k": 3,
+            "rff_k": 1,
+            "score_threshold": 0.5,
+            "max_tokens": 50,
+            "max_tokens_intermediate": 20,
+            "model": "albert-small",
+        }
+        response = client.post_with_permissions(f"/v1{ENDPOINT__MULTIAGENTS}", json=payload)
+        assert response.status_code == 200, response.text
+        data = response.json()
+
+        # Check response schema
+        assert "answer" in data
+        assert "choice" in data
+        assert "choice_desc" in data
+        assert "n_retry" in data
+        assert "sources_refs" in data
+        assert "sources_content" in data