Skip to content

Commit d5fb9b3

Browse files
committed
Make document upload method flexible (#238)
Conflicts: ragna/_api/core.py ragna/deploy/_api/core.py
1 parent f650cdd commit d5fb9b3

File tree

7 files changed

+37
-25
lines changed

7 files changed

+37
-25
lines changed

ragna/_api/core.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -100,16 +100,16 @@ def get_session() -> Iterator[database.Session]:
100100
async def create_document_upload_info(
101101
user: UserDependency,
102102
name: Annotated[str, Body(..., embed=True)],
103-
) -> schemas.DocumentUploadInfo:
103+
) -> schemas.DocumentUpload:
104104
with get_session() as session:
105105
document = schemas.Document(name=name)
106-
url, data, metadata = await config.core.document.get_upload_info(
106+
metadata, parameters = await config.core.document.get_upload_info(
107107
config=config, user=user, id=document.id, name=document.name
108108
)
109109
database.add_document(
110110
session, user=user, document=document, metadata=metadata
111111
)
112-
return schemas.DocumentUploadInfo(url=url, data=data, document=document)
112+
return schemas.DocumentUpload(parameters=parameters, document=document)
113113

114114
@app.put("/document")
115115
async def upload_document(

ragna/_api/schemas.py

+2-3
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,8 @@ def from_core(cls, document: ragna.core.Document) -> Document:
2727
)
2828

2929

30-
class DocumentUploadInfo(BaseModel):
31-
url: str
32-
data: dict
30+
class DocumentUpload(BaseModel):
31+
parameters: ragna.core.DocumentUploadParameters
3332
document: Document
3433

3534

ragna/_ui/resources/upload.js

+6-5
Original file line numberDiff line numberDiff line change
@@ -12,18 +12,19 @@ async function uploadFile(file, token, informationEndpoint) {
1212
headers: { "Content-Type": "application/json", Authorization: `Bearer ${token}` },
1313
body: JSON.stringify({ name: file.name }),
1414
});
15-
const documentInfo = await response.json();
15+
const documentUpload = await response.json();
1616

17+
const parameters = documentUpload.parameters;
1718
var body = new FormData();
18-
for (const [key, value] of Object.entries(documentInfo.data)) {
19+
for (const [key, value] of Object.entries(parameters.data)) {
1920
body.append(key, value);
2021
}
2122
body.append("file", file);
2223

23-
await fetch(documentInfo.url, {
24-
method: "PUT",
24+
await fetch(parameters.url, {
25+
method: parameters.method,
2526
body: body,
2627
});
2728

28-
return documentInfo.document;
29+
return documentUpload.document;
2930
}

ragna/core/__init__.py

+2
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
"Config",
77
"Document",
88
"DocumentHandler",
9+
"DocumentUploadParameters",
910
"EnvVarRequirement",
1011
"LocalDocument",
1112
"Message",
@@ -35,6 +36,7 @@
3536
from ._document import (
3637
Document,
3738
DocumentHandler,
39+
DocumentUploadParameters,
3840
LocalDocument,
3941
Page,
4042
PdfDocumentHandler,

ragna/core/_document.py

+9-3
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,12 @@
1717
from ._config import Config
1818

1919

20+
class DocumentUploadParameters(BaseModel):
21+
method: str
22+
url: str
23+
data: dict
24+
25+
2026
class Document(RequirementsMixin, abc.ABC):
2127
"""Abstract base class for all documents."""
2228

@@ -59,7 +65,7 @@ def get_handler(name: str) -> DocumentHandler:
5965
@abc.abstractmethod
6066
async def get_upload_info(
6167
cls, *, config: Config, user: str, id: uuid.UUID, name: str
62-
) -> tuple[str, dict[str, Any], dict[str, Any]]:
68+
) -> tuple[dict[str, Any], DocumentUploadParameters]:
6369
pass
6470

6571
@abc.abstractmethod
@@ -138,7 +144,7 @@ def read(self) -> bytes:
138144
@classmethod
139145
async def get_upload_info(
140146
cls, *, config: Config, user: str, id: uuid.UUID, name: str
141-
) -> tuple[str, dict[str, Any], dict[str, Any]]:
147+
) -> tuple[dict[str, Any], DocumentUploadParameters]:
142148
url = f"{config.api.url}/document"
143149
data = {
144150
"token": jwt.encode(
@@ -152,7 +158,7 @@ async def get_upload_info(
152158
)
153159
}
154160
metadata = {"path": str(config.local_cache_root / "documents" / str(id))}
155-
return url, data, metadata
161+
return metadata, DocumentUploadParameters(method="PUT", url=url, data=data)
156162

157163
@classmethod
158164
def decode_upload_token(cls, token: str) -> tuple[str, uuid.UUID]:

scripts/add_chats.py

+7-5
Original file line numberDiff line numberDiff line change
@@ -34,15 +34,17 @@ def main():
3434
documents = []
3535
for i in range(5):
3636
name = f"document{i}.txt"
37-
document_info = (
37+
document_upload = (
3838
client.post("/document", json={"name": name}).raise_for_status().json()
3939
)
40-
client.put(
41-
document_info["url"],
42-
data=document_info["data"],
40+
parameters = document_upload["parameters"]
41+
client.request(
42+
parameters["method"],
43+
parameters["url"],
44+
data=parameters["data"],
4345
files={"file": f"Content of {name}".encode()},
4446
).raise_for_status()
45-
documents.append(document_info["document"])
47+
documents.append(document_upload["document"])
4648

4749
## chat 1
4850

tests/api/test_e2e.py

+8-6
Original file line numberDiff line numberDiff line change
@@ -75,20 +75,22 @@ def check_api(config):
7575

7676
assert client.get("/chats").raise_for_status().json() == []
7777

78-
document_info = (
78+
document_upload = (
7979
client.post("/document", json={"name": document_path.name})
8080
.raise_for_status()
8181
.json()
8282
)
83-
document = document_info["document"]
83+
document = document_upload["document"]
8484
assert document["name"] == document_path.name
8585

86+
parameters = document_upload["parameters"]
8687
with open(document_path, "rb") as file:
87-
client.put(
88-
document_info["url"],
89-
data=document_info["data"],
88+
client.request(
89+
parameters["method"],
90+
parameters["url"],
91+
data=parameters["data"],
9092
files={"file": file},
91-
).raise_for_status()
93+
)
9294

9395
components = client.get("/components").raise_for_status().json()
9496
documents = components["documents"]

0 commit comments

Comments
 (0)