From e539cf5c28c25e81efebd16c511ebaa1d9a78a73 Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Sat, 23 Dec 2023 15:38:52 -0800 Subject: [PATCH] feat: gguf upload --- backend/.gitignore | 4 +- backend/apps/web/main.py | 3 +- backend/apps/web/routers/utils.py | 78 ++++ backend/utils/misc.py | 8 + src/lib/components/chat/SettingsModal.svelte | 457 +++++++++++++++---- src/lib/constants.ts | 3 +- 6 files changed, 457 insertions(+), 96 deletions(-) create mode 100644 backend/apps/web/routers/utils.py diff --git a/backend/.gitignore b/backend/.gitignore index 6d1787081..bbb8ba188 100644 --- a/backend/.gitignore +++ b/backend/.gitignore @@ -1,2 +1,4 @@ __pycache__ -.env \ No newline at end of file +.env +_old +uploads \ No newline at end of file diff --git a/backend/apps/web/main.py b/backend/apps/web/main.py index 0380f0c34..854f16269 100644 --- a/backend/apps/web/main.py +++ b/backend/apps/web/main.py @@ -1,7 +1,7 @@ from fastapi import FastAPI, Request, Depends, HTTPException from fastapi.middleware.cors import CORSMiddleware -from apps.web.routers import auths, users +from apps.web.routers import auths, users, utils from config import WEBUI_VERSION, WEBUI_AUTH app = FastAPI() @@ -19,6 +19,7 @@ app.add_middleware( app.include_router(auths.router, prefix="/auths", tags=["auths"]) app.include_router(users.router, prefix="/users", tags=["users"]) +app.include_router(utils.router, prefix="/utils", tags=["utils"]) @app.get("/") diff --git a/backend/apps/web/routers/utils.py b/backend/apps/web/routers/utils.py new file mode 100644 index 000000000..8c822f022 --- /dev/null +++ b/backend/apps/web/routers/utils.py @@ -0,0 +1,78 @@ +from fastapi import APIRouter, UploadFile, File, BackgroundTasks +from fastapi import Depends, HTTPException, status +from starlette.responses import StreamingResponse + +from pydantic import BaseModel + +from utils.misc import calculate_sha256 +import requests + + +import os +import asyncio +import json +from config import OLLAMA_API_BASE_URL + + +router = APIRouter() + + +class UploadBlobForm(BaseModel): + filename: str + + +@router.post("/upload") +async def upload(file: UploadFile = File(...)): + os.makedirs("./uploads", exist_ok=True) + file_path = os.path.join("./uploads", file.filename) + + def file_write_stream(): + total = 0 + total_size = file.size + chunk_size = 1024 * 1024 + + done = False + try: + with open(file_path, "wb") as f: + while True: + chunk = file.file.read(chunk_size) + if not chunk: + break + f.write(chunk) + total += len(chunk) + done = total_size == total + + res = { + "total": total_size, + "uploaded": total, + } + + yield f"data: {json.dumps(res)}\n\n" + + if done: + with open(file_path, "rb") as f: + hashed = calculate_sha256(f) + + f.seek(0) + file_data = f.read() + + url = f"{OLLAMA_API_BASE_URL}/blobs/sha256:{hashed}" + + response = requests.post(url, data=file_data) + + if response.ok: + res = { + "done": done, + "blob": f"sha256:{hashed}", + } + os.remove(file_path) + + yield f"data: {json.dumps(res)}\n\n" + else: + raise "Ollama: Could not create blob, Please try again." + + except Exception as e: + res = {"error": str(e)} + yield f"data: {json.dumps(res)}\n\n" + + return StreamingResponse(file_write_stream(), media_type="text/event-stream") diff --git a/backend/utils/misc.py b/backend/utils/misc.py index e4011b34c..c6508487a 100644 --- a/backend/utils/misc.py +++ b/backend/utils/misc.py @@ -13,3 +13,11 @@ def get_gravatar_url(email): # Grab the actual image URL return f"https://www.gravatar.com/avatar/{hash_hex}?d=mp" + + +def calculate_sha256(file): + sha256 = hashlib.sha256() + # Read the file in chunks to efficiently handle large files + for chunk in iter(lambda: file.read(8192), b""): + sha256.update(chunk) + return sha256.hexdigest() diff --git a/src/lib/components/chat/SettingsModal.svelte b/src/lib/components/chat/SettingsModal.svelte index 5d4334c4a..a17c955d7 100644 --- a/src/lib/components/chat/SettingsModal.svelte +++ b/src/lib/components/chat/SettingsModal.svelte @@ -1,12 +1,18 @@