Skip to content

Commit

Permalink
Make document upload method flexible (#238)
Browse files Browse the repository at this point in the history
Conflicts:
	ragna/_api/core.py
	ragna/deploy/_api/core.py
  • Loading branch information
pmeier committed Dec 7, 2023
1 parent f650cdd commit d5fb9b3
Show file tree
Hide file tree
Showing 7 changed files with 37 additions and 25 deletions.
6 changes: 3 additions & 3 deletions ragna/_api/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,16 +100,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.core.document.get_upload_info(
metadata, parameters = await config.core.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(
Expand Down
5 changes: 2 additions & 3 deletions ragna/_api/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
11 changes: 6 additions & 5 deletions ragna/_ui/resources/upload.js
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
2 changes: 2 additions & 0 deletions ragna/core/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
"Config",
"Document",
"DocumentHandler",
"DocumentUploadParameters",
"EnvVarRequirement",
"LocalDocument",
"Message",
Expand Down Expand Up @@ -35,6 +36,7 @@
from ._document import (
Document,
DocumentHandler,
DocumentUploadParameters,
LocalDocument,
Page,
PdfDocumentHandler,
Expand Down
12 changes: 9 additions & 3 deletions ragna/core/_document.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,12 @@
from ._config import Config


class DocumentUploadParameters(BaseModel):
method: str
url: str
data: dict


class Document(RequirementsMixin, abc.ABC):
"""Abstract base class for all documents."""

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -138,7 +144,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(
Expand All @@ -152,7 +158,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]:
Expand Down
12 changes: 7 additions & 5 deletions scripts/add_chats.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
14 changes: 8 additions & 6 deletions tests/api/test_e2e.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,20 +75,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"]
Expand Down

0 comments on commit d5fb9b3

Please sign in to comment.