Update model_suggestions.py

model_suggestions.py

@@ -1,110 +1,121 @@
 from nc_py_api import Nextcloud
 import json
-from typing import Dict, List
-import os
-import time
+from typing import Dict
 from datetime import datetime
-import threading
-import arena_config
+import os
 import re
+import config
 
 # Initialize Nextcloud client
 nc = Nextcloud(
-    nextcloud_url=arena_config.NEXTCLOUD_URL,
-    nc_auth_user=arena_config.NEXTCLOUD_USERNAME,
-    nc_auth_pass=arena_config.NEXTCLOUD_PASSWORD
+    nextcloud_url=config.NEXTCLOUD_URL,
+    nc_auth_user=config.NEXTCLOUD_USERNAME,
+    nc_auth_pass=config.NEXTCLOUD_PASSWORD
 )
 
-SUGGESTIONS_FILE = "model_suggestions.json"
-NEXTCLOUD_SUGGESTIONS_PATH = "/gpu_poor_model_suggestions.json"
-
 def validate_model_url(url: str) -> bool:
-    """Validate if the provided URL matches the expected HuggingFace format."""
-    pattern = r'^hf\.co/[\w-]+/[\w\.-]+(?:-GGUF)?:Q[0-9]+(?:_[A-Z0-9_]+)?$'
-    return bool(re.match(pattern, url))
+    """
+    Validate if the provided URL matches the expected format.
+    Accepts both direct Ollama models and HuggingFace GGUF models.
+    """
+    # Pattern for HuggingFace GGUF models
+    hf_pattern = r'^hf\.co/[\w-]+/[\w\.-]+(?:-GGUF)?:Q[0-9]+(?:_[A-Z0-9_]+)?$'
+    # Pattern for direct Ollama models
+    ollama_pattern = r'^[\w\.-]+(?::\d+(?:\.\d+)?[b])?(?:-[\w-]+)?:Q[0-9]+(?:_[A-Z0-9_]+)?$'
+    
+    return bool(re.match(hf_pattern, url) or re.match(ollama_pattern, url))
 
 def load_suggestions() -> Dict:
-    """Load suggestions from Nextcloud, fallback to local file, or initialize if neither exists."""
-    suggestions = None
-    
-    # First try to load from Nextcloud
+    """Load suggestions from Nextcloud with local file fallback."""
     try:
-        remote_data = nc.files.download(NEXTCLOUD_SUGGESTIONS_PATH)
+        # Try to load from Nextcloud
+        remote_data = nc.files.download(config.NEXTCLOUD_SUGGESTIONS_PATH)
         if remote_data:
             suggestions = json.loads(remote_data.decode('utf-8'))
-            # Update local file with Nextcloud data
-            with open(SUGGESTIONS_FILE, 'w') as f:
+            # Update local cache
+            with open('model_suggestions.json', 'w') as f:
                 json.dump(suggestions, f, indent=2)
             return suggestions
     except Exception as e:
         print(f"Could not load from Nextcloud: {e}")
+        
+        # Try local cache
+        if os.path.exists('model_suggestions.json'):
+            try:
+                with open('model_suggestions.json', 'r') as f:
+                    return json.load(f)
+            except Exception as e:
+                print(f"Could not load from local cache: {e}")
     
-    # If Nextcloud fails, try local file
-    if os.path.exists(SUGGESTIONS_FILE):
-        try:
-            with open(SUGGESTIONS_FILE, 'r') as f:
-                suggestions = json.load(f)
-                return suggestions
-        except Exception as e:
-            print(f"Could not load from local file: {e}")
-    
-    # If both fail, initialize new suggestions
-    suggestions = {
+    # Initialize new suggestions if both attempts fail
+    return {
         "suggestions": {},
-        "last_updated": datetime.now().isoformat()
+        "last_updated": datetime.now().isoformat(),
+        "total_suggestions": 0
     }
-    
-    # Save the new suggestions to both local and Nextcloud
-    save_suggestions(suggestions)
-    
-    return suggestions
 
-def save_suggestions(suggestions: Dict) -> None:
-    """Save suggestions to both local file and Nextcloud."""
-    with open(SUGGESTIONS_FILE, 'w') as f:
-        json.dump(suggestions, f, indent=2)
-    
-    # Upload to Nextcloud
+def save_suggestions(suggestions: Dict) -> bool:
+    """Save suggestions to both Nextcloud and local cache."""
     try:
-        nc.files.upload(
-            NEXTCLOUD_SUGGESTIONS_PATH,
-            json.dumps(suggestions, indent=2).encode('utf-8')
-        )
+        # Update metadata
+        suggestions["last_updated"] = datetime.now().isoformat()
+        suggestions["total_suggestions"] = sum(s["count"] for s in suggestions["suggestions"].values())
+        
+        # Save to Nextcloud
+        json_data = json.dumps(suggestions, indent=2)
+        nc.files.upload(config.NEXTCLOUD_SUGGESTIONS_PATH, json_data.encode('utf-8'))
+        
+        # Update local cache
+        with open('model_suggestions.json', 'w') as f:
+            json.dump(suggestions, f, indent=2)
+        
+        return True
     except Exception as e:
-        print(f"Error uploading to Nextcloud: {e}")
+        print(f"Error saving suggestions: {e}")
+        return False
 
 def add_suggestion(model_url: str) -> str:
-    """Add a new model suggestion or update existing one."""
+    """Add or update a model suggestion with validation."""
+    # Validate model URL format
     if not validate_model_url(model_url):
-        return "❌ Invalid model URL format. Please use the format: hf.co/username/model-name-GGUF:Q4_K_M"
+        return ("❌ Invalid model URL format. Please use either:\n"
+                "- Ollama format: model-name:Q4_K_M\n"
+                "- HuggingFace format: hf.co/username/model-name-GGUF:Q4_K_M")
+    
+    # Check if model is already approved
+    if model_url in dict(config.get_approved_models()):
+        return "ℹ️ This model is already in the arena!"
     
     suggestions = load_suggestions()
+    current_time = datetime.now().isoformat()
     
     if model_url in suggestions["suggestions"]:
-        suggestions["suggestions"][model_url]["count"] += 1
-        suggestions["suggestions"][model_url]["last_suggested"] = datetime.now().isoformat()
-        message = f"✨ Model suggestion updated! This model has been suggested {suggestions['suggestions'][model_url]['count']} times."
+        suggestions["suggestions"][model_url].update({
+            "count": suggestions["suggestions"][model_url]["count"] + 1,
+            "last_suggested": current_time
+        })
+        message = (f"✨ Model suggestion updated! "
+                  f"This model has been suggested {suggestions['suggestions'][model_url]['count']} times.")
     else:
         suggestions["suggestions"][model_url] = {
             "count": 1,
-            "first_suggested": datetime.now().isoformat(),
-            "last_suggested": datetime.now().isoformat()
+            "first_suggested": current_time,
+            "last_suggested": current_time
         }
         message = "✅ New model suggestion recorded successfully!"
     
-    suggestions["last_updated"] = datetime.now().isoformat()
-    save_suggestions(suggestions)
-    
-    return message
+    if save_suggestions(suggestions):
+        return message
+    return "❌ Error saving suggestion. Please try again later."
 
 def get_suggestions_html() -> str:
-    """Generate HTML table of model suggestions."""
+    """Generate HTML table of model suggestions with improved styling."""
     suggestions = load_suggestions()
     
-    # Sort suggestions by count (descending)
+    # Sort suggestions by count (descending) and last suggested date
     sorted_suggestions = sorted(
         suggestions["suggestions"].items(),
-        key=lambda x: x[1]["count"],
+        key=lambda x: (x[1]["count"], x[1]["last_suggested"]),
         reverse=True
     )
     
@@ -114,10 +125,11 @@ def get_suggestions_html() -> str:
             width: 100%;
             border-collapse: collapse;
             font-family: Arial, sans-serif;
+            margin-top: 20px;
         }
         .suggestions-table th, .suggestions-table td {
             border: 1px solid #ddd;
-            padding: 8px;
+            padding: 12px;
             text-align: left;
         }
         .suggestions-table th {
@@ -135,11 +147,15 @@ def get_suggestions_html() -> str:
             border-radius: 12px;
             font-size: 0.9em;
         }
-        .description-column {
+        .stats-header {
             font-size: 0.9em;
             color: #888;
+            margin-bottom: 10px;
         }
     </style>
+    <div class="stats-header">
+        Total Suggestions: {total_suggestions} | Last Updated: {last_updated}
+    </div>
     <table class='suggestions-table'>
     <tr>
         <th class='rank-column'>Rank</th>
@@ -148,7 +164,10 @@ def get_suggestions_html() -> str:
         <th>First Suggested</th>
         <th>Last Suggested</th>
     </tr>
-    """
+    """.format(
+        total_suggestions=suggestions.get("total_suggestions", 0),
+        last_updated=suggestions.get("last_updated", "Never").split("T")[0]
+    )
     
     for index, (model_url, data) in enumerate(sorted_suggestions, start=1):
         rank_display = {1: "🥇", 2: "🥈", 3: "🥉"}.get(index, f"{index}")