Sfoglia il codice sorgente

Merge branch 'main' into rag

Timothy Jaeryang Baek 1 anno fa
parent
commit
ee1559378d
41 ha cambiato i file con 1192 aggiunte e 369 eliminazioni
  1. 291 0
      .gitignore
  2. 3 0
      .prettierignore
  3. 35 0
      INSTALLATION.md
  4. 64 27
      README.md
  5. 2 1
      backend/.gitignore
  6. 7 3
      backend/apps/ollama/main.py
  7. 3 5
      backend/apps/web/main.py
  8. 0 10
      backend/apps/web/models/users.py
  9. 12 25
      backend/apps/web/routers/auths.py
  10. 36 97
      backend/apps/web/routers/chats.py
  11. 63 115
      backend/apps/web/routers/modelfiles.py
  12. 31 61
      backend/apps/web/routers/users.py
  13. 2 0
      backend/requirements.txt
  14. 18 14
      backend/utils/utils.py
  15. 2 2
      docker-compose.api.yaml
  16. 6 0
      docker-compose.data.yaml
  17. 3 3
      docker-compose.gpu.yaml
  18. 3 3
      docker-compose.yaml
  19. 0 0
      kubernetes/helm/.helmignore
  20. 5 0
      kubernetes/helm/Chart.yaml
  21. 4 0
      kubernetes/helm/templates/ollama-namespace.yaml
  22. 13 0
      kubernetes/helm/templates/ollama-service.yaml
  23. 55 0
      kubernetes/helm/templates/ollama-statefulset.yaml
  24. 38 0
      kubernetes/helm/templates/webui-deployment.yaml
  25. 23 0
      kubernetes/helm/templates/webui-ingress.yaml
  26. 12 0
      kubernetes/helm/templates/webui-pvc.yaml
  27. 15 0
      kubernetes/helm/templates/webui-service.yaml
  28. 38 0
      kubernetes/helm/values.yaml
  29. 4 0
      kubernetes/manifest/base/ollama-namespace.yaml
  30. 12 0
      kubernetes/manifest/base/ollama-service.yaml
  31. 37 0
      kubernetes/manifest/base/ollama-statefulset.yaml
  32. 28 0
      kubernetes/manifest/base/webui-deployment.yaml
  33. 20 0
      kubernetes/manifest/base/webui-ingress.yaml
  34. 15 0
      kubernetes/manifest/base/webui-service.yaml
  35. 12 0
      kubernetes/manifest/kustomization.yaml
  36. 17 0
      kubernetes/manifest/patches/ollama-statefulset-gpu.yaml
  37. 237 0
      run-compose.sh
  38. 1 1
      src/app.css
  39. 18 0
      src/lib/components/chat/MessageInput.svelte
  40. 1 0
      src/lib/components/chat/Messages/ResponseMessage.svelte
  41. 6 2
      src/lib/components/chat/Messages/UserMessage.svelte

+ 291 - 0
.gitignore

@@ -8,3 +8,294 @@ node_modules
 !.env.example
 !.env.example
 vite.config.js.timestamp-*
 vite.config.js.timestamp-*
 vite.config.ts.timestamp-*
 vite.config.ts.timestamp-*
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Logs
+logs
+*.log
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+lerna-debug.log*
+.pnpm-debug.log*
+
+# Diagnostic reports (https://nodejs.org/api/report.html)
+report.[0-9]*.[0-9]*.[0-9]*.[0-9]*.json
+
+# Runtime data
+pids
+*.pid
+*.seed
+*.pid.lock
+
+# Directory for instrumented libs generated by jscoverage/JSCover
+lib-cov
+
+# Coverage directory used by tools like istanbul
+coverage
+*.lcov
+
+# nyc test coverage
+.nyc_output
+
+# Grunt intermediate storage (https://gruntjs.com/creating-plugins#storing-task-files)
+.grunt
+
+# Bower dependency directory (https://bower.io/)
+bower_components
+
+# node-waf configuration
+.lock-wscript
+
+# Compiled binary addons (https://nodejs.org/api/addons.html)
+build/Release
+
+# Dependency directories
+node_modules/
+jspm_packages/
+
+# Snowpack dependency directory (https://snowpack.dev/)
+web_modules/
+
+# TypeScript cache
+*.tsbuildinfo
+
+# Optional npm cache directory
+.npm
+
+# Optional eslint cache
+.eslintcache
+
+# Optional stylelint cache
+.stylelintcache
+
+# Microbundle cache
+.rpt2_cache/
+.rts2_cache_cjs/
+.rts2_cache_es/
+.rts2_cache_umd/
+
+# Optional REPL history
+.node_repl_history
+
+# Output of 'npm pack'
+*.tgz
+
+# Yarn Integrity file
+.yarn-integrity
+
+# dotenv environment variable files
+.env
+.env.development.local
+.env.test.local
+.env.production.local
+.env.local
+
+# parcel-bundler cache (https://parceljs.org/)
+.cache
+.parcel-cache
+
+# Next.js build output
+.next
+out
+
+# Nuxt.js build / generate output
+.nuxt
+dist
+
+# Gatsby files
+.cache/
+# Comment in the public line in if your project uses Gatsby and not Next.js
+# https://nextjs.org/blog/next-9-1#public-directory-support
+# public
+
+# vuepress build output
+.vuepress/dist
+
+# vuepress v2.x temp and cache directory
+.temp
+.cache
+
+# Docusaurus cache and generated files
+.docusaurus
+
+# Serverless directories
+.serverless/
+
+# FuseBox cache
+.fusebox/
+
+# DynamoDB Local files
+.dynamodb/
+
+# TernJS port file
+.tern-port
+
+# Stores VSCode versions used for testing VSCode extensions
+.vscode-test
+
+# yarn v2
+.yarn/cache
+.yarn/unplugged
+.yarn/build-state.yml
+.yarn/install-state.gz
+.pnp.*

+ 3 - 0
.prettierignore

@@ -11,3 +11,6 @@ node_modules
 pnpm-lock.yaml
 pnpm-lock.yaml
 package-lock.json
 package-lock.json
 yarn.lock
 yarn.lock
+
+# Ignore kubernetes files
+kubernetes

+ 35 - 0
INSTALLATION.md

@@ -0,0 +1,35 @@
+### Installing Both Ollama and Ollama Web UI Using Kustomize
+
+For cpu-only pod
+
+```bash
+kubectl apply -f ./kubernetes/manifest/base
+```
+
+For gpu-enabled pod
+
+```bash
+kubectl apply -k ./kubernetes/manifest
+```
+
+### Installing Both Ollama and Ollama Web UI Using Helm
+
+Package Helm file first
+
+```bash
+helm package ./kubernetes/helm/
+```
+
+For cpu-only pod
+
+```bash
+helm install ollama-webui ./ollama-webui-*.tgz
+```
+
+For gpu-enabled pod
+
+```bash
+helm install ollama-webui ./ollama-webui-*.tgz --set ollama.resources.limits.nvidia.com/gpu="1"
+```
+
+Check the `kubernetes/helm/values.yaml` file to know which parameters are available for customization

+ 64 - 27
README.md

@@ -27,7 +27,7 @@ Also check our sibling project, [OllamaHub](https://ollamahub.com/), where you c
 
 
 - ⚡ **Swift Responsiveness**: Enjoy fast and responsive performance.
 - ⚡ **Swift Responsiveness**: Enjoy fast and responsive performance.
 
 
-- 🚀 **Effortless Setup**: Install seamlessly using Docker for a hassle-free experience.
+- 🚀 **Effortless Setup**: Install seamlessly using Docker or Kubernetes (kubectl, kustomize or helm) for a hassle-free experience.
 
 
 - 💻 **Code Syntax Highlighting**: Enjoy enhanced code readability with our syntax highlighting feature.
 - 💻 **Code Syntax Highlighting**: Enjoy enhanced code readability with our syntax highlighting feature.
 
 
@@ -79,32 +79,6 @@ Don't forget to explore our sibling project, [OllamaHub](https://ollamahub.com/)
 
 
 - **Privacy and Data Security:** We prioritize your privacy and data security above all. Please be reassured that all data entered into the Ollama Web UI is stored locally on your device. Our system is designed to be privacy-first, ensuring that no external requests are made, and your data does not leave your local environment. We are committed to maintaining the highest standards of data privacy and security, ensuring that your information remains confidential and under your control.
 - **Privacy and Data Security:** We prioritize your privacy and data security above all. Please be reassured that all data entered into the Ollama Web UI is stored locally on your device. Our system is designed to be privacy-first, ensuring that no external requests are made, and your data does not leave your local environment. We are committed to maintaining the highest standards of data privacy and security, ensuring that your information remains confidential and under your control.
 
 
-### Installing Both Ollama and Ollama Web UI Using Docker Compose
-
-If you don't have Ollama installed yet, you can use the provided Docker Compose file for a hassle-free installation. Simply run the following command:
-
-```bash
-docker compose up -d --build
-```
-
-This command will install both Ollama and Ollama Web UI on your system.
-
-#### Enable GPU
-
-Use the additional Docker Compose file designed to enable GPU support by running the following command:
-
-```bash
-docker compose -f docker-compose.yml -f docker-compose.gpu.yml up -d --build
-```
-
-#### Expose Ollama API outside the container stack
-
-Deploy the service with an additional Docker Compose file designed for API exposure:
-
-```bash
-docker compose -f docker-compose.yml -f docker-compose.api.yml up -d --build
-```
-
 ### Installing Ollama Web UI Only
 ### Installing Ollama Web UI Only
 
 
 #### Prerequisites
 #### Prerequisites
@@ -149,6 +123,69 @@ docker build -t ollama-webui .
 docker run -d -p 3000:8080 -e OLLAMA_API_BASE_URL=https://example.com/api -v ollama-webui:/app/backend/data --name ollama-webui --restart always ollama-webui
 docker run -d -p 3000:8080 -e OLLAMA_API_BASE_URL=https://example.com/api -v ollama-webui:/app/backend/data --name ollama-webui --restart always ollama-webui
 ```
 ```
 
 
+### Installing Both Ollama and Ollama Web UI
+
+#### Using Docker Compose
+
+If you don't have Ollama installed yet, you can use the provided Docker Compose file for a hassle-free installation. Simply run the following command:
+
+```bash
+docker compose up -d --build
+```
+
+This command will install both Ollama and Ollama Web UI on your system.
+
+##### Enable GPU
+
+Use the additional Docker Compose file designed to enable GPU support by running the following command:
+
+```bash
+docker compose -f docker-compose.yaml -f docker-compose.gpu.yaml up -d --build
+```
+
+##### Expose Ollama API outside the container stack
+
+Deploy the service with an additional Docker Compose file designed for API exposure:
+
+```bash
+docker compose -f docker-compose.yaml -f docker-compose.api.yaml up -d --build
+```
+
+#### Using Provided `run-compose.sh` Script (Linux)
+
+Also available on Windows under any docker-enabled WSL2 linux distro (you have to enable it from Docker Desktop)
+
+Simply run the following command to grant execute permission to script:
+
+```bash
+chmod +x run-compose.sh
+```
+
+##### For CPU only container
+
+```bash
+./run-compose.sh
+```
+
+##### Enable GPU
+
+For GPU enabled container (to enable this you must have your gpu driver for docker, it mostly works with nvidia so this is the official install guide: [nvidia-container-toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html))
+Warning! A GPU-enabled installation has only been tested using linux and nvidia GPU, full functionalities are not guaranteed under Windows or Macos or using a different GPU
+
+```bash
+./run-compose.sh --enable-gpu
+```
+
+Note that both the above commands will use the latest production docker image in repository, to be able to build the latest local version you'll need to append the `--build` parameter, for example:
+
+```bash
+./run-compose.sh --enable-gpu --build
+```
+
+#### Using Alternative Methods (Kustomize or Helm)
+
+See [INSTALLATION.md](/INSTALLATION.md) for information on how to install and/or join our [Ollama Web UI Discord community](https://discord.gg/5rJgQTnV4s).
+
 ## How to Install Without Docker
 ## How to Install Without Docker
 
 
 While we strongly recommend using our convenient Docker container installation for optimal support, we understand that some situations may require a non-Docker setup, especially for development purposes. Please note that non-Docker installations are not officially supported, and you might need to troubleshoot on your own.
 While we strongly recommend using our convenient Docker container installation for optimal support, we understand that some situations may require a non-Docker setup, especially for development purposes. Please note that non-Docker installations are not officially supported, and you might need to troubleshoot on your own.

+ 2 - 1
backend/.gitignore

@@ -4,4 +4,5 @@ _old
 uploads
 uploads
 .ipynb_checkpoints
 .ipynb_checkpoints
 *.db
 *.db
-_test
+_test
+Pipfile

+ 7 - 3
backend/apps/ollama/main.py

@@ -8,7 +8,7 @@ import json
 
 
 from apps.web.models.users import Users
 from apps.web.models.users import Users
 from constants import ERROR_MESSAGES
 from constants import ERROR_MESSAGES
-from utils.utils import extract_token_from_auth_header
+from utils.utils import decode_token
 from config import OLLAMA_API_BASE_URL, WEBUI_AUTH
 from config import OLLAMA_API_BASE_URL, WEBUI_AUTH
 
 
 app = Flask(__name__)
 app = Flask(__name__)
@@ -34,8 +34,12 @@ def proxy(path):
     # Basic RBAC support
     # Basic RBAC support
     if WEBUI_AUTH:
     if WEBUI_AUTH:
         if "Authorization" in headers:
         if "Authorization" in headers:
-            token = extract_token_from_auth_header(headers["Authorization"])
-            user = Users.get_user_by_token(token)
+            _, credentials = headers["Authorization"].split()
+            token_data = decode_token(credentials)
+            if token_data is None or "email" not in token_data:
+                return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
+
+            user = Users.get_user_by_email(token_data["email"])
             if user:
             if user:
                 # Only user and admin roles can access
                 # Only user and admin roles can access
                 if user.role in ["user", "admin"]:
                 if user.role in ["user", "admin"]:

+ 3 - 5
backend/apps/web/main.py

@@ -1,6 +1,6 @@
-from fastapi import FastAPI, Request, Depends, HTTPException
+from fastapi import FastAPI, Depends
+from fastapi.routing import APIRoute
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.cors import CORSMiddleware
-
 from apps.web.routers import auths, users, chats, modelfiles, utils
 from apps.web.routers import auths, users, chats, modelfiles, utils
 from config import WEBUI_VERSION, WEBUI_AUTH
 from config import WEBUI_VERSION, WEBUI_AUTH
 
 
@@ -16,13 +16,11 @@ app.add_middleware(
     allow_headers=["*"],
     allow_headers=["*"],
 )
 )
 
 
-
 app.include_router(auths.router, prefix="/auths", tags=["auths"])
 app.include_router(auths.router, prefix="/auths", tags=["auths"])
+
 app.include_router(users.router, prefix="/users", tags=["users"])
 app.include_router(users.router, prefix="/users", tags=["users"])
 app.include_router(chats.router, prefix="/chats", tags=["chats"])
 app.include_router(chats.router, prefix="/chats", tags=["chats"])
 app.include_router(modelfiles.router, prefix="/modelfiles", tags=["modelfiles"])
 app.include_router(modelfiles.router, prefix="/modelfiles", tags=["modelfiles"])
-
-
 app.include_router(utils.router, prefix="/utils", tags=["utils"])
 app.include_router(utils.router, prefix="/utils", tags=["utils"])
 
 
 
 

+ 0 - 10
backend/apps/web/models/users.py

@@ -3,8 +3,6 @@ from peewee import *
 from playhouse.shortcuts import model_to_dict
 from playhouse.shortcuts import model_to_dict
 from typing import List, Union, Optional
 from typing import List, Union, Optional
 import time
 import time
-
-from utils.utils import decode_token
 from utils.misc import get_gravatar_url
 from utils.misc import get_gravatar_url
 
 
 from apps.web.internal.db import DB
 from apps.web.internal.db import DB
@@ -85,14 +83,6 @@ class UsersTable:
         except:
         except:
             return None
             return None
 
 
-    def get_user_by_token(self, token: str) -> Optional[UserModel]:
-        data = decode_token(token)
-
-        if data != None and "email" in data:
-            return self.get_user_by_email(data["email"])
-        else:
-            return None
-
     def get_users(self, skip: int = 0, limit: int = 50) -> List[UserModel]:
     def get_users(self, skip: int = 0, limit: int = 50) -> List[UserModel]:
         return [
         return [
             UserModel(**model_to_dict(user))
             UserModel(**model_to_dict(user))

+ 12 - 25
backend/apps/web/routers/auths.py

@@ -19,11 +19,7 @@ from apps.web.models.auths import (
 from apps.web.models.users import Users
 from apps.web.models.users import Users
 
 
 
 
-from utils.utils import (
-    get_password_hash,
-    bearer_scheme,
-    create_token,
-)
+from utils.utils import get_password_hash, get_current_user, create_token
 from utils.misc import get_gravatar_url
 from utils.misc import get_gravatar_url
 from constants import ERROR_MESSAGES
 from constants import ERROR_MESSAGES
 
 
@@ -36,22 +32,14 @@ router = APIRouter()
 
 
 
 
 @router.get("/", response_model=UserResponse)
 @router.get("/", response_model=UserResponse)
-async def get_session_user(cred=Depends(bearer_scheme)):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-    if user:
-        return {
-            "id": user.id,
-            "email": user.email,
-            "name": user.name,
-            "role": user.role,
-            "profile_image_url": user.profile_image_url,
-        }
-    else:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
-        )
+async def get_session_user(user=Depends(get_current_user)):
+    return {
+        "id": user.id,
+        "email": user.email,
+        "name": user.name,
+        "role": user.role,
+        "profile_image_url": user.profile_image_url,
+    }
 
 
 
 
 ############################
 ############################
@@ -60,10 +48,9 @@ async def get_session_user(cred=Depends(bearer_scheme)):
 
 
 
 
 @router.post("/update/password", response_model=bool)
 @router.post("/update/password", response_model=bool)
-async def update_password(form_data: UpdatePasswordForm, cred=Depends(bearer_scheme)):
-    token = cred.credentials
-    session_user = Users.get_user_by_token(token)
-
+async def update_password(
+    form_data: UpdatePasswordForm, session_user=Depends(get_current_user)
+):
     if session_user:
     if session_user:
         user = Auths.authenticate_user(session_user.email, form_data.password)
         user = Auths.authenticate_user(session_user.email, form_data.password)
 
 

+ 36 - 97
backend/apps/web/routers/chats.py

@@ -1,8 +1,7 @@
-from fastapi import Response
-from fastapi import Depends, FastAPI, HTTPException, status
+from fastapi import Depends, Request, HTTPException, status
 from datetime import datetime, timedelta
 from datetime import datetime, timedelta
 from typing import List, Union, Optional
 from typing import List, Union, Optional
-
+from utils.utils import get_current_user
 from fastapi import APIRouter
 from fastapi import APIRouter
 from pydantic import BaseModel
 from pydantic import BaseModel
 import json
 import json
@@ -30,17 +29,10 @@ router = APIRouter()
 
 
 
 
 @router.get("/", response_model=List[ChatTitleIdResponse])
 @router.get("/", response_model=List[ChatTitleIdResponse])
-async def get_user_chats(skip: int = 0, limit: int = 50, cred=Depends(bearer_scheme)):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        return Chats.get_chat_lists_by_user_id(user.id, skip, limit)
-    else:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
-        )
+async def get_user_chats(
+    user=Depends(get_current_user), skip: int = 0, limit: int = 50
+):
+    return Chats.get_chat_lists_by_user_id(user.id, skip, limit)
 
 
 
 
 ############################
 ############################
@@ -49,20 +41,11 @@ async def get_user_chats(skip: int = 0, limit: int = 50, cred=Depends(bearer_sch
 
 
 
 
 @router.get("/all", response_model=List[ChatResponse])
 @router.get("/all", response_model=List[ChatResponse])
-async def get_all_user_chats(cred=Depends(bearer_scheme)):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        return [
-            ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
-            for chat in Chats.get_all_chats_by_user_id(user.id)
-        ]
-    else:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
-        )
+async def get_all_user_chats(user=Depends(get_current_user)):
+    return [
+        ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
+        for chat in Chats.get_all_chats_by_user_id(user.id)
+    ]
 
 
 
 
 ############################
 ############################
@@ -71,18 +54,9 @@ async def get_all_user_chats(cred=Depends(bearer_scheme)):
 
 
 
 
 @router.post("/new", response_model=Optional[ChatResponse])
 @router.post("/new", response_model=Optional[ChatResponse])
-async def create_new_chat(form_data: ChatForm, cred=Depends(bearer_scheme)):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        chat = Chats.insert_new_chat(user.id, form_data)
-        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
-    else:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
-        )
+async def create_new_chat(form_data: ChatForm, user=Depends(get_current_user)):
+    chat = Chats.insert_new_chat(user.id, form_data)
+    return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
 
 
 
 
 ############################
 ############################
@@ -91,24 +65,14 @@ async def create_new_chat(form_data: ChatForm, cred=Depends(bearer_scheme)):
 
 
 
 
 @router.get("/{id}", response_model=Optional[ChatResponse])
 @router.get("/{id}", response_model=Optional[ChatResponse])
-async def get_chat_by_id(id: str, cred=Depends(bearer_scheme)):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        chat = Chats.get_chat_by_id_and_user_id(id, user.id)
-
-        if chat:
-            return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
-        else:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail=ERROR_MESSAGES.NOT_FOUND,
-            )
+async def get_chat_by_id(id: str, user=Depends(get_current_user)):
+    chat = Chats.get_chat_by_id_and_user_id(id, user.id)
+
+    if chat:
+        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
     else:
     else:
         raise HTTPException(
         raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
+            status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
         )
         )
 
 
 
 
@@ -118,26 +82,19 @@ async def get_chat_by_id(id: str, cred=Depends(bearer_scheme)):
 
 
 
 
 @router.post("/{id}", response_model=Optional[ChatResponse])
 @router.post("/{id}", response_model=Optional[ChatResponse])
-async def update_chat_by_id(id: str, form_data: ChatForm, cred=Depends(bearer_scheme)):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        chat = Chats.get_chat_by_id_and_user_id(id, user.id)
-        if chat:
-            updated_chat = {**json.loads(chat.chat), **form_data.chat}
-
-            chat = Chats.update_chat_by_id(id, updated_chat)
-            return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
-        else:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
-            )
+async def update_chat_by_id(
+    id: str, form_data: ChatForm, user=Depends(get_current_user)
+):
+    chat = Chats.get_chat_by_id_and_user_id(id, user.id)
+    if chat:
+        updated_chat = {**json.loads(chat.chat), **form_data.chat}
+
+        chat = Chats.update_chat_by_id(id, updated_chat)
+        return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
     else:
     else:
         raise HTTPException(
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
+            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
         )
         )
 
 
 
 
@@ -147,18 +104,9 @@ async def update_chat_by_id(id: str, form_data: ChatForm, cred=Depends(bearer_sc
 
 
 
 
 @router.delete("/{id}", response_model=bool)
 @router.delete("/{id}", response_model=bool)
-async def delete_chat_by_id(id: str, cred=Depends(bearer_scheme)):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        result = Chats.delete_chat_by_id_and_user_id(id, user.id)
-        return result
-    else:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
-        )
+async def delete_chat_by_id(id: str, user=Depends(get_current_user)):
+    result = Chats.delete_chat_by_id_and_user_id(id, user.id)
+    return result
 
 
 
 
 ############################
 ############################
@@ -167,15 +115,6 @@ async def delete_chat_by_id(id: str, cred=Depends(bearer_scheme)):
 
 
 
 
 @router.delete("/", response_model=bool)
 @router.delete("/", response_model=bool)
-async def delete_all_user_chats(cred=Depends(bearer_scheme)):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        result = Chats.delete_chats_by_user_id(user.id)
-        return result
-    else:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
-        )
+async def delete_all_user_chats(user=Depends(get_current_user)):
+    result = Chats.delete_chats_by_user_id(user.id)
+    return result

+ 63 - 115
backend/apps/web/routers/modelfiles.py

@@ -1,4 +1,3 @@
-from fastapi import Response
 from fastapi import Depends, FastAPI, HTTPException, status
 from fastapi import Depends, FastAPI, HTTPException, status
 from datetime import datetime, timedelta
 from datetime import datetime, timedelta
 from typing import List, Union, Optional
 from typing import List, Union, Optional
@@ -6,8 +5,6 @@ from typing import List, Union, Optional
 from fastapi import APIRouter
 from fastapi import APIRouter
 from pydantic import BaseModel
 from pydantic import BaseModel
 import json
 import json
-
-from apps.web.models.users import Users
 from apps.web.models.modelfiles import (
 from apps.web.models.modelfiles import (
     Modelfiles,
     Modelfiles,
     ModelfileForm,
     ModelfileForm,
@@ -16,9 +13,7 @@ from apps.web.models.modelfiles import (
     ModelfileResponse,
     ModelfileResponse,
 )
 )
 
 
-from utils.utils import (
-    bearer_scheme,
-)
+from utils.utils import get_current_user
 from constants import ERROR_MESSAGES
 from constants import ERROR_MESSAGES
 
 
 router = APIRouter()
 router = APIRouter()
@@ -29,17 +24,8 @@ router = APIRouter()
 
 
 
 
 @router.get("/", response_model=List[ModelfileResponse])
 @router.get("/", response_model=List[ModelfileResponse])
-async def get_modelfiles(skip: int = 0, limit: int = 50, cred=Depends(bearer_scheme)):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        return Modelfiles.get_modelfiles(skip, limit)
-    else:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
-        )
+async def get_modelfiles(skip: int = 0, limit: int = 50, user=Depends(get_current_user)):
+    return Modelfiles.get_modelfiles(skip, limit)
 
 
 
 
 ############################
 ############################
@@ -48,36 +34,28 @@ async def get_modelfiles(skip: int = 0, limit: int = 50, cred=Depends(bearer_sch
 
 
 
 
 @router.post("/create", response_model=Optional[ModelfileResponse])
 @router.post("/create", response_model=Optional[ModelfileResponse])
-async def create_new_modelfile(form_data: ModelfileForm, cred=Depends(bearer_scheme)):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        # Admin Only
-        if user.role == "admin":
-            modelfile = Modelfiles.insert_new_modelfile(user.id, form_data)
-
-            if modelfile:
-                return ModelfileResponse(
-                    **{
-                        **modelfile.model_dump(),
-                        "modelfile": json.loads(modelfile.modelfile),
-                    }
-                )
-            else:
-                raise HTTPException(
-                    status_code=status.HTTP_401_UNAUTHORIZED,
-                    detail=ERROR_MESSAGES.DEFAULT(),
-                )
-        else:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
-            )
+async def create_new_modelfile(
+    form_data: ModelfileForm, user=Depends(get_current_user)
+):
+    if user.role != "admin":
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
+        )
+
+    modelfile = Modelfiles.insert_new_modelfile(user.id, form_data)
+
+    if modelfile:
+        return ModelfileResponse(
+            **{
+                **modelfile.model_dump(),
+                "modelfile": json.loads(modelfile.modelfile),
+            }
+        )
     else:
     else:
         raise HTTPException(
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
+            detail=ERROR_MESSAGES.DEFAULT(),
         )
         )
 
 
 
 
@@ -87,31 +65,20 @@ async def create_new_modelfile(form_data: ModelfileForm, cred=Depends(bearer_sch
 
 
 
 
 @router.post("/", response_model=Optional[ModelfileResponse])
 @router.post("/", response_model=Optional[ModelfileResponse])
-async def get_modelfile_by_tag_name(
-    form_data: ModelfileTagNameForm, cred=Depends(bearer_scheme)
-):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        modelfile = Modelfiles.get_modelfile_by_tag_name(form_data.tag_name)
-
-        if modelfile:
-            return ModelfileResponse(
-                **{
-                    **modelfile.model_dump(),
-                    "modelfile": json.loads(modelfile.modelfile),
-                }
-            )
-        else:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail=ERROR_MESSAGES.NOT_FOUND,
-            )
+async def get_modelfile_by_tag_name(form_data: ModelfileTagNameForm, user=Depends(get_current_user)):
+    modelfile = Modelfiles.get_modelfile_by_tag_name(form_data.tag_name)
+
+    if modelfile:
+        return ModelfileResponse(
+            **{
+                **modelfile.model_dump(),
+                "modelfile": json.loads(modelfile.modelfile),
+            }
+        )
     else:
     else:
         raise HTTPException(
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
+            detail=ERROR_MESSAGES.NOT_FOUND,
         )
         )
 
 
 
 
@@ -122,44 +89,34 @@ async def get_modelfile_by_tag_name(
 
 
 @router.post("/update", response_model=Optional[ModelfileResponse])
 @router.post("/update", response_model=Optional[ModelfileResponse])
 async def update_modelfile_by_tag_name(
 async def update_modelfile_by_tag_name(
-    form_data: ModelfileUpdateForm, cred=Depends(bearer_scheme)
+    form_data: ModelfileUpdateForm, user=Depends(get_current_user)
 ):
 ):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        if user.role == "admin":
-            modelfile = Modelfiles.get_modelfile_by_tag_name(form_data.tag_name)
-            if modelfile:
-                updated_modelfile = {
-                    **json.loads(modelfile.modelfile),
-                    **form_data.modelfile,
-                }
-
-                modelfile = Modelfiles.update_modelfile_by_tag_name(
-                    form_data.tag_name, updated_modelfile
-                )
-
-                return ModelfileResponse(
-                    **{
-                        **modelfile.model_dump(),
-                        "modelfile": json.loads(modelfile.modelfile),
-                    }
-                )
-            else:
-                raise HTTPException(
-                    status_code=status.HTTP_401_UNAUTHORIZED,
-                    detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
-                )
-        else:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
-            )
+    if user.role != "admin":
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
+        )
+    modelfile = Modelfiles.get_modelfile_by_tag_name(form_data.tag_name)
+    if modelfile:
+        updated_modelfile = {
+            **json.loads(modelfile.modelfile),
+            **form_data.modelfile,
+        }
+
+        modelfile = Modelfiles.update_modelfile_by_tag_name(
+            form_data.tag_name, updated_modelfile
+        )
+
+        return ModelfileResponse(
+            **{
+                **modelfile.model_dump(),
+                "modelfile": json.loads(modelfile.modelfile),
+            }
+        )
     else:
     else:
         raise HTTPException(
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
+            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
         )
         )
 
 
 
 
@@ -170,22 +127,13 @@ async def update_modelfile_by_tag_name(
 
 
 @router.delete("/delete", response_model=bool)
 @router.delete("/delete", response_model=bool)
 async def delete_modelfile_by_tag_name(
 async def delete_modelfile_by_tag_name(
-    form_data: ModelfileTagNameForm, cred=Depends(bearer_scheme)
+    form_data: ModelfileTagNameForm, user=Depends(get_current_user)
 ):
 ):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        if user.role == "admin":
-            result = Modelfiles.delete_modelfile_by_tag_name(form_data.tag_name)
-            return result
-        else:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
-            )
-    else:
+    if user.role != "admin":
         raise HTTPException(
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
+            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
         )
         )
+
+    result = Modelfiles.delete_modelfile_by_tag_name(form_data.tag_name)
+    return result

+ 31 - 61
backend/apps/web/routers/users.py

@@ -12,11 +12,7 @@ from apps.web.models.users import UserModel, UserRoleUpdateForm, Users
 from apps.web.models.auths import Auths
 from apps.web.models.auths import Auths
 
 
 
 
-from utils.utils import (
-    get_password_hash,
-    bearer_scheme,
-    create_token,
-)
+from utils.utils import get_current_user
 from constants import ERROR_MESSAGES
 from constants import ERROR_MESSAGES
 
 
 router = APIRouter()
 router = APIRouter()
@@ -27,23 +23,13 @@ router = APIRouter()
 
 
 
 
 @router.get("/", response_model=List[UserModel])
 @router.get("/", response_model=List[UserModel])
-async def get_users(skip: int = 0, limit: int = 50, cred=Depends(bearer_scheme)):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        if user.role == "admin":
-            return Users.get_users(skip, limit)
-        else:
-            raise HTTPException(
-                status_code=status.HTTP_403_FORBIDDEN,
-                detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
-            )
-    else:
+async def get_users(skip: int = 0, limit: int = 50, user=Depends(get_current_user)):
+    if user.role != "admin":
         raise HTTPException(
         raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
         )
         )
+    return Users.get_users(skip, limit)
 
 
 
 
 ############################
 ############################
@@ -52,28 +38,21 @@ async def get_users(skip: int = 0, limit: int = 50, cred=Depends(bearer_scheme))
 
 
 
 
 @router.post("/update/role", response_model=Optional[UserModel])
 @router.post("/update/role", response_model=Optional[UserModel])
-async def update_user_role(form_data: UserRoleUpdateForm, cred=Depends(bearer_scheme)):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        if user.role == "admin":
-            if user.id != form_data.id:
-                return Users.update_user_role_by_id(form_data.id, form_data.role)
-            else:
-                raise HTTPException(
-                    status_code=status.HTTP_403_FORBIDDEN,
-                    detail=ERROR_MESSAGES.ACTION_PROHIBITED,
-                )
-        else:
-            raise HTTPException(
-                status_code=status.HTTP_403_FORBIDDEN,
-                detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
-            )
+async def update_user_role(
+    form_data: UserRoleUpdateForm, user=Depends(get_current_user)
+):
+    if user.role != "admin":
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
+        )
+
+    if user.id != form_data.id:
+        return Users.update_user_role_by_id(form_data.id, form_data.role)
     else:
     else:
         raise HTTPException(
         raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail=ERROR_MESSAGES.ACTION_PROHIBITED,
         )
         )
 
 
 
 
@@ -83,34 +62,25 @@ async def update_user_role(form_data: UserRoleUpdateForm, cred=Depends(bearer_sc
 
 
 
 
 @router.delete("/{user_id}", response_model=bool)
 @router.delete("/{user_id}", response_model=bool)
-async def delete_user_by_id(user_id: str, cred=Depends(bearer_scheme)):
-    token = cred.credentials
-    user = Users.get_user_by_token(token)
-
-    if user:
-        if user.role == "admin":
-            if user.id != user_id:
-                result = Auths.delete_auth_by_id(user_id)
-
-                if result:
-                    return True
-                else:
-                    raise HTTPException(
-                        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-                        detail=ERROR_MESSAGES.DELETE_USER_ERROR,
-                    )
+async def delete_user_by_id(user_id: str, user=Depends(get_current_user)):
+    if user.role == "admin":
+        if user.id != user_id:
+            result = Auths.delete_auth_by_id(user_id)
+
+            if result:
+                return True
             else:
             else:
                 raise HTTPException(
                 raise HTTPException(
-                    status_code=status.HTTP_403_FORBIDDEN,
-                    detail=ERROR_MESSAGES.ACTION_PROHIBITED,
+                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                    detail=ERROR_MESSAGES.DELETE_USER_ERROR,
                 )
                 )
         else:
         else:
             raise HTTPException(
             raise HTTPException(
                 status_code=status.HTTP_403_FORBIDDEN,
                 status_code=status.HTTP_403_FORBIDDEN,
-                detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
+                detail=ERROR_MESSAGES.ACTION_PROHIBITED,
             )
             )
     else:
     else:
         raise HTTPException(
         raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.INVALID_TOKEN,
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
         )
         )

+ 2 - 0
backend/requirements.txt

@@ -20,3 +20,5 @@ chromadb
 
 
 PyJWT
 PyJWT
 pyjwt[crypto]
 pyjwt[crypto]
+
+black

+ 18 - 14
backend/utils/utils.py

@@ -1,7 +1,9 @@
-from fastapi.security import HTTPBasicCredentials, HTTPBearer
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from fastapi import HTTPException, status, Depends
+from apps.web.models.users import Users
 from pydantic import BaseModel
 from pydantic import BaseModel
 from typing import Union, Optional
 from typing import Union, Optional
-
+from constants import ERROR_MESSAGES
 from passlib.context import CryptContext
 from passlib.context import CryptContext
 from datetime import datetime, timedelta
 from datetime import datetime, timedelta
 import requests
 import requests
@@ -53,16 +55,18 @@ def extract_token_from_auth_header(auth_header: str):
     return auth_header[len("Bearer ") :]
     return auth_header[len("Bearer ") :]
 
 
 
 
-def verify_token(request):
-    try:
-        bearer = request.headers["authorization"]
-        if bearer:
-            token = bearer[len("Bearer ") :]
-            decoded = jwt.decode(
-                token, JWT_SECRET_KEY, options={"verify_signature": False}
+def get_current_user(auth_token: HTTPAuthorizationCredentials = Depends(HTTPBearer())):
+    data = decode_token(auth_token.credentials)
+    if data != None and "email" in data:
+        user = Users.get_user_by_email(data["email"])
+        if user is None:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail=ERROR_MESSAGES.INVALID_TOKEN,
             )
             )
-            return decoded
-        else:
-            return None
-    except Exception as e:
-        return None
+        return user
+    else:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail=ERROR_MESSAGES.UNAUTHORIZED,
+        )

+ 2 - 2
docker-compose.api.yml → docker-compose.api.yaml

@@ -1,7 +1,7 @@
-version: '3.6'
+version: '3.8'
 
 
 services:
 services:
   ollama:
   ollama:
     # Expose Ollama API outside the container stack
     # Expose Ollama API outside the container stack
     ports:
     ports:
-      - 11434:11434
+      - ${OLLAMA_WEBAPI_PORT-11434}:11434

+ 6 - 0
docker-compose.data.yaml

@@ -0,0 +1,6 @@
+version: '3.8'
+
+services:
+  ollama:
+    volumes:
+      - ${OLLAMA_DATA_DIR-./ollama-data}:/root/.ollama

+ 3 - 3
docker-compose.gpu.yml → docker-compose.gpu.yaml

@@ -1,4 +1,4 @@
-version: '3.6'
+version: '3.8'
 
 
 services:
 services:
   ollama:
   ollama:
@@ -7,7 +7,7 @@ services:
       resources:
       resources:
         reservations:
         reservations:
           devices:
           devices:
-            - driver: nvidia
-              count: 1
+            - driver: ${OLLAMA_GPU_DRIVER-nvidia}
+              count: ${OLLAMA_GPU_COUNT-1}
               capabilities:
               capabilities:
                 - gpu
                 - gpu

+ 3 - 3
docker-compose.yml → docker-compose.yaml

@@ -1,4 +1,4 @@
-version: '3.6'
+version: '3.8'
 
 
 services:
 services:
   ollama:
   ollama:
@@ -16,14 +16,14 @@ services:
       args:
       args:
         OLLAMA_API_BASE_URL: '/ollama/api'
         OLLAMA_API_BASE_URL: '/ollama/api'
       dockerfile: Dockerfile
       dockerfile: Dockerfile
-    image: ollama-webui:latest
+    image: ghcr.io/ollama-webui/ollama-webui:main
     container_name: ollama-webui
     container_name: ollama-webui
     volumes:
     volumes:
       - ollama-webui:/app/backend/data
       - ollama-webui:/app/backend/data
     depends_on:
     depends_on:
       - ollama
       - ollama
     ports:
     ports:
-      - 3000:8080
+      - ${OLLAMA_WEBUI_PORT-3000}:8080
     environment:
     environment:
       - "OLLAMA_API_BASE_URL=http://ollama:11434/api"
       - "OLLAMA_API_BASE_URL=http://ollama:11434/api"
     extra_hosts:
     extra_hosts:

+ 0 - 0
kubernetes/helm/.helmignore


+ 5 - 0
kubernetes/helm/Chart.yaml

@@ -0,0 +1,5 @@
+apiVersion: v2
+name: ollama-webui
+description: "Ollama Web UI: A User-Friendly Web Interface for Chat Interactions 👋"
+version: 1.0.0
+icon: https://raw.githubusercontent.com/ollama-webui/ollama-webui/main/static/favicon.png

+ 4 - 0
kubernetes/helm/templates/ollama-namespace.yaml

@@ -0,0 +1,4 @@
+apiVersion: v1
+kind: Namespace
+metadata:
+  name: {{ .Values.namespace }}

+ 13 - 0
kubernetes/helm/templates/ollama-service.yaml

@@ -0,0 +1,13 @@
+apiVersion: v1
+kind: Service
+metadata:
+  name: ollama-service
+  namespace: {{ .Values.namespace }}
+spec:
+  type: {{ .Values.ollama.service.type }}
+  selector:
+    app: ollama
+  ports:
+  - protocol: TCP
+    port: {{ .Values.ollama.servicePort }}
+    targetPort: {{ .Values.ollama.servicePort }}

+ 55 - 0
kubernetes/helm/templates/ollama-statefulset.yaml

@@ -0,0 +1,55 @@
+apiVersion: apps/v1
+kind: StatefulSet
+metadata:
+  name: ollama
+  namespace: {{ .Values.namespace }}
+spec:
+  serviceName: "ollama"
+  replicas: {{ .Values.ollama.replicaCount }}
+  selector:
+    matchLabels:
+      app: ollama
+  template:
+    metadata:
+      labels:
+        app: ollama
+    spec:
+      containers:
+      - name: ollama
+        image: {{ .Values.ollama.image }}
+        ports:
+        - containerPort: {{ .Values.ollama.servicePort }}
+        env:
+        {{- if .Values.ollama.gpu.enabled }}
+          - name: PATH
+            value: /usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
+          - name: LD_LIBRARY_PATH
+            value: /usr/local/nvidia/lib:/usr/local/nvidia/lib64
+          - name: NVIDIA_DRIVER_CAPABILITIES
+            value: compute,utility
+        {{- end}}
+        {{- if .Values.ollama.resources }}
+        resources: {{- toYaml .Values.ollama.resources | nindent 10 }}
+        {{- end }}
+        volumeMounts:
+        - name: ollama-volume
+          mountPath: /root/.ollama
+        tty: true
+      {{- with .Values.ollama.nodeSelector }}
+      nodeSelector:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      tolerations:
+        {{- if .Values.ollama.gpu.enabled }}
+        - key: nvidia.com/gpu
+          operator: Exists
+          effect: NoSchedule
+        {{- end }}
+  volumeClaimTemplates:
+  - metadata:
+      name: ollama-volume
+    spec:
+      accessModes: [ "ReadWriteOnce" ]
+      resources:
+        requests:
+          storage: {{ .Values.ollama.volumeSize }}

+ 38 - 0
kubernetes/helm/templates/webui-deployment.yaml

@@ -0,0 +1,38 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: ollama-webui-deployment
+  namespace: {{ .Values.namespace }}
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: ollama-webui
+  template:
+    metadata:
+      labels:
+        app: ollama-webui
+    spec:
+      containers:
+      - name: ollama-webui
+        image: {{ .Values.webui.image }}
+        ports:
+        - containerPort: 8080
+        {{- if .Values.webui.resources }}
+        resources: {{- toYaml .Values.webui.resources | nindent 10 }}
+        {{- end }}
+        volumeMounts:
+        - name: webui-volume
+          mountPath: /app/backend/data
+        env:
+        - name: OLLAMA_API_BASE_URL
+          value: "http://ollama-service.{{ .Values.namespace }}.svc.cluster.local:{{ .Values.ollama.servicePort }}/api"
+        tty: true
+      {{- with .Values.webui.nodeSelector }}
+      nodeSelector:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      volumes:
+      - name: webui-volume
+        persistentVolumeClaim:
+          claimName: ollama-webui-pvc

+ 23 - 0
kubernetes/helm/templates/webui-ingress.yaml

@@ -0,0 +1,23 @@
+{{- if .Values.webui.ingress.enabled }}
+apiVersion: networking.k8s.io/v1
+kind: Ingress
+metadata:
+  name: ollama-webui-ingress
+  namespace: {{ .Values.namespace }}
+{{- if .Values.webui.ingress.annotations }}
+  annotations:
+{{ toYaml .Values.webui.ingress.annotations | trimSuffix "\n" | indent 4 }}
+{{- end }}
+spec:
+  rules:
+  - host: {{ .Values.webui.ingress.host }}
+    http:
+      paths:
+      - path: /
+        pathType: Prefix
+        backend:
+          service:
+            name: ollama-webui-service
+            port:
+              number: {{ .Values.webui.servicePort }}
+{{- end }}

+ 12 - 0
kubernetes/helm/templates/webui-pvc.yaml

@@ -0,0 +1,12 @@
+apiVersion: v1
+kind: PersistentVolumeClaim
+metadata:
+  labels:
+    app: ollama-webui
+  name: ollama-webui-pvc
+  namespace: {{ .Values.namespace }}  
+spec:
+  accessModes: [ "ReadWriteOnce" ]
+  resources:
+    requests:
+      storage: {{ .Values.webui.volumeSize }}

+ 15 - 0
kubernetes/helm/templates/webui-service.yaml

@@ -0,0 +1,15 @@
+apiVersion: v1
+kind: Service
+metadata:
+  name: ollama-webui-service
+  namespace: {{ .Values.namespace }}
+spec:
+  type: {{ .Values.webui.service.type }} # Default: NodePort  # Use LoadBalancer if you're on a cloud that supports it
+  selector:
+    app: ollama-webui
+  ports:
+    - protocol: TCP
+      port: {{ .Values.webui.servicePort }}
+      targetPort: {{ .Values.webui.servicePort }}
+      # If using NodePort, you can optionally specify the nodePort:
+      # nodePort: 30000

+ 38 - 0
kubernetes/helm/values.yaml

@@ -0,0 +1,38 @@
+namespace: ollama-namespace
+
+ollama:
+  replicaCount: 1
+  image: ollama/ollama:latest
+  servicePort: 11434
+  resources:
+    limits:
+      cpu: "2000m"
+      memory: "2Gi"
+      nvidia.com/gpu: "0"
+  volumeSize: 1Gi
+  nodeSelector: {}
+  tolerations: []
+  service:
+    type: ClusterIP
+  gpu:
+    enabled: false
+
+webui:
+  replicaCount: 1
+  image: ghcr.io/ollama-webui/ollama-webui:main
+  servicePort: 8080
+  resources:
+    limits:
+      cpu: "500m"
+      memory: "500Mi"
+  ingress:
+    enabled: true
+    annotations:
+    # Use appropriate annotations for your Ingress controller, e.g., for NGINX:
+      # nginx.ingress.kubernetes.io/rewrite-target: /
+    host: ollama.minikube.local
+  volumeSize: 1Gi
+  nodeSelector: {}
+  tolerations: []
+  service:
+    type: NodePort

+ 4 - 0
kubernetes/manifest/base/ollama-namespace.yaml

@@ -0,0 +1,4 @@
+apiVersion: v1
+kind: Namespace
+metadata:
+  name: ollama-namespace

+ 12 - 0
kubernetes/manifest/base/ollama-service.yaml

@@ -0,0 +1,12 @@
+apiVersion: v1
+kind: Service
+metadata:
+  name: ollama-service
+  namespace: ollama-namespace
+spec:
+  selector:
+    app: ollama
+  ports:
+  - protocol: TCP
+    port: 11434
+    targetPort: 11434

+ 37 - 0
kubernetes/manifest/base/ollama-statefulset.yaml

@@ -0,0 +1,37 @@
+apiVersion: apps/v1
+kind: StatefulSet
+metadata:
+  name: ollama
+  namespace: ollama-namespace
+spec:
+  serviceName: "ollama"
+  replicas: 1
+  selector:
+    matchLabels:
+      app: ollama
+  template:
+    metadata:
+      labels:
+        app: ollama
+    spec:
+      containers:
+      - name: ollama
+        image: ollama/ollama:latest
+        ports:
+        - containerPort: 11434
+        resources:
+          limits:
+            cpu: "2000m"
+            memory: "2Gi"
+        volumeMounts:
+        - name: ollama-volume
+          mountPath: /root/.ollama
+        tty: true
+  volumeClaimTemplates:
+  - metadata:
+      name: ollama-volume
+    spec:
+      accessModes: [ "ReadWriteOnce" ]
+      resources:
+        requests:
+          storage: 1Gi

+ 28 - 0
kubernetes/manifest/base/webui-deployment.yaml

@@ -0,0 +1,28 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: ollama-webui-deployment
+  namespace: ollama-namespace
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: ollama-webui
+  template:
+    metadata:
+      labels:
+        app: ollama-webui
+    spec:
+      containers:
+      - name: ollama-webui
+        image: ghcr.io/ollama-webui/ollama-webui:main
+        ports:
+        - containerPort: 8080
+        resources:
+          limits:
+            cpu: "500m"
+            memory: "500Mi"
+        env:
+        - name: OLLAMA_API_BASE_URL
+          value: "http://ollama-service.ollama-namespace.svc.cluster.local:11434/api"
+        tty: true

+ 20 - 0
kubernetes/manifest/base/webui-ingress.yaml

@@ -0,0 +1,20 @@
+apiVersion: networking.k8s.io/v1
+kind: Ingress
+metadata:
+  name: ollama-webui-ingress
+  namespace: ollama-namespace
+  #annotations:
+    # Use appropriate annotations for your Ingress controller, e.g., for NGINX:
+    # nginx.ingress.kubernetes.io/rewrite-target: /
+spec:
+  rules:
+  - host: ollama.minikube.local
+    http:
+      paths:
+      - path: /
+        pathType: Prefix
+        backend:
+          service:
+            name: ollama-webui-service
+            port:
+              number: 8080

+ 15 - 0
kubernetes/manifest/base/webui-service.yaml

@@ -0,0 +1,15 @@
+apiVersion: v1
+kind: Service
+metadata:
+  name: ollama-webui-service
+  namespace: ollama-namespace
+spec:
+  type: NodePort  # Use LoadBalancer if you're on a cloud that supports it
+  selector:
+    app: ollama-webui
+  ports:
+    - protocol: TCP
+      port: 8080
+      targetPort: 8080
+      # If using NodePort, you can optionally specify the nodePort:
+      # nodePort: 30000

+ 12 - 0
kubernetes/manifest/kustomization.yaml

@@ -0,0 +1,12 @@
+resources:
+- base/ollama-namespace.yaml
+- base/ollama-service.yaml
+- base/ollama-statefulset.yaml
+- base/webui-deployment.yaml
+- base/webui-service.yaml
+- base/webui-ingress.yaml
+
+apiVersion: kustomize.config.k8s.io/v1beta1
+kind: Kustomization
+patches:
+- path: patches/ollama-statefulset-gpu.yaml

+ 17 - 0
kubernetes/manifest/patches/ollama-statefulset-gpu.yaml

@@ -0,0 +1,17 @@
+apiVersion: apps/v1
+kind: StatefulSet
+metadata:
+  name: ollama
+  namespace: ollama-namespace
+spec:
+  selector:
+    matchLabels:
+      app: ollama
+  serviceName: "ollama"
+  template:
+    spec:
+      containers:
+      - name: ollama
+        resources:
+          limits:
+            nvidia.com/gpu: "1"

+ 237 - 0
run-compose.sh

@@ -0,0 +1,237 @@
+#!/bin/bash
+
+# Define color and formatting codes
+BOLD='\033[1m'
+GREEN='\033[1;32m'
+WHITE='\033[1;37m'
+RED='\033[0;31m'
+NC='\033[0m' # No Color
+# Unicode character for tick mark
+TICK='\u2713'
+
+# Detect GPU driver
+get_gpu_driver() {
+    # Detect NVIDIA GPUs
+    if lspci | grep -i nvidia >/dev/null; then
+        echo "nvidia"
+        return
+    fi
+
+    # Detect AMD GPUs (including GCN architecture check for amdgpu vs radeon)
+    if lspci | grep -i amd >/dev/null; then
+        # List of known GCN and later architecture cards
+        # This is a simplified list, and in a real-world scenario, you'd want a more comprehensive one
+        local gcn_and_later=("Radeon HD 7000" "Radeon HD 8000" "Radeon R5" "Radeon R7" "Radeon R9" "Radeon RX")
+
+        # Get GPU information
+        local gpu_info=$(lspci | grep -i 'vga.*amd')
+
+        for model in "${gcn_and_later[@]}"; do
+            if echo "$gpu_info" | grep -iq "$model"; then
+                echo "amdgpu"
+                return
+            fi
+        done
+
+        # Default to radeon if no GCN or later architecture is detected
+        echo "radeon"
+        return
+    fi
+
+    # Detect Intel GPUs
+    if lspci | grep -i intel >/dev/null; then
+        echo "i915"
+        return
+    fi
+
+    # If no known GPU is detected
+    echo "Unknown or unsupported GPU driver"
+    exit 1
+}
+
+# Function for rolling animation
+show_loading() {
+    local spin='-\|/'
+    local i=0
+
+    printf " "
+
+    while kill -0 $1 2>/dev/null; do
+        i=$(( (i+1) %4 ))
+        printf "\b${spin:$i:1}"
+        sleep .1
+    done
+
+    # Replace the spinner with a tick
+    printf "\b${GREEN}${TICK}${NC}"
+}
+
+# Usage information
+usage() {
+    echo "Usage: $0 [OPTIONS]"
+    echo "Options:"
+    echo "  --enable-gpu[count=COUNT]  Enable GPU support with the specified count."
+    echo "  --enable-api[port=PORT]    Enable API and expose it on the specified port."
+    echo "  --webui[port=PORT]         Set the port for the web user interface."
+    echo "  --data[folder=PATH]        Bind mount for ollama data folder (by default will create the 'ollama' volume)."
+    echo "  --build                    Build the docker image before running the compose project."
+    echo "  --drop                     Drop the compose project."
+    echo "  -q, --quiet                Run script in headless mode."
+    echo "  -h, --help                 Show this help message."
+    echo ""
+    echo "Examples:"
+    echo "  $0 --drop"
+    echo "  $0 --enable-gpu[count=1]"
+    echo "  $0 --enable-api[port=11435]"
+    echo "  $0 --enable-gpu[count=1] --enable-api[port=12345] --webui[port=3000]"
+    echo "  $0 --enable-gpu[count=1] --enable-api[port=12345] --webui[port=3000] --data[folder=./ollama-data]"
+    echo "  $0 --enable-gpu[count=1] --enable-api[port=12345] --webui[port=3000] --data[folder=./ollama-data] --build"
+    echo ""
+    echo "This script configures and runs a docker-compose setup with optional GPU support, API exposure, and web UI configuration."
+    echo "About the gpu to use, the script automatically detects it using the "lspci" command."
+    echo "In this case the gpu detected is: $(get_gpu_driver)"
+}
+
+# Default values
+gpu_count=1
+api_port=11435
+webui_port=3000
+headless=false
+build_image=false
+kill_compose=false
+
+# Function to extract value from the parameter
+extract_value() {
+    echo "$1" | sed -E 's/.*\[.*=(.*)\].*/\1/; t; s/.*//'
+}
+
+# Parse arguments
+while [[ $# -gt 0 ]]; do
+    key="$1"
+
+    case $key in
+        --enable-gpu*)
+            enable_gpu=true
+            value=$(extract_value "$key")
+            gpu_count=${value:-1}
+            ;;
+        --enable-api*)
+            enable_api=true
+            value=$(extract_value "$key")
+            api_port=${value:-11435}
+            ;;
+        --webui*)
+            value=$(extract_value "$key")
+            webui_port=${value:-3000}
+            ;;
+        --data*)
+            value=$(extract_value "$key")
+            data_dir=${value:-"./ollama-data"}
+            ;;
+        --drop)
+            kill_compose=true
+            ;;
+        --build)
+            build_image=true
+            ;;
+        -q|--quiet)
+            headless=true
+            ;;
+        -h|--help)
+            usage
+            exit
+            ;;
+        *)
+            # Unknown option
+            echo "Unknown option: $key"
+            usage
+            exit 1
+            ;;
+    esac
+    shift # past argument or value
+done
+
+if [[ $kill_compose == true ]]; then
+    docker compose down --remove-orphans
+    echo -e "${GREEN}${BOLD}Compose project dropped successfully.${NC}"
+    exit
+else
+    DEFAULT_COMPOSE_COMMAND="docker compose -f docker-compose.yaml"
+    if [[ $enable_gpu == true ]]; then
+        # Validate and process command-line arguments
+        if [[ -n $gpu_count ]]; then
+            if ! [[ $gpu_count =~ ^[0-9]+$ ]]; then
+                echo "Invalid GPU count: $gpu_count"
+                exit 1
+            fi
+            echo "Enabling GPU with $gpu_count GPUs"
+            # Add your GPU allocation logic here
+            export OLLAMA_GPU_DRIVER=$(get_gpu_driver)
+            export OLLAMA_GPU_COUNT=$gpu_count # Set OLLAMA_GPU_COUNT environment variable
+        fi
+        DEFAULT_COMPOSE_COMMAND+=" -f docker-compose.gpu.yaml"
+    fi
+    if [[ $enable_api == true ]]; then
+        DEFAULT_COMPOSE_COMMAND+=" -f docker-compose.api.yaml"
+        if [[ -n $api_port ]]; then
+            export OLLAMA_WEBAPI_PORT=$api_port # Set OLLAMA_WEBAPI_PORT environment variable
+        fi
+    fi
+    if [[ -n $data_dir ]]; then
+        DEFAULT_COMPOSE_COMMAND+=" -f docker-compose.data.yaml"
+        export OLLAMA_DATA_DIR=$data_dir # Set OLLAMA_DATA_DIR environment variable
+    fi
+    DEFAULT_COMPOSE_COMMAND+=" up -d"
+    DEFAULT_COMPOSE_COMMAND+=" --remove-orphans"
+    DEFAULT_COMPOSE_COMMAND+=" --force-recreate"
+    if [[ $build_image == true ]]; then
+        DEFAULT_COMPOSE_COMMAND+=" --build"
+    fi
+fi
+
+# Recap of environment variables
+echo
+echo -e "${WHITE}${BOLD}Current Setup:${NC}"
+echo -e "   ${GREEN}${BOLD}GPU Driver:${NC} ${OLLAMA_GPU_DRIVER:-Not Enabled}"
+echo -e "   ${GREEN}${BOLD}GPU Count:${NC} ${OLLAMA_GPU_COUNT:-Not Enabled}"
+echo -e "   ${GREEN}${BOLD}WebAPI Port:${NC} ${OLLAMA_WEBAPI_PORT:-Not Enabled}"
+echo -e "   ${GREEN}${BOLD}Data Folder:${NC} ${data_dir:-Using ollama volume}"
+echo -e "   ${GREEN}${BOLD}WebUI Port:${NC} $webui_port"
+echo
+
+if [[ $headless == true ]]; then
+    echo -ne "${WHITE}${BOLD}Running in headless mode... ${NC}"
+    choice="y"
+else
+    # Ask for user acceptance
+    echo -ne "${WHITE}${BOLD}Do you want to proceed with current setup? (Y/n): ${NC}"
+    read -n1 -s choice
+fi
+
+echo
+
+if [[ $choice == "" || $choice == "y" ]]; then
+    # Execute the command with the current user
+    eval "$DEFAULT_COMPOSE_COMMAND" &
+
+    # Capture the background process PID
+    PID=$!
+
+    # Display the loading animation
+    #show_loading $PID
+
+    # Wait for the command to finish
+    wait $PID
+
+    echo
+    # Check exit status
+    if [ $? -eq 0 ]; then
+        echo -e "${GREEN}${BOLD}Compose project started successfully.${NC}"
+    else
+        echo -e "${RED}${BOLD}There was an error starting the compose project.${NC}"
+    fi
+else
+    echo "Aborted."
+fi
+
+echo

+ 1 - 1
src/app.css

@@ -16,7 +16,7 @@ html {
 
 
 code {
 code {
 	/* white-space-collapse: preserve !important; */
 	/* white-space-collapse: preserve !important; */
-	white-space: pre;
+	overflow-x: auto;
 	width: auto;
 	width: auto;
 }
 }
 
 

+ 18 - 0
src/lib/components/chat/MessageInput.svelte

@@ -298,6 +298,24 @@
 									submitPrompt(prompt);
 									submitPrompt(prompt);
 								}
 								}
 							}}
 							}}
+							on:keydown={(e) => {
+								if (prompt === '' && e.key == 'ArrowUp') {
+									e.preventDefault();
+
+									const userMessageElement = [
+										...document.getElementsByClassName('user-message')
+									]?.at(-1);
+
+									const editButton = [
+										...document.getElementsByClassName('edit-user-message-button')
+									]?.at(-1);
+
+									console.log(userMessageElement);
+
+									userMessageElement.scrollIntoView({ block: 'center' });
+									editButton?.click();
+								}
+							}}
 							rows="1"
 							rows="1"
 							on:input={(e) => {
 							on:input={(e) => {
 								e.target.style.height = '';
 								e.target.style.height = '';

+ 1 - 0
src/lib/components/chat/Messages/ResponseMessage.svelte

@@ -88,6 +88,7 @@
 				let code = block.querySelector('code');
 				let code = block.querySelector('code');
 				code.style.borderTopRightRadius = 0;
 				code.style.borderTopRightRadius = 0;
 				code.style.borderTopLeftRadius = 0;
 				code.style.borderTopLeftRadius = 0;
+				code.style.whiteSpace = 'pre';
 
 
 				let topBarDiv = document.createElement('div');
 				let topBarDiv = document.createElement('div');
 				topBarDiv.style.backgroundColor = '#202123';
 				topBarDiv.style.backgroundColor = '#202123';

+ 6 - 2
src/lib/components/chat/Messages/UserMessage.svelte

@@ -24,6 +24,8 @@
 
 
 		editElement.style.height = '';
 		editElement.style.height = '';
 		editElement.style.height = `${editElement.scrollHeight}px`;
 		editElement.style.height = `${editElement.scrollHeight}px`;
+
+		editElement?.focus();
 	};
 	};
 
 
 	const editMessageConfirmHandler = async () => {
 	const editMessageConfirmHandler = async () => {
@@ -43,7 +45,9 @@
 	<ProfileImage src={user?.profile_image_url ?? '/user.png'} />
 	<ProfileImage src={user?.profile_image_url ?? '/user.png'} />
 
 
 	<div class="w-full overflow-hidden">
 	<div class="w-full overflow-hidden">
-		<Name>You</Name>
+		<div class="user-message">
+			<Name>You</Name>
+		</div>
 
 
 		<div
 		<div
 			class="prose chat-{message.role} w-full max-w-full dark:prose-invert prose-headings:my-0 prose-p:my-0 prose-p:-mb-4 prose-pre:my-0 prose-table:my-0 prose-blockquote:my-0 prose-img:my-0 prose-ul:-my-4 prose-ol:-my-4 prose-li:-my-3 prose-ul:-mb-6 prose-ol:-mb-6 prose-li:-mb-4 whitespace-pre-line"
 			class="prose chat-{message.role} w-full max-w-full dark:prose-invert prose-headings:my-0 prose-p:my-0 prose-p:-mb-4 prose-pre:my-0 prose-table:my-0 prose-blockquote:my-0 prose-img:my-0 prose-ul:-my-4 prose-ol:-my-4 prose-li:-my-3 prose-ul:-mb-6 prose-ol:-mb-6 prose-li:-mb-4 whitespace-pre-line"
@@ -145,7 +149,7 @@
 						{/if}
 						{/if}
 
 
 						<button
 						<button
-							class="invisible group-hover:visible p-1 rounded dark:hover:bg-gray-800 transition"
+							class="invisible group-hover:visible p-1 rounded dark:hover:bg-gray-800 transition edit-user-message-button"
 							on:click={() => {
 							on:click={() => {
 								editMessageHandler();
 								editMessageHandler();
 							}}
 							}}