Make document upload method flexible (#238)

Author: pmeier

ragna/core/__init__.py

@@ -4,6 +4,7 @@
     "Component",
     "Document",
     "DocumentHandler",
+    "DocumentUploadParameters",
     "EnvVarRequirement",
     "LocalDocument",
     "Message",
@@ -31,6 +32,7 @@
 from ._document import (
     Document,
     DocumentHandler,
+    DocumentUploadParameters,
     LocalDocument,
     Page,
     PdfDocumentHandler,

ragna/core/_document.py

@@ -17,6 +17,12 @@
     from ragna.deploy import Config
 
 
+class DocumentUploadParameters(BaseModel):
+    method: str
+    url: str
+    data: dict
+
+
 class Document(RequirementsMixin, abc.ABC):
     """Abstract base class for all documents."""
 
@@ -59,7 +65,7 @@ def get_handler(name: str) -> DocumentHandler:
     @abc.abstractmethod
     async def get_upload_info(
         cls, *, config: Config, user: str, id: uuid.UUID, name: str
-    ) -> tuple[str, dict[str, Any], dict[str, Any]]:
+    ) -> tuple[dict[str, Any], DocumentUploadParameters]:
         pass
 
     @abc.abstractmethod
@@ -136,7 +142,7 @@ def read(self) -> bytes:
     @classmethod
     async def get_upload_info(
         cls, *, config: Config, user: str, id: uuid.UUID, name: str
-    ) -> tuple[str, dict[str, Any], dict[str, Any]]:
+    ) -> tuple[dict[str, Any], DocumentUploadParameters]:
         url = f"{config.api.url}/document"
         data = {
             "token": jwt.encode(
@@ -150,7 +156,7 @@ async def get_upload_info(
             )
         }
         metadata = {"path": str(config.local_cache_root / "documents" / str(id))}
-        return url, data, metadata
+        return metadata, DocumentUploadParameters(method="PUT", url=url, data=data)
 
     @classmethod
     def decode_upload_token(cls, token: str) -> tuple[str, uuid.UUID]:

ragna/deploy/_api/core.py

@@ -116,16 +116,16 @@ def get_session() -> Iterator[database.Session]:
     async def create_document_upload_info(
         user: UserDependency,
         name: Annotated[str, Body(..., embed=True)],
-    ) -> schemas.DocumentUploadInfo:
+    ) -> schemas.DocumentUpload:
         with get_session() as session:
             document = schemas.Document(name=name)
-            url, data, metadata = await config.document.get_upload_info(
+            metadata, parameters = await config.document.get_upload_info(
                 config=config, user=user, id=document.id, name=document.name
             )
             database.add_document(
                 session, user=user, document=document, metadata=metadata
             )
-            return schemas.DocumentUploadInfo(url=url, data=data, document=document)
+            return schemas.DocumentUpload(parameters=parameters, document=document)
 
     @app.put("/document")
     async def upload_document(

ragna/deploy/_api/schemas.py

@@ -27,9 +27,8 @@ def from_core(cls, document: ragna.core.Document) -> Document:
         )
 
 
-class DocumentUploadInfo(BaseModel):
-    url: str
-    data: dict
+class DocumentUpload(BaseModel):
+    parameters: ragna.core.DocumentUploadParameters
     document: Document
 
 

ragna/deploy/_ui/resources/upload.js

@@ -12,18 +12,19 @@ async function uploadFile(file, token, informationEndpoint) {
     headers: { "Content-Type": "application/json", Authorization: `Bearer ${token}` },
     body: JSON.stringify({ name: file.name }),
   });
-  const documentInfo = await response.json();
+  const documentUpload = await response.json();
 
+  const parameters = documentUpload.parameters;
   var body = new FormData();
-  for (const [key, value] of Object.entries(documentInfo.data)) {
+  for (const [key, value] of Object.entries(parameters.data)) {
     body.append(key, value);
   }
   body.append("file", file);
 
-  await fetch(documentInfo.url, {
-    method: "PUT",
+  await fetch(parameters.url, {
+    method: parameters.method,
     body: body,
   });
 
-  return documentInfo.document;
+  return documentUpload.document;
 }

scripts/add_chats.py

@@ -34,15 +34,17 @@ def main():
     documents = []
     for i in range(5):
         name = f"document{i}.txt"
-        document_info = (
+        document_upload = (
             client.post("/document", json={"name": name}).raise_for_status().json()
         )
-        client.put(
-            document_info["url"],
-            data=document_info["data"],
+        parameters = document_upload["parameters"]
+        client.request(
+            parameters["method"],
+            parameters["url"],
+            data=parameters["data"],
             files={"file": f"Content of {name}".encode()},
         ).raise_for_status()
-        documents.append(document_info["document"])
+        documents.append(document_upload["document"])
 
     ## chat 1
 

tests/deploy/api/test_e2e.py

@@ -49,20 +49,22 @@ def check_api(config):
 
         assert client.get("/chats").raise_for_status().json() == []
 
-        document_info = (
+        document_upload = (
             client.post("/document", json={"name": document_path.name})
             .raise_for_status()
             .json()
         )
-        document = document_info["document"]
+        document = document_upload["document"]
         assert document["name"] == document_path.name
 
+        parameters = document_upload["parameters"]
         with open(document_path, "rb") as file:
-            client.put(
-                document_info["url"],
-                data=document_info["data"],
+            client.request(
+                parameters["method"],
+                parameters["url"],
+                data=parameters["data"],
                 files={"file": file},
-            ).raise_for_status()
+            )
 
         components = client.get("/components").raise_for_status().json()
         documents = components["documents"]