Replace simulated download with real ollama pull subprocess, add Ollama status check, installed models endpoint, WS wss support

Dockerfile

@@ -1,5 +1,9 @@
 FROM python:3.11-slim
 
+RUN apt-get update && apt-get install -y curl && \
+    curl -fsSL https://ollama.com/install.sh | sh && \
+    rm -rf /var/lib/apt/lists/*
+
 WORKDIR /app
 
 COPY requirements.txt .
@@ -9,4 +13,6 @@ COPY . .
 
 EXPOSE 7860
 
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ollama serve > /tmp/ollama.log 2>&1 & \
+    sleep 3 && \
+    uvicorn app:app --host 0.0.0.0 --port 7860

app.py

@@ -1,6 +1,8 @@
 import json
 import asyncio
 import uuid
+import re
+import shutil
 from datetime import datetime, timedelta
 from typing import Optional
 
@@ -151,6 +153,24 @@ async def me(user: dict = Depends(get_current_user)):
 async def get_models():
     return {"models": settings.recommended_models}
 
+@app.get("/api/installed-models")
+async def installed_models():
+    ollama_path = shutil.which("ollama")
+    if not ollama_path:
+        return {"models": []}
+    try:
+        proc = await asyncio.create_subprocess_exec(ollama_path, "list", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
+        stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=5)
+        models = []
+        for line in stdout.decode().strip().split("\n")[1:]:
+            if line.strip():
+                parts = line.split()
+                if len(parts) >= 1:
+                    models.append({"name": parts[0], "size": parts[2] if len(parts) > 2 else "?"})
+        return {"models": models}
+    except:
+        return {"models": []}
+
 
 # --- API Key Management ---
 
@@ -171,11 +191,32 @@ async def delete_api_key(key_id: int, user: dict = Depends(get_current_user)):
 
 # --- Download Sessions ---
 
+@app.get("/api/ollama-status")
+async def ollama_status():
+    ollama_path = shutil.which("ollama")
+    if not ollama_path:
+        return {"available": False, "message": "Ollama not found on this server"}
+    try:
+        proc = await asyncio.create_subprocess_exec(ollama_path, "list", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
+        stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=5)
+        models = []
+        for line in stdout.decode().strip().split("\n")[1:]:
+            if line.strip():
+                parts = line.split()
+                if len(parts) >= 1:
+                    models.append(parts[0])
+        return {"available": True, "message": "Ollama is running", "models": models}
+    except Exception as e:
+        return {"available": False, "message": f"Ollama server not reachable: {str(e)}"}
+
 @app.post("/api/download-sessions")
 async def start_download_session(req: CreateSessionRequest, user: dict = Depends(get_current_user)):
+    ollama_path = shutil.which("ollama")
+    if not ollama_path:
+        raise HTTPException(status_code=400, detail="Ollama is not installed on this server. Run locally with Ollama installed.")
     session_id = str(uuid.uuid4())
     create_download_session(user["id"], session_id, req.model_name)
-    asyncio.create_task(simulate_download(session_id, req.model_name))
+    asyncio.create_task(ollama_download(session_id, req.model_name, ollama_path))
     return {"session_id": session_id, "model_name": req.model_name, "status": "started"}
 
 @app.get("/api/download-sessions")
@@ -204,28 +245,112 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str):
         manager.disconnect(session_id, websocket)
 
 
-# --- Simulated Download (mocked until Ollama works) ---
-
-async def simulate_download(session_id: str, model_name: str):
-    total_mb = 4500
-    downloaded = 0
-    speed = 5000
+# --- Real Ollama Download ---
+
+def parse_ollama_progress(line: str) -> Optional[dict]:
+    m = re.search(r'pulling\s+\S+\s*:\s*(\d+)%\s*.*?(\d+\.?\d*)\s*(KB|MB|GB)\s*/\s*(\d+\.?\d*)\s*(KB|MB|GB)\s*(\d+\.?\d*)\s*(KB|MB)/s\s*(.+)?', line)
+    if m:
+        pct = int(m.group(1))
+        dl_val = float(m.group(2))
+        dl_unit = m.group(3)
+        total_val = float(m.group(4))
+        total_unit = m.group(5)
+        speed_val = float(m.group(6))
+        speed_unit = m.group(7)
+        eta_str = m.group(8).strip() if m.group(8) else ""
+
+        def to_mb(val, unit):
+            if unit == "GB": return int(val * 1024)
+            if unit == "KB": return int(val / 1024)
+            return int(val)
+
+        downloaded_mb = to_mb(dl_val, dl_unit)
+        total_mb = to_mb(total_val, total_unit)
+        speed_kbps = int(speed_val * 1024 if speed_unit == "MB" else speed_val)
+
+        eta_seconds = 0
+        if eta_str:
+            eta_parts = eta_str.split(":")
+            if len(eta_parts) == 2:
+                eta_seconds = int(eta_parts[0]) * 60 + int(eta_parts[1])
+            elif "m" in eta_str and "s" in eta_str:
+                parts = eta_str.replace("m", " ").replace("s", "").split()
+                eta_seconds = int(parts[0]) * 60 + (int(parts[1]) if len(parts) > 1 else 0)
+            elif "h" in eta_str:
+                parts = eta_str.replace("h", " ").replace("m", "").split()
+                eta_seconds = int(parts[0]) * 3600 + (int(parts[1]) * 60 if len(parts) > 1 else 0)
+
+        return {
+            "progress": pct,
+            "downloaded_mb": downloaded_mb,
+            "total_size_mb": total_mb,
+            "speed_kbps": speed_kbps,
+            "eta_seconds": eta_seconds,
+            "current_step": f"Downloading layer ({pct}%)"
+        }
+    if "pulling manifest" in line:
+        return {"current_step": "Pulling manifest...", "progress": 0}
+    if "verifying sha256" in line:
+        return {"current_step": "Verifying...", "progress": 95}
+    if "writing manifest" in line:
+        return {"current_step": "Writing manifest...", "progress": 98}
+    if "success" in line:
+        return {"current_step": "Done", "progress": 100, "status": "completed"}
+    return None
+
+async def ollama_download(session_id: str, model_name: str, ollama_path: str):
+    update_download_progress(session_id, status="starting", progress=0, current_step="Starting download...")
+    await manager.broadcast(session_id, {"status": "starting", "progress": 0, "current_step": "Starting download..."})
 
-    update_download_progress(session_id, status="downloading", progress=0, total_size_mb=total_mb, downloaded_mb=0, speed_kbps=speed, current_step="Connecting...", eta_seconds=total_mb * 8 // speed)
-    await manager.broadcast(session_id, {"status": "downloading", "progress": 0, "current_step": "Connecting...", "speed_kbps": speed, "eta_seconds": total_mb * 8 // speed})
-
-    await asyncio.sleep(2)
-
-    for i in range(1, 101):
-        await asyncio.sleep(0.5)
-        downloaded = int(total_mb * i / 100)
-        eta = int((total_mb - downloaded) * 8 / speed)
-        step = "Downloading manifest" if i < 10 else ("Downloading layers" if i < 90 else "Verifying")
-        update_download_progress(session_id, progress=i, downloaded_mb=downloaded, current_step=step, speed_kbps=speed, eta_seconds=eta)
-        await manager.broadcast(session_id, {"status": "downloading", "progress": i, "downloaded_mb": downloaded, "total_size_mb": total_mb, "current_step": step, "speed_kbps": speed, "eta_seconds": eta})
-
-    update_download_progress(session_id, status="completed", progress=100, downloaded_mb=total_mb, current_step="Done", speed_kbps=0, eta_seconds=0, completed_at=datetime.utcnow().isoformat())
-    await manager.broadcast(session_id, {"status": "completed", "progress": 100, "current_step": "Done"})
+    try:
+        proc = await asyncio.create_subprocess_exec(
+            ollama_path, "pull", model_name,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE
+        )
+
+        last_progress = {}
+
+        async def read_stderr():
+            nonlocal last_progress
+            while True:
+                line = await proc.stderr.readline()
+                if not line:
+                    break
+                decoded = line.decode(errors="replace").strip()
+                if not decoded:
+                    continue
+                parsed = parse_ollama_progress(decoded)
+                if parsed:
+                    last_progress.update({k: v for k, v in parsed.items() if v is not None})
+                    data = dict(last_progress)
+                    data["status"] = "downloading"
+                    update_download_progress(session_id, **data)
+                    await manager.broadcast(session_id, data)
+
+        async def read_stdout():
+            async for line in proc.stdout:
+                decoded = line.decode(errors="replace").strip()
+                if decoded:
+                    last_progress["current_step"] = decoded
+                    data = dict(last_progress)
+                    data["status"] = "downloading"
+                    update_download_progress(session_id, **data)
+                    await manager.broadcast(session_id, data)
+
+        await asyncio.gather(read_stderr(), read_stdout())
+        await proc.wait()
+
+        if proc.returncode == 0:
+            update_download_progress(session_id, status="completed", progress=100, current_step="Done", speed_kbps=0, eta_seconds=0, completed_at=datetime.utcnow().isoformat())
+            await manager.broadcast(session_id, {"status": "completed", "progress": 100, "current_step": "Done"})
+        else:
+            update_download_progress(session_id, status="error", current_step="Download failed", error_message=f"Process exited with code {proc.returncode}")
+            await manager.broadcast(session_id, {"status": "error", "current_step": "Download failed"})
+
+    except Exception as e:
+        update_download_progress(session_id, status="error", current_step="Error", error_message=str(e))
+        await manager.broadcast(session_id, {"status": "error", "current_step": f"Error: {str(e)}"})
 
 
 # --- Startup ---

templates/dashboard.html

@@ -10,9 +10,18 @@
     <div class="container">
         <header>
             <h1>Dashboard</h1>
-            <button class="btn btn-secondary" onclick="logout()">Logout</button>
+            <div>
+                <span id="ollama-badge" class="status status-pending">Checking Ollama...</span>
+                <button class="btn btn-secondary" onclick="logout()">Logout</button>
+            </div>
         </header>
 
+        <div id="ollama-warning" class="card hidden" style="border: 2px solid #e63946;">
+            <h2>Ollama Not Available</h2>
+            <p id="ollama-msg">Ollama is not installed or not running on this server.</p>
+            <p>To use this app locally: <code>ollama serve</code> then <code>python app.py</code></p>
+        </div>
+
         <div class="dashboard-grid">
             <div class="card">
                 <h2>Your Profile</h2>
@@ -29,11 +38,17 @@
             </div>
 
             <div class="card">
+                <h2>Installed Models</h2>
+                <ul id="installed-models" class="model-list"><li>Loading...</li></ul>
+            </div>
+
+            <div class="card full-width">
                 <h2>Start Download</h2>
                 <form onsubmit="startDownload(event)">
                     <select id="model-select" required></select>
                     <button type="submit" class="btn btn-primary">Download</button>
                 </form>
+                <p id="download-error" class="error"></p>
             </div>
 
             <div class="card full-width">
@@ -72,7 +87,31 @@
             `;
         });
 
-        // Load models
+        // Check Ollama status
+        fetch('/api/ollama-status').then(r => r.json()).then(data => {
+            const badge = document.getElementById('ollama-badge');
+            if (data.available) {
+                badge.className = 'status status-completed';
+                badge.textContent = 'Ollama Ready';
+            } else {
+                badge.className = 'status status-error';
+                badge.textContent = 'Ollama Unavailable';
+                document.getElementById('ollama-warning').classList.remove('hidden');
+                document.getElementById('ollama-msg').textContent = data.message;
+            }
+        });
+
+        // Load installed models
+        fetch('/api/installed-models').then(r => r.json()).then(data => {
+            const ul = document.getElementById('installed-models');
+            if (data.models.length === 0) {
+                ul.innerHTML = '<li>No models installed yet. Download one below.</li>';
+            } else {
+                ul.innerHTML = data.models.map(m => `<li><strong>${m.name}</strong> (${m.size})</li>`).join('');
+            }
+        });
+
+        // Load recommended models
         fetch('/api/models').then(r => r.json()).then(data => {
             const select = document.getElementById('model-select');
             data.models.filter(m => !m.name.includes('70b')).forEach(m => {
@@ -118,6 +157,10 @@
             const model = document.getElementById('model-select').value;
             const res = await api('/api/download-sessions', {method: 'POST', body: {model_name: model}});
             const data = await res.json();
+            if (!res.ok) {
+                document.getElementById('download-error').textContent = data.detail || 'Download failed';
+                return;
+            }
             window.location.href = `/progress/${data.session_id}`;
         }
 

templates/progress.html

@@ -44,7 +44,8 @@
         });
 
         // WebSocket
-        const ws = new WebSocket(`ws://${window.location.host}/ws/${sessionId}`);
+        const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
+        const ws = new WebSocket(`${protocol}//${window.location.host}/ws/${sessionId}`);
 
         ws.onmessage = function(event) {
             const data = JSON.parse(event.data);
@@ -55,9 +56,9 @@
 
             if (data.status) document.getElementById('status-text').textContent = data.status;
             if (data.current_step) document.getElementById('step-text').textContent = data.current_step;
-            if (data.downloaded_mb) document.getElementById('downloaded-text').textContent = data.downloaded_mb + ' MB';
-            if (data.total_size_mb) document.getElementById('total-text').textContent = data.total_size_mb + ' MB';
-            if (data.speed_kbps) document.getElementById('speed-text').textContent = data.speed_kbps;
+            if (data.downloaded_mb !== undefined) document.getElementById('downloaded-text').textContent = data.downloaded_mb + ' MB';
+            if (data.total_size_mb !== undefined) document.getElementById('total-text').textContent = data.total_size_mb + ' MB';
+            if (data.speed_kbps !== undefined) document.getElementById('speed-text').textContent = data.speed_kbps;
             if (data.eta_seconds !== undefined) {
                 const eta = data.eta_seconds;
                 if (eta <= 0) document.getElementById('eta-text').textContent = 'Done';
@@ -68,11 +69,25 @@
             if (data.status === 'completed') {
                 document.getElementById('status-text').textContent = 'Completed';
                 document.getElementById('step-text').textContent = 'Done';
+                document.getElementById('progress-bar').style.width = '100%';
+                document.getElementById('progress-percent').textContent = '100%';
+            }
+
+            if (data.status === 'error') {
+                document.getElementById('status-text').textContent = 'Error';
+                document.getElementById('progress-bar').style.background = 'linear-gradient(90deg, #e63946, #c1121f)';
             }
         };
 
         ws.onerror = function() {
-            document.getElementById('status-text').textContent = 'Connection lost';
+            document.getElementById('status-text').textContent = 'Connecting...';
+        };
+
+        ws.onclose = function() {
+            if (document.getElementById('status-text').textContent !== 'Completed' &&
+                document.getElementById('status-text').textContent !== 'Error') {
+                document.getElementById('status-text').textContent = 'Disconnected';
+            }
         };
     </script>
 </body>