Update app.py

app.py

@@ -9,9 +9,9 @@ import requests
 import numpy as np
 import asyncio
 from typing import List, Dict, Any, Optional, Union
-from fastapi import FastAPI, HTTPException, Depends, Request, File, UploadFile
+from fastapi import FastAPI, HTTPException, Depends, Request
 from fastapi.concurrency import run_in_threadpool
-from fastapi.responses import HTMLResponse, JSONResponse
+from fastapi.responses import HTMLResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 from pydantic import BaseModel, Field
@@ -345,14 +345,23 @@ def classify_image(image_data):
             "nsfw_score": 0.0
         }
 
-def process_content_item(item: Union[str, Dict], text_model: str = "gemma") -> Dict:
+def process_content_item(item: Union[str, TextContent, ImageContent], text_model: str = "gemma") -> Dict:
+    work_item = {}
     if isinstance(item, str):
-        item = {"type": "text", "text": item}
+        work_item = {"type": "text", "text": item}
+    elif isinstance(item, (TextContent, ImageContent)):
+        work_item = item.model_dump()
+    else:
+        # This case should ideally not be hit with proper Pydantic validation
+        return {
+            "flagged": False,
+            "error": f"Unsupported item type: {type(item).__name__}"
+        }
 
-    content_type = item.get("type")
+    content_type = work_item.get("type")
 
     if content_type == "text":
-        text = item.get("text", "")
+        text = work_item.get("text", "")
         if text_model == "gemma":
             gemma_result = classify_text_with_gemma(text)
             flagged = gemma_result["classification"] == "u"
@@ -365,12 +374,7 @@ def process_content_item(item: Union[str, Dict], text_model: str = "gemma") -> D
                 "violence": 0.9 if flagged else 0.1, "violence/graphic": 0.9 if flagged else 0.1,
                 "nsfw": 0.1,
             }
-            return {
-                "flagged": flagged,
-                "categories": {k: (v > 0.5) for k, v in scores.items()},
-                "category_scores": scores,
-                "text": text,
-            }
+            return {"flagged": flagged, "categories": {k: (v > 0.5) for k, v in scores.items()}, "category_scores": scores, "text": text}
         elif text_model == "detoxify":
             d = classify_text_with_detoxify(text)
             scores = {
@@ -381,12 +385,7 @@ def process_content_item(item: Union[str, Dict], text_model: str = "gemma") -> D
                 "violence": d["category_scores"].get("threat", 0.1), "violence/graphic": d["category_scores"].get("threat", 0.1),
                 "nsfw": d["category_scores"].get("sexual_explicit", 0.1),
             }
-            return {
-                "flagged": d["flagged"],
-                "categories": {k: (v > 0.5) for k, v in scores.items()},
-                "category_scores": scores,
-                "text": text,
-            }
+            return {"flagged": d["flagged"], "categories": {k: (v > 0.5) for k, v in scores.items()}, "category_scores": scores, "text": text}
         elif text_model == "both":
             gemma_result = classify_text_with_gemma(text)
             detoxify_result = classify_text_with_detoxify(text)
@@ -405,17 +404,12 @@ def process_content_item(item: Union[str, Dict], text_model: str = "gemma") -> D
                 "violence/graphic": max(0.9 if gemma_result["classification"] == "u" else 0.1, detoxify_result["category_scores"].get("threat", 0.1)),
                 "nsfw": detoxify_result["category_scores"].get("sexual_explicit", 0.1),
             }
-            return {
-                "flagged": flagged,
-                "categories": {k: (v > 0.5) for k, v in scores.items()},
-                "category_scores": scores,
-                "text": text,
-            }
+            return {"flagged": flagged, "categories": {k: (v > 0.5) for k, v in scores.items()}, "category_scores": scores, "text": text}
 
     elif content_type == "image":
         image_data = None
-        image_url = item.get("url")
-        image_base64 = item.get("base64")
+        image_url = work_item.get("url")
+        image_base64 = work_item.get("base64")
         
         if image_url:
             try:
@@ -449,7 +443,7 @@ def process_content_item(item: Union[str, Dict], text_model: str = "gemma") -> D
                 "categories": {k: (v > 0.5) for k, v in scores.items()},
                 "category_scores": scores,
                 "image_url": image_url,
-                "image_base64": image_base64[:50] + "..." if isinstance(image_base64, str) and len(image_base64) > 50 else None,
+                "image_base64": work_item.get("base64"),
             }
     
     default_scores = {
@@ -462,7 +456,7 @@ def process_content_item(item: Union[str, Dict], text_model: str = "gemma") -> D
         "flagged": False,
         "categories": {k: False for k in default_scores},
         "category_scores": default_scores,
-        "error": f"Invalid or unprocessable item: {item}"
+        "error": f"Invalid or unprocessable item: {work_item}"
     }
 
 def get_api_key(request: Request):
@@ -508,8 +502,8 @@ async def moderate_content(
             for item in items:
                 if isinstance(item, str):
                     total_tokens += count_tokens(item)
-                elif isinstance(item, dict) and item.get("type") == "text":
-                    total_tokens += count_tokens(item.get("text", ""))
+                elif isinstance(item, TextContent):
+                    total_tokens += count_tokens(item.text)
         else:
             raise HTTPException(status_code=400, detail="Invalid input format")
         
@@ -1032,21 +1026,13 @@ with open("templates/index.html", "w", encoding='utf-8') as f:
             }
         });
         
-        analyzeTextBtn.addEventListener('click', async () => {
-            const text = textInput.value.trim();
-            if (!text) {
-                showNotification('Please enter text to analyze', 'error');
-                return;
-            }
-            
+        async function analyze(payload) {
             const apiKey = apiKeyInput.value.trim();
             if (!apiKey) {
                 showNotification('Please enter your API key', 'error');
                 return;
             }
             
-            const textModel = textModelSelect.value;
-            
             showLoading(true);
             try {
                 const response = await fetch('/v1/moderations', {
@@ -1055,17 +1041,14 @@ with open("templates/index.html", "w", encoding='utf-8') as f:
                         'Content-Type': 'application/json',
                         'Authorization': `Bearer ${apiKey}`
                     },
-                    body: JSON.stringify({
-                        input: text,
-                        model: textModel
-                    })
+                    body: JSON.stringify(payload)
                 });
-                
+
                 if (!response.ok) {
                     const errorData = await response.json();
-                    throw new Error(errorData.detail || 'An error occurred');
+                    throw new Error(errorData.detail || `HTTP error! status: ${response.status}`);
                 }
-                
+
                 const data = await response.json();
                 displayResults(data.results);
                 updateMetrics();
@@ -1074,158 +1057,90 @@ with open("templates/index.html", "w", encoding='utf-8') as f:
             } finally {
                 showLoading(false);
             }
+        }
+        
+        analyzeTextBtn.addEventListener('click', async () => {
+            const text = textInput.value.trim();
+            if (!text) {
+                showNotification('Please enter text to analyze', 'error');
+                return;
+            }
+            await analyze({
+                input: text,
+                model: textModelSelect.value
+            });
         });
         
         analyzeImageBtn.addEventListener('click', async () => {
             const url = imageUrl.value.trim();
-            const fileInput = document.querySelector('#imageUpload');
-            const file = fileInput.files[0];
-            
+            const file = imageUpload.files[0];
+
             if (!url && !file) {
                 showNotification('Please provide an image URL or upload an image', 'error');
                 return;
             }
-            
-            const apiKey = apiKeyInput.value.trim();
-            if (!apiKey) {
-                showNotification('Please enter your API key', 'error');
-                return;
-            }
-            
+
             let imageInput;
-            
             if (url) {
-                imageInput = {
-                    type: "image",
-                    url: url
-                };
+                imageInput = { type: "image", url: url };
             } else {
-                const reader = new FileReader();
-                const base64Promise = new Promise((resolve) => {
+                const base64 = await new Promise((resolve) => {
+                    const reader = new FileReader();
                     reader.onload = (event) => resolve(event.target.result);
+                    reader.readAsDataURL(file);
                 });
-                reader.readAsDataURL(file);
-                const base64 = await base64Promise;
-                
-                imageInput = {
-                    type: "image",
-                    base64: base64
-                };
-            }
-            
-            showLoading(true);
-            try {
-                const response = await fetch('/v1/moderations', {
-                    method: 'POST',
-                    headers: {
-                        'Content-Type': 'application/json',
-                        'Authorization': `Bearer ${apiKey}`
-                    },
-                    body: JSON.stringify({
-                        input: [imageInput],
-                        model: textModelSelect.value
-                    })
-                });
-                
-                if (!response.ok) {
-                    const errorData = await response.json();
-                    throw new Error(errorData.detail || 'An error occurred');
-                }
-                
-                const data = await response.json();
-                displayResults(data.results);
-                updateMetrics();
-            } catch (error) {
-                showNotification(`Error: ${error.message}`, 'error');
-            } finally {
-                showLoading(false);
+                imageInput = { type: "image", base64: base64 };
             }
+
+            await analyze({
+                input: [imageInput],
+                model: textModelSelect.value
+            });
         });
-        
+
         analyzeMixedBtn.addEventListener('click', async () => {
             const items = Array.from(mixedItemsContainer.querySelectorAll('.mixed-item'));
             if (items.length === 0) {
                 showNotification('Please add at least one item to analyze', 'error');
                 return;
             }
-            
-            const apiKey = apiKeyInput.value.trim();
-            if (!apiKey) {
-                showNotification('Please enter your API key', 'error');
-                return;
-            }
-            
-            const inputItems = [];
-            
-            for (const item of items) {
+
+            const inputPromises = items.map(async (item) => {
                 const type = item.querySelector('.item-type').value;
                 const contentDiv = item.querySelector('.item-content');
-                
+
                 if (type === 'text') {
-                    const textarea = contentDiv.querySelector('textarea');
-                    const text = textarea.value.trim();
-                    if (text) {
-                        inputItems.push({
-                            type: 'text',
-                            text: text
-                        });
-                    }
+                    const text = contentDiv.querySelector('textarea').value.trim();
+                    return text ? { type: 'text', text: text } : null;
                 } else {
-                    const urlInput = contentDiv.querySelector('input[type="text"]');
-                    const fileInput = contentDiv.querySelector('input[type="file"]');
-                    const preview = contentDiv.querySelector('.image-preview');
-                    const previewImg = contentDiv.querySelector('.image-preview img');
-                    
-                    const url = urlInput.value.trim();
-                    const file = fileInput.files[0];
+                    const url = contentDiv.querySelector('input[type="text"]').value.trim();
+                    const file = contentDiv.querySelector('input[type="file"]').files[0];
                     
                     if (url) {
-                        inputItems.push({
-                            type: 'image',
-                            url: url
-                        });
-                    } else if (file || !preview.classList.contains('hidden')) {
-                        const imgSrc = previewImg.src;
-                        inputItems.push({
-                            type: 'image',
-                            base64: imgSrc
+                        return { type: 'image', url: url };
+                    } else if (file) {
+                        const base64 = await new Promise((resolve) => {
+                           const reader = new FileReader();
+                           reader.onload = (event) => resolve(event.target.result);
+                           reader.readAsDataURL(file);
                         });
+                        return { type: 'image', base64: base64 };
                     }
+                    return null;
                 }
-            }
-            
+            });
+
+            const inputItems = (await Promise.all(inputPromises)).filter(Boolean);
+
             if (inputItems.length === 0) {
                 showNotification('Please add content to at least one item', 'error');
                 return;
             }
             
-            showLoading(true);
-            try {
-                const response = await fetch('/v1/moderations', {
-                    method: 'POST',
-                    headers: {
-                        'Content-Type': 'application/json',
-                        'Authorization': `Bearer ${apiKey}`
-                    },
-                    body: JSON.stringify({
-                        input: inputItems,
-                        model: textModelSelect.value
-                    })
-                });
-                
-                if (!response.ok) {
-                    const errorData = await response.json();
-                    throw new Error(errorData.detail || 'An error occurred');
-                }
-                
-                const data = await response.json();
-                displayResults(data.results);
-                updateMetrics();
-            } catch (error) {
-                showNotification(`Error: ${error.message}`, 'error');
-            } finally {
-                showLoading(false);
-            }
+            await analyze({
+                input: inputItems,
+                model: textModelSelect.value
+            });
         });
         
         function displayResults(results) {
@@ -1240,10 +1155,10 @@ with open("templates/index.html", "w", encoding='utf-8') as f:
                     'Content may contain inappropriate or harmful material.' : 
                     'Content appears to be safe and appropriate.';
                 
-                const categories = Object.entries(result.categories)
+                const categories = result.categories ? Object.entries(result.categories)
                     .filter(([_, value]) => value)
                     .map(([key, _]) => key.replace('/', ' '))
-                    .join(', ');
+                    .join(', ') : 'N/A';
                 
                 let contentPreview = '';
                 
@@ -1256,7 +1171,7 @@ with open("templates/index.html", "w", encoding='utf-8') as f:
                         <img src="${result.image_url}" class="max-h-48 mx-auto rounded" />
                     </div>`;
                 } else if (result.image_base64) {
-                    contentPreview = `<div class="bg-black/20 p-4 rounded-lg mb-4 text-center">
+                     contentPreview = `<div class="bg-black/20 p-4 rounded-lg mb-4 text-center">
                         <img src="${result.image_base64}" class="max-h-48 mx-auto rounded" />
                     </div>`;
                 }
@@ -1282,7 +1197,7 @@ with open("templates/index.html", "w", encoding='utf-8') as f:
                                 </div>
                             ` : ''}
                             <div class="grid grid-cols-2 md:grid-cols-4 gap-2 text-xs">
-                                ${Object.entries(result.category_scores).map(([category, score]) => `
+                                ${result.category_scores ? Object.entries(result.category_scores).map(([category, score]) => `
                                     <div class="bg-black/20 p-2 rounded">
                                         <div class="font-medium">${category.replace('/', ' ')}</div>
                                         <div class="w-full bg-gray-700 rounded-full h-1.5 mt-1">
@@ -1290,7 +1205,7 @@ with open("templates/index.html", "w", encoding='utf-8') as f:
                                         </div>
                                         <div class="text-right mt-1">${(score * 100).toFixed(0)}%</div>
                                     </div>
-                                `).join('')}
+                                `).join('') : ''}
                             </div>
                         </div>
                     </div>
@@ -1304,7 +1219,8 @@ with open("templates/index.html", "w", encoding='utf-8') as f:
         }
         
         async function updateMetrics() {
-            const apiKey = apiKeyInput.value.trim() || 'temp-key-for-metrics';
+            const apiKey = apiKeyInput.value.trim();
+            if (!apiKey) return;
             try {
                 const response = await fetch('/v1/metrics', {
                     headers: { 'Authorization': 'Bearer ' + apiKey }
@@ -1323,11 +1239,7 @@ with open("templates/index.html", "w", encoding='utf-8') as f:
         }
         
         function showLoading(show) {
-            if (show) {
-                loadingModal.classList.remove('hidden');
-            } else {
-                loadingModal.classList.add('hidden');
-            }
+            loadingModal.style.display = show ? 'flex' : 'none';
         }
         
         function showNotification(message, type = 'info') {
@@ -1335,12 +1247,7 @@ with open("templates/index.html", "w", encoding='utf-8') as f:
             notification.className = `fixed top-4 right-4 p-4 rounded-lg shadow-lg z-50 ${
                 type === 'error' ? 'bg-red-500' : 'bg-indigo-500'
             } text-white`;
-            notification.innerHTML = `
-                <div class="flex items-center">
-                    <i class="fas ${type === 'error' ? 'fa-exclamation-circle' : 'fa-info-circle'} mr-2"></i>
-                    <span>${message}</span>
-                </div>
-            `;
+            notification.innerHTML = `<div class="flex items-center"><i class="fas ${type === 'error' ? 'fa-exclamation-circle' : 'fa-info-circle'} mr-2"></i><span>${message}</span></div>`;
             
             document.body.appendChild(notification);