Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py
 from flask import Flask, request, jsonify, render_template_string
 import os
 import requests
@@ -10,32 +11,33 @@ app = Flask(__name__)
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Configuration
-OLLAMA_API_URL = os.getenv('OLLAMA_API_URL', 'https://huggingface.co/spaces/tommytracx/ollama-api')
-DEFAULT_MODEL = os.getenv('DEFAULT_MODEL', 'llama2,llama2:13b,llama2:70b,codellama,neural-chat,gemma-3-270m').split(',')
+OLLAMA_BASE_URL = os.getenv('OLLAMA_BASE_URL', 'http://localhost:11434')
+ALLOWED_MODELS = os.getenv('ALLOWED_MODELS', 'llama2,llama2:13b,llama2:70b,codellama,neural-chat,gemma-3-270m').split(',')
 MAX_TOKENS = int(os.getenv('MAX_TOKENS', '2048'))
 TEMPERATURE = float(os.getenv('TEMPERATURE', '0.7'))
 
-class OllamaClient:
-    def __init__(self, api_url: str):
-        self.api_url = api_url.rstrip('/')
-        self.available_models = DEFAULT_MODEL  # Initialize with default models
+class OllamaManager:
+    def __init__(self, base_url: str):
+        self.base_url = base_url.rstrip('/')
+        self.available_models = ALLOWED_MODELS  # Initialize with allowed models
         self.refresh_models()
     
     def refresh_models(self) -> None:
-        """Refresh the list of available models from the API, falling back to defaults on failure."""
+        """Refresh the list of available models from Ollama API, falling back to allowed models."""
         try:
-            response = requests.get(f"{self.api_url}/api/models", timeout=10)
+            response = requests.get(f"{self.base_url}/api/tags", timeout=10)
             response.raise_for_status()
             data = response.json()
-            if data.get('status') == 'success' and isinstance(data.get('models'), list):
-                self.available_models = data['models']
-                logging.info(f"Successfully fetched models: {self.available_models}")
-            else:
-                logging.warning(f"Invalid response format from API: {data}")
-                self.available_models = DEFAULT_MODEL
+            models = [model['name'] for model in data.get('models', [])]
+            # Filter models to only include those in ALLOWED_MODELS
+            self.available_models = [model for model in models if model in ALLOWED_MODELS]
+            if not self.available_models:
+                self.available_models = ALLOWED_MODELS
+                logging.warning("No allowed models found in API response, using ALLOWED_MODELS")
+            logging.info(f"Available models: {self.available_models}")
         except Exception as e:
             logging.error(f"Error refreshing models: {e}")
-            self.available_models = DEFAULT_MODEL
+            self.available_models = ALLOWED_MODELS
     
     def list_models(self) -> List[str]:
         """Return the list of available models."""
@@ -53,18 +55,15 @@ class OllamaClient:
                 "stream": False,
                 **kwargs
             }
-            response = requests.post(f"{self.api_url}/api/generate", json=payload, timeout=120)
+            response = requests.post(f"{self.base_url}/api/generate", json=payload, timeout=120)
             response.raise_for_status()
             data = response.json()
-            
-            if data.get('status') == 'success':
-                return {
-                    "status": "success",
-                    "response": data.get('response', ''),
-                    "model": model_name,
-                    "usage": data.get('usage', {})
-                }
-            return {"status": "error", "message": data.get('message', 'Unknown error')}
+            return {
+                "status": "success",
+                "response": data.get('response', ''),
+                "model": model_name,
+                "usage": data.get('usage', {})
+            }
         except Exception as e:
             logging.error(f"Error generating response: {e}")
             return {"status": "error", "message": str(e)}
@@ -72,15 +71,15 @@ class OllamaClient:
     def health_check(self) -> Dict[str, Any]:
         """Check the health of the Ollama API."""
         try:
-            response = requests.get(f"{self.api_url}/health", timeout=10)
+            response = requests.get(f"{self.base_url}/api/tags", timeout=10)
             response.raise_for_status()
-            return response.json()
+            return {"status": "healthy", "available_models": len(self.available_models)}
         except Exception as e:
             logging.error(f"Health check failed: {e}")
             return {"status": "unhealthy", "error": str(e)}
 
-# Initialize Ollama client
-ollama_client = OllamaClient(OLLAMA_API_URL)
+# Initialize Ollama manager
+ollama_manager = OllamaManager(OLLAMA_BASE_URL)
 
 # HTML template for the chat interface
 HTML_TEMPLATE = '''
@@ -91,7 +90,195 @@ HTML_TEMPLATE = '''
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
     <title>OpenWebUI - Ollama Chat</title>
     <style>
-        /* [Previous CSS unchanged] */
+        * {
+            margin: 0;
+            padding: 0;
+            box-sizing: border-box;
+        }
+        body {
+            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            min-height: 100vh;
+            padding: 20px;
+        }
+        .container {
+            max-width: 1200px;
+            margin: 0 auto;
+            background: white;
+            border-radius: 20px;
+            box-shadow: 0 20px 40px rgba(0,0,0,0.1);
+            overflow: hidden;
+        }
+        .header {
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            color: white;
+            padding: 30px;
+            text-align: center;
+        }
+        .header h1 {
+            font-size: 2.5rem;
+            margin-bottom: 10px;
+            font-weight: 700;
+        }
+        .header p {
+            font-size: 1.1rem;
+            opacity: 0.9;
+        }
+        .controls {
+            padding: 20px 30px;
+            background: #f8f9fa;
+            border-bottom: 1px solid #e9ecef;
+            display: flex;
+            gap: 15px;
+            align-items: center;
+            flex-wrap: wrap;
+        }
+        .control-group {
+            display: flex;
+            align-items: center;
+            gap: 8px;
+        }
+        .control-group label {
+            font-weight: 600;
+            color: #495057;
+            min-width: 80px;
+        }
+        .control-group select,
+        .control-group input {
+            padding: 8px 12px;
+            border: 2px solid #e9ecef;
+            border-radius: 8px;
+            font-size: 14px;
+            transition: border-color 0.3s;
+        }
+        .control-group select:focus,
+        .control-group input:focus {
+            outline: none;
+            border-color: #667eea;
+        }
+        .chat-container {
+            height: 500px;
+            overflow-y: auto;
+            padding: 20px;
+            background: #fafbfc;
+        }
+        .message {
+            margin-bottom: 20px;
+            display: flex;
+            gap: 15px;
+        }
+        .message.user {
+            flex-direction: row-reverse;
+        }
+        .message-avatar {
+            width: 40px;
+            height: 40px;
+            border-radius: 50%;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            font-weight: bold;
+            color: white;
+            flex-shrink: 0;
+        }
+        .message.user .message-avatar {
+            background: #667eea;
+        }
+        .message.assistant .message-avatar {
+            background: #28a745;
+        }
+        .message-content {
+            background: white;
+            padding: 15px 20px;
+            border-radius: 18px;
+            max-width: 70%;
+            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
+            line-height: 1.5;
+        }
+        .message.user .message-content {
+            background: #667eea;
+            color: white;
+        }
+        .message.assistant .message-content {
+            background: white;
+            color: #333;
+        }
+        .input-container {
+            padding: 20px 30px;
+            background: white;
+            border-top: 1px solid #e9ecef;
+        }
+        .input-form {
+            display: flex;
+            gap: 15px;
+        }
+        .input-field {
+            flex: 1;
+            padding: 15px 20px;
+            border: 2px solid #e9ecef;
+            border-radius: 25px;
+            font-size: 16px;
+            transition: border-color 0.3s;
+            resize: none;
+            min-height: 50px;
+            max-height: 120px;
+        }
+        .input-field:focus {
+            outline: none;
+            border-color: #667eea;
+        }
+        .send-button {
+            padding: 15px 30px;
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            color: white;
+            border: none;
+            border-radius: 25px;
+            font-size: 16px;
+            font-weight: 600;
+            cursor: pointer;
+            transition: transform 0.2s;
+            min-width: 100px;
+        }
+        .send-button:hover {
+            transform: translateY(-2px);
+        }
+        .send-button:disabled {
+            opacity: 0.6;
+            cursor: not-allowed;
+            transform: none;
+        }
+        .status {
+            text-align: center;
+            padding: 10px;
+            font-size: 14px;
+            color: #6c757d;
+        }
+        .status.error {
+            color: #dc3545;
+        }
+        .status.success {
+            color: #28a745;
+        }
+        .typing-indicator {
+            display: none;
+            padding: 15px 20px;
+            background: white;
+            border-radius: 18px;
+            color: #6c757d;
+            font-style: italic;
+        }
+        @media (max-width: 768px) {
+            .controls {
+                flex-direction: column;
+                align-items: stretch;
+            }
+            .control-group {
+                justify-content: space-between;
+            }
+            .message-content {
+                max-width: 85%;
+            }
+        }
     </style>
 </head>
 <body>
@@ -173,7 +360,7 @@ HTML_TEMPLATE = '''
                         const option = document.createElement('option');
                         option.value = model;
                         option.textContent = model;
-                        if (model === '{{ default_model[0] }}') {
+                        if (model === 'gemma-3-270m') {
                             option.selected = true;
                         }
                         modelSelect.appendChild(option);
@@ -230,7 +417,7 @@ HTML_TEMPLATE = '''
                 const response = await fetch('/api/chat', {
                     method: 'POST',
                     headers: { 'Content-Type': 'application/json' },
-                    body: JSON.stringify({ model, message, temperature, max_tokens: maxTokens })
+                    body: JSON.stringify({ model, prompt: message, temperature, max_tokens: maxTokens })
                 });
                 const data = await response.json();
                 
@@ -297,22 +484,22 @@ HTML_TEMPLATE = '''
 @app.route('/')
 def home():
     """Main chat interface."""
-    return render_template_string(HTML_TEMPLATE, ollama_api_url=OLLAMA_API_URL, default_model=DEFAULT_MODEL)
+    return render_template_string(HTML_TEMPLATE, ollama_base_url=OLLAMA_BASE_URL, default_model=ALLOWED_MODELS)
 
 @app.route('/api/chat', methods=['POST'])
 def chat():
     """Chat API endpoint."""
     try:
         data = request.get_json()
-        if not data or 'message' not in data or 'model' not in data:
-            return jsonify({"status": "error", "message": "Message and model are required"}), 400
+        if not data or 'prompt' not in data or 'model' not in data:
+            return jsonify({"status": "error", "message": "Prompt and model are required"}), 400
         
-        message = data['message']
+        prompt = data['prompt']
         model = data['model']
         temperature = data.get('temperature', TEMPERATURE)
         max_tokens = data.get('max_tokens', MAX_TOKENS)
         
-        result = ollama_client.generate(model, message, temperature=temperature, max_tokens=max_tokens)
+        result = ollama_manager.generate(model, prompt, temperature=temperature, max_tokens=max_tokens)
         return jsonify(result), 200 if result["status"] == "success" else 500
     except Exception as e:
         logging.error(f"Chat endpoint error: {e}")
@@ -322,7 +509,7 @@ def chat():
 def get_models():
     """Get available models."""
     try:
-        models = ollama_client.list_models()
+        models = ollama_manager.list_models()
         return jsonify({
             "status": "success",
             "models": models,
@@ -336,7 +523,7 @@ def get_models():
 def health_check():
     """Health check endpoint."""
     try:
-        ollama_health = ollama_client.health_check()
+        ollama_health = ollama_manager.health_check()
         return jsonify({
             "status": "healthy",
             "ollama_api": ollama_health,