Update app.py

app.py

@@ -1,96 +1,20 @@
-import subprocess
-import time
-import requests
-import threading
 from flask import Flask, render_template, request, jsonify, Response
-import signal
-import sys
-import os
+from llama_cpp import Llama
+import json
 
 app = Flask(__name__)
 
-# ── Config ──────────────────────────────────────────────────────────
-LLAMA_SERVER_HOST = "127.0.0.1"
-LLAMA_SERVER_PORT = 8080
-MODEL_PATH = "/opt/models/smollm2-1.7b-instruct-q4_k_m.gguf"
-
-# Find the correct binary path (name changed across versions)
-POSSIBLE_BINS = [
-    "/opt/llama.cpp/build/bin/llama-server",
-    "/opt/llama.cpp/build/bin/server",
-    "/opt/llama.cpp/build/server",
-]
-
-LLAMA_BIN = None
-for path in POSSIBLE_BINS:
-    if os.path.isfile(path):
-        LLAMA_BIN = path
-        break
-
-if LLAMA_BIN is None:
-    print("ERROR: Could not find llama server binary!")
-    # Try to find it dynamically
-    result = subprocess.run(
-        ["find", "/opt/llama.cpp/build", "-name", "*server*", "-type", "f"],
-        capture_output=True, text=True
-    )
-    print(f"Available server-like files:\n{result.stdout}")
-    sys.exit(1)
-
-print(f"Using binary: {LLAMA_BIN}")
-
-llama_process = None
-
-
-def start_llama_server():
-    """Start llama.cpp server as a subprocess."""
-    global llama_process
-
-    cmd = [
-        LLAMA_BIN,
-        "-m", MODEL_PATH,
-        "--host", LLAMA_SERVER_HOST,
-        "--port", str(LLAMA_SERVER_PORT),
-        "-c", "2048",        # context size
-        "-t", "4",           # threads (HF free spaces usually have 2-4 cores)
-        "--mlock",           # lock model in RAM
-    ]
-
-    print(f"Starting llama.cpp server: {' '.join(cmd)}")
-    llama_process = subprocess.Popen(
-        cmd,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        text=True,
-    )
-
-    # Stream llama.cpp logs in background
-    def stream_logs():
-        for line in llama_process.stdout:
-            print(f"[llama.cpp] {line}", end="")
-
-    threading.Thread(target=stream_logs, daemon=True).start()
-
-
-def wait_for_server(timeout=120):
-    """Wait until the llama.cpp server is ready."""
-    url = f"http://{LLAMA_SERVER_HOST}:{LLAMA_SERVER_PORT}/health"
-    start = time.time()
-    while time.time() - start < timeout:
-        try:
-            r = requests.get(url, timeout=2)
-            if r.status_code == 200:
-                print("llama.cpp server is ready!")
-                return True
-        except requests.ConnectionError:
-            pass
-        time.sleep(2)
-    print("Timeout waiting for llama.cpp server")
-    return False
+# ── Load model once at startup (takes ~2-5 seconds) ──
+print("Loading model...")
+llm = Llama(
+    model_path="/opt/models/model.gguf",
+    n_ctx=2048,
+    n_threads=4,
+    verbose=False,
+)
+print("Model loaded!")
 
 
-# ── Routes ──────────────────────────────────────────────────────────
-
 @app.route("/")
 def index():
     return render_template("index.html")
@@ -98,138 +22,58 @@ def index():
 
 @app.route("/chat", methods=["POST"])
 def chat():
-    """Non-streaming chat endpoint."""
     data = request.json
-    user_message = data.get("message", "").strip()
-
-    if not user_message:
+    msg = data.get("message", "").strip()
+    if not msg:
         return jsonify({"error": "Empty message"}), 400
 
-    # Build prompt (SmolLM2-Instruct format)
-    prompt = (
-        f"<|im_start|>system\n"
-        f"You are a helpful assistant.<|im_end|>\n"
-        f"<|im_start|>user\n"
-        f"{user_message}<|im_end|>\n"
-        f"<|im_start|>assistant\n"
+    output = llm.create_chat_completion(
+        messages=[
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": msg},
+        ],
+        max_tokens=512,
+        temperature=0.7,
+        stream=False,
     )
 
-    payload = {
-        "prompt": prompt,
-        "n_predict": 512,
-        "temperature": 0.7,
-        "top_p": 0.9,
-        "stop": ["<|im_end|>", "<|im_start|>"],
-        "stream": False,
-    }
-
-    try:
-        resp = requests.post(
-            f"http://{LLAMA_SERVER_HOST}:{LLAMA_SERVER_PORT}/completion",
-            json=payload,
-            timeout=120,
-        )
-        result = resp.json()
-        return jsonify({"response": result.get("content", "").strip()})
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
+    reply = output["choices"][0]["message"]["content"]
+    return jsonify({"response": reply})
 
 
 @app.route("/chat/stream", methods=["POST"])
 def chat_stream():
-    """Streaming chat endpoint using SSE."""
     data = request.json
-    user_message = data.get("message", "").strip()
-
-    if not user_message:
+    msg = data.get("message", "").strip()
+    if not msg:
         return jsonify({"error": "Empty message"}), 400
 
-    prompt = (
-        f"<|im_start|>system\n"
-        f"You are a helpful assistant.<|im_end|>\n"
-        f"<|im_start|>user\n"
-        f"{user_message}<|im_end|>\n"
-        f"<|im_start|>assistant\n"
-    )
+    def generate():
+        stream = llm.create_chat_completion(
+            messages=[
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": msg},
+            ],
+            max_tokens=512,
+            temperature=0.7,
+            stream=True,
+        )
 
-    payload = {
-        "prompt": prompt,
-        "n_predict": 512,
-        "temperature": 0.7,
-        "top_p": 0.9,
-        "stop": ["<|im_end|>", "<|im_start|>"],
-        "stream": True,
-    }
+        for chunk in stream:
+            delta = chunk["choices"][0].get("delta", {})
+            content = delta.get("content", "")
+            if content:
+                yield f"data: {json.dumps({'content': content})}\n\n"
 
-    def generate():
-        try:
-            with requests.post(
-                f"http://{LLAMA_SERVER_HOST}:{LLAMA_SERVER_PORT}/completion",
-                json=payload,
-                stream=True,
-                timeout=120,
-            ) as resp:
-                import json as json_module
-                for line in resp.iter_lines():
-                    if line:
-                        decoded = line.decode("utf-8")
-                        if decoded.startswith("data: "):
-                            json_str = decoded[6:]
-                            if json_str.strip() == "[DONE]":
-                                yield f"data: [DONE]\n\n"
-                                break
-                            try:
-                                chunk = json_module.loads(json_str)
-                                content = chunk.get("content", "")
-                                if content:
-                                    yield f"data: {json_module.dumps({'content': content})}\n\n"
-                                if chunk.get("stop", False):
-                                    yield f"data: [DONE]\n\n"
-                                    break
-                            except json_module.JSONDecodeError:
-                                pass
-        except Exception as e:
-            import json as json_module
-            yield f"data: {json_module.dumps({'error': str(e)})}\n\n"
+        yield "data: [DONE]\n\n"
 
     return Response(generate(), mimetype="text/event-stream")
 
 
 @app.route("/health")
 def health():
-    try:
-        r = requests.get(
-            f"http://{LLAMA_SERVER_HOST}:{LLAMA_SERVER_PORT}/health",
-            timeout=5,
-        )
-        return jsonify({"flask": "ok", "llama": r.json()})
-    except Exception:
-        return jsonify({"flask": "ok", "llama": "not ready"}), 503
-
-
-# ── Cleanup ─────────────────────────────────────────────────────────
-
-def cleanup(signum=None, frame=None):
-    global llama_process
-    if llama_process:
-        print("Shutting down llama.cpp server...")
-        llama_process.terminate()
-        llama_process.wait()
-    sys.exit(0)
-
-
-signal.signal(signal.SIGTERM, cleanup)
-signal.signal(signal.SIGINT, cleanup)
-
+    return jsonify({"status": "ok"})
 
-# ── Main ────────────────────────────────────────────────────────────
 
 if __name__ == "__main__":
-    start_llama_server()
-
-    if wait_for_server(timeout=180):
-        print("Starting Flask on port 7860...")
-        app.run(host="0.0.0.0", port=7860, debug=False)
-    else:
-        print("Failed to start llama.cpp server. Exiting.")
-        cleanup()
+    app.run(host="0.0.0.0", port=7860)