Save file in chunks

AI-Engineer-Foundation · Aug 8, 2023 · d8ce887 · d8ce887
1 parent 9978c34
commit d8ce887
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 3 deletions.
diff --git a/agent/python/agent_protocol/agent.py b/agent/python/agent_protocol/agent.py
@@ -1,6 +1,7 @@
 import asyncio
 import os
 
+import aiofiles
 from fastapi import APIRouter, UploadFile
 from fastapi.responses import FileResponse
 from hypercorn.asyncio import serve
@@ -146,8 +147,9 @@ async def upload_agent_task_artifacts(
     if not os.path.exists(path):
         os.makedirs(path)
 
-    with open(os.path.join(path, file.filename), "wb") as f:
-        f.write(await file.read())
+    async with aiofiles.open(os.path.join(path, file.filename), "wb") as f:
+        while content := await file.read(1024 * 1024):  # async read chunk ~1MiB
+            await f.write(content)
 
     return artifact
 

diff --git a/agent/python/poetry.lock b/agent/python/poetry.lock
diff --git a/agent/python/pyproject.toml b/agent/python/pyproject.toml
@@ -18,6 +18,7 @@ pydantic = "^1.10.5, <2"
 click = "^8.1.6"
 requests = "^2.31.0"
 python-multipart = "^0.0.6"
+aiofiles = "^23.1.0"
 
 [tool.poetry.group.dev.dependencies]
 fastapi-code-generator = "^0.4.2"