From d81319c8ec303bbc1a2c51e702a5f90cbc4c387a Mon Sep 17 00:00:00 2001
From: Keannu Bernasol <noehbernasol0@gmail.com>
Date: Sat, 7 Dec 2024 14:03:17 +0800
Subject: [PATCH] Update Ollama API to use different model

---
 .env.sample                                            |  4 +++-
 .../config/management/commands/start_watcher.py        | 10 +++++-----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/.env.sample b/.env.sample
index 69f9163..72470fd 100644
--- a/.env.sample
+++ b/.env.sample
@@ -23,6 +23,8 @@ ADMIN_PASSWORD = ''
 TEST_DATA = "True"
 
 # Ollama for Categorization 
-OLLAMA_URL = ""
+OLLAMA_URL = "localhost:11434"
+OLLAMA_USE_AUTH = "False"
+OLLAMA_MODEL = "knoopx/mobile-vlm:3b-fp16"
 OLLAMA_USERNAME = ""
 OLLAMA_PASSWORD = ""
\ No newline at end of file
diff --git a/docmanager_backend/config/management/commands/start_watcher.py b/docmanager_backend/config/management/commands/start_watcher.py
index 0eb0998..0bc3e79 100644
--- a/docmanager_backend/config/management/commands/start_watcher.py
+++ b/docmanager_backend/config/management/commands/start_watcher.py
@@ -1,7 +1,7 @@
 from ollama import ChatResponse
 import base64
 import httpx
-from django.core.management.base import BaseCommand, CommandError
+from django.core.management.base import BaseCommand
 
 from io import BytesIO
 from documents.models import Document
@@ -89,7 +89,7 @@ class PDFHandler(FileSystemEventHandler):
                         client = Client(
                             host=get_secret("OLLAMA_URL"),
                             auth=httpx.BasicAuth(
-                                username=get_secret("OLLAMA_USERNAME"), password=get_secret("OLLAMA_PASSWORD"))
+                                username=get_secret("OLLAMA_USERNAME"), password=get_secret("OLLAMA_PASSWORD")) if get_secret("OLLAMA_USE_AUTH") else None
                         )
 
                         encoded_image = base64.b64encode(
@@ -111,15 +111,15 @@ class PDFHandler(FileSystemEventHandler):
                             """
 
                             response: ChatResponse = client.chat(
-                                model="llama3.2-vision",
+                                model=get_secret("OLLAMA_MODEL"),
                                 messages=[
                                     {"role": "user", "content": content,
                                         "images": [encoded_image]},
                                 ],
                             )
 
-                            document_type = response["message"]["content"].split(":")[
-                                0].replace("*", "").replace(".", "")
+                            document_type = response["message"]["content"].replace(
+                                "*", "").replace(".", "")
 
                             # A few safety checks if the model does not follow through with output instructions
                             if len(document_type) > 16: