nicer image upload

docker-compose.yml

@@ -8,6 +8,7 @@ services:
   vllm:
     image: vllm/vllm-openai:latest
     command: "--model Qwen/Qwen2-VL-2B-Instruct-GPTQ-Int4 --quantization gptq "
+    #command: "--model HuggingFaceTB/SmolVLM-Instruct --max_model_len 4098"
     volumes:
       - ~/.cache/huggingface:/root/.cache/huggingface
     ports:

penparse/webui/tasks.py

@@ -1,5 +1,6 @@
 import base64
 import litellm
+import openai
 
 from loguru import logger
 from celery import shared_task
@@ -26,7 +27,6 @@ Please include whitespace and formatting for headings too.
 """
 
 
-
 @shared_task
 def process_memo(memo_id: str):
     """Run OCR on a memo and store the output"""
@@ -70,15 +70,24 @@ def process_memo(memo_id: str):
     litellm.api_base = settings.OPENAI_API_BASE  # os.environ.get("OPENAI_API_BASE")
     litellm.api_key = settings.OPENAI_API_KEY
 
-    response = litellm.completion(
-        model=settings.OPENAI_MODEL, #os.getenv("MODEL", "openai/gpt-4o"),
-        messages=[message],
-        temperature=0.01
-    )
+    try:
+        response = litellm.completion(
+            model=settings.OPENAI_MODEL, #os.getenv("MODEL", "openai/gpt-4o"),
+            messages=[message],
+            temperature=0.01
+        )
 
-    response.choices[0].message["content"]
+        response.choices[0].message["content"]
 
-    with transaction.atomic():
-        memo.content = response.choices[0].message["content"]
-        memo.status = MemoStatus.Done
-        memo.save()
+        with transaction.atomic():
+            memo.content = response.choices[0].message["content"]
+            memo.status = MemoStatus.Done
+            memo.model_name = settings.OPENAI_MODEL
+            memo.save()
+    except openai.OpenAIError as e:
+
+        with transaction.atomic():
+            memo.status = MemoStatus.Error
+            memo.error_message = e.__repr__()
+            memo.save()
+            logger.error(e)

penparse/webui/templates/dashboard.html

@@ -90,29 +90,91 @@
     action="{% url 'upload_document' %}"
     method="post"
     enctype="multipart/form-data"
+    id="upload-form"
   >
     {% csrf_token %}
-    <div class="mb-4">
+    <div 
+      id="drop-area"
+      class="border-2 border-dashed border-gray-300 rounded-lg p-8 mb-4 transition-colors duration-300 ease-in-out hover:border-blue-500 cursor-pointer"
+    >
       <input
         type="file"
         name="document"
-        id="document"
+        id="fileElem"
         class="hidden"
         accept=".png,.jpg,.jpeg"
       />
-      <label
-        for="document"
-        class="bg-blue-500 text-white px-6 py-3 rounded-lg hover:bg-blue-600 transition duration-300 cursor-pointer inline-block"
-      >
-        Choose File
+      <label for="fileElem" class="cursor-pointer">
+        <div class="text-gray-500 mb-2">
+          <svg class="w-12 h-12 mx-auto mb-2" fill="none" stroke="currentColor" viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg">
+            <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M7 16a4 4 0 01-.88-7.903A5 5 0 1115.9 6L16 6a5 5 0 011 9.9M15 13l-3-3m0 0l-3 3m3-3v12"></path>
+          </svg>
+          <p class="text-lg font-medium">Drag & drop your file here</p>
+          <p class="text-sm">or click to select a file</p>
+        </div>
       </label>
     </div>
+    <p id="file-name" class="text-gray-600 mb-4 hidden"></p>
     <button
       type="submit"
-      class="bg-blue-500 text-white px-6 py-3 rounded-lg hover:bg-blue-600 transition duration-300"
+      id="upload-button"
+      class="bg-blue-500 text-white px-6 py-3 rounded-lg hover:bg-blue-600 transition duration-300 hidden"
     >
       Upload Document
     </button>
   </form>
 </section>
+
+<script>
+  const dropArea = document.getElementById('drop-area');
+  const fileInput = document.getElementById('fileElem');
+  const fileName = document.getElementById('file-name');
+  const uploadButton = document.getElementById('upload-button');
+
+  ['dragenter', 'dragover', 'dragleave', 'drop'].forEach(eventName => {
+    dropArea.addEventListener(eventName, preventDefaults, false);
+  });
+
+  function preventDefaults(e) {
+    e.preventDefault();
+    e.stopPropagation();
+  }
+
+  ['dragenter', 'dragover'].forEach(eventName => {
+    dropArea.addEventListener(eventName, highlight, false);
+  });
+
+  ['dragleave', 'drop'].forEach(eventName => {
+    dropArea.addEventListener(eventName, unhighlight, false);
+  });
+
+  function highlight() {
+    dropArea.classList.add('border-blue-500', 'bg-blue-50');
+  }
+
+  function unhighlight() {
+    dropArea.classList.remove('border-blue-500', 'bg-blue-50');
+  }
+
+  dropArea.addEventListener('drop', handleDrop, false);
+
+  function handleDrop(e) {
+    const dt = e.dataTransfer;
+    const files = dt.files;
+    handleFiles(files);
+  }
+
+  fileInput.addEventListener('change', function() {
+    handleFiles(this.files);
+  });
+
+  function handleFiles(files) {
+    if (files.length > 0) {
+      fileName.textContent = `Selected file: ${files[0].name}`;
+      fileName.classList.remove('hidden');
+      uploadButton.classList.remove('hidden');
+    }
+  }
+</script>
+
 {% endblock %}

penparse/webui/test/test_tasks.py

@@ -0,0 +1,104 @@
+import pytest
+from django.core.files.base import ContentFile
+from django.core.files.storage import default_storage
+from unittest.mock import patch, MagicMock
+from penparse.webui.tasks import process_memo
+from penparse.webui.models import ImageMemo, MemoStatus
+
+
+@pytest.fixture
+def sample_image_memo(db):
+    memo = ImageMemo.objects.create(
+        status=MemoStatus.Pending, image_mimetype="image/jpeg"
+    )
+    memo.image.save("test_image.jpg", ContentFile(b"fake image content"))
+    return memo
+
+
+@pytest.mark.django_db
+def test_process_memo_success(sample_image_memo):
+    with patch("penparse.webui.tasks.litellm") as mock_litellm:
+        mock_response = MagicMock()
+        mock_response.choices[0].message = {"content": "Transcribed content"}
+        mock_litellm.completion.return_value = mock_response
+
+        process_memo(sample_image_memo.id)
+
+        processed_memo = ImageMemo.objects.get(id=sample_image_memo.id)
+        assert processed_memo.status == MemoStatus.Done
+        assert processed_memo.content == "Transcribed content"
+        assert processed_memo.error_message == ""
+
+
+@pytest.mark.django_db
+def test_process_memo_missing_image(sample_image_memo):
+    default_storage.delete(sample_image_memo.image.name)
+
+    process_memo(sample_image_memo.id)
+
+    processed_memo = ImageMemo.objects.get(id=sample_image_memo.id)
+    assert processed_memo.status == MemoStatus.Error
+    assert "Image file" in processed_memo.error_message
+
+
+@pytest.mark.django_db
+def test_process_memo_api_error(sample_image_memo):
+    with patch("penparse.webui.tasks.litellm") as mock_litellm:
+        mock_litellm.completion.side_effect = mock_litellm.APIError("API Error")
+
+        process_memo(sample_image_memo.id)
+
+        processed_memo = ImageMemo.objects.get(id=sample_image_memo.id)
+        assert processed_memo.status == MemoStatus.Error
+        assert "API Error" in processed_memo.error_message
+
+
+@pytest.mark.django_db
+def test_process_memo_sets_model_name(sample_image_memo):
+    with (
+        patch("penparse.webui.tasks.litellm") as mock_litellm,
+        patch("penparse.webui.tasks.settings") as mock_settings,
+    ):
+        mock_response = MagicMock()
+        mock_response.choices[0].message = {"content": "Transcribed content"}
+        mock_litellm.completion.return_value = mock_response
+        mock_settings.OPENAI_MODEL = "test-model"
+
+        process_memo(sample_image_memo.id)
+
+        processed_memo = ImageMemo.objects.get(id=sample_image_memo.id)
+        assert processed_memo.model_name == "test-model"
+
+
+@pytest.mark.django_db
+def test_process_memo_uses_correct_api_settings(sample_image_memo):
+    with (
+        patch("penparse.webui.tasks.litellm") as mock_litellm,
+        patch("penparse.webui.tasks.settings") as mock_settings,
+    ):
+        mock_response = MagicMock()
+        mock_response.choices[0].message = {"content": "Transcribed content"}
+        mock_litellm.completion.return_value = mock_response
+        mock_settings.OPENAI_API_BASE = "https://test-api-base.com"
+        mock_settings.OPENAI_API_KEY = "test-api-key"
+        mock_settings.OPENAI_MODEL = "test-model"
+
+        process_memo(sample_image_memo.id)
+
+        assert mock_litellm.api_base == "https://test-api-base.com"
+        assert mock_litellm.api_key == "test-api-key"
+        mock_litellm.completion.assert_called_once_with(
+            model="test-model",
+            messages=pytest.approx(
+                [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": pytest.ANY},
+                            {"type": "image_url", "image_url": {"url": pytest.ANY}},
+                        ],
+                    }
+                ]
+            ),
+            temperature=0.01,
+        )