Update server.py

server.py

@@ -9,6 +9,7 @@ import os
 import json
 import time
 import hashlib
+import threading
 
 app = FastAPI()
 
@@ -26,12 +27,13 @@ LLAMA_SERVER = "./llama.cpp/build/bin/llama-server"
 
 
 # ---------------------------
-# REQUEST SCHEMAS
+# REQUEST MODELS
 # ---------------------------
 
 class ChatRequest(BaseModel):
     model: str
     messages: list
+    stream: bool = True
 
 
 class GenerateRequest(BaseModel):
@@ -75,7 +77,7 @@ os.makedirs("models", exist_ok=True)
 
 if not os.path.exists(MODEL_PATH):
 
-    print("Downloading model from HuggingFace")
+    print("Downloading model from HuggingFace...")
 
     downloaded = hf_hub_download(
         repo_id=MODEL_REPO,
@@ -91,15 +93,29 @@ if not os.path.exists(MODEL_PATH):
 # START LLAMA SERVER
 # ---------------------------
 
-print("Starting llama-server...")
+def start_llama():
 
-subprocess.Popen([
-    LLAMA_SERVER,
-    "-m", MODEL_PATH,
-    "--host", "0.0.0.0",
-    "--port", "8080",
-    "-c", "2048"
-])
+    print("Starting llama-server...")
+
+    subprocess.Popen([
+        LLAMA_SERVER,
+        "-m", MODEL_PATH,
+        "--host", "0.0.0.0",
+        "--port", "8080",
+        "-c", "2048"
+    ])
+
+    # wait for server to start
+    for _ in range(30):
+        try:
+            requests.get("http://localhost:8080/health")
+            print("llama-server ready")
+            return
+        except:
+            time.sleep(1)
+
+
+threading.Thread(target=start_llama, daemon=True).start()
 
 
 # ---------------------------
@@ -112,120 +128,65 @@ def root():
 
 
 # ---------------------------
-# OLLAMA MODEL LIST
+# MODEL LIST (Ollama style)
 # ---------------------------
 
-@app.post("/api/chat")
-def chat(req: ChatRequest):
-
-    prompt = build_prompt(req.messages)
-
-    stream = getattr(req, "stream", False)
-
-    r = requests.post(
-        "http://localhost:8080/completion",
-        json={
-            "prompt": prompt,
-            "stream": stream,
-            "n_predict": 512
-        },
-        stream=stream
-    )
+@app.get("/api/tags")
+def tags():
 
-    # NON STREAM MODE (normal JSON)
-    if not stream:
+    size = os.path.getsize(MODEL_PATH)
 
-        data = r.json()
-        text = data.get("content", "")
+    with open(MODEL_PATH, "rb") as f:
+        digest = hashlib.sha256(f.read()).hexdigest()
 
-        return JSONResponse({
-            "model": req.model,
-            "message": {
-                "role": "assistant",
-                "content": text
-            },
-            "done": True
-        })
-
-    # STREAM MODE (NDJSON like Ollama)
-    def stream_generator():
-
-        for line in r.iter_lines():
-
-            if not line:
-                continue
-
-            line = line.decode("utf-8").strip()
-
-            if line.startswith("data:"):
-                line = line[5:].strip()
-
-            try:
-                data = json.loads(line)
-            except:
-                continue
-
-            token = data.get("content", "")
-
-            yield json.dumps({
-                "model": req.model,
-                "message": {
-                    "role": "assistant",
-                    "content": token
-                },
-                "done": False
-            }) + "\n"
-
-        yield json.dumps({
-            "model": req.model,
-            "done": True
-        }) + "\n"
+    return {
+        "models": [
+            {
+                "name": MODEL_NAME,
+                "model": MODEL_NAME,
+                "modified_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
+                "size": size,
+                "digest": digest,
+                "details": {
+                    "format": "gguf",
+                    "family": "llama",
+                    "families": ["llama"],
+                    "parameter_size": "1.1B",
+                    "quantization_level": "Q4_K_M"
+                }
+            }
+        ]
+    }
 
-    return StreamingResponse(
-        stream_generator(),
-        media_type="application/x-ndjson"
-    )
 
 # ---------------------------
-# GENERATE (NON STREAM)
+# GENERATE (non-stream)
 # ---------------------------
 
 @app.post("/api/generate")
 def generate(req: GenerateRequest):
 
-    start = time.time()
-
-    response = requests.post(
+    r = requests.post(
         "http://localhost:8080/completion",
         json={
             "prompt": req.prompt,
-            "n_predict": 200
+            "n_predict": 256
         }
     )
 
-    data = response.json()
+    data = r.json()
 
     text = data.get("content", "").strip()
 
-    duration = int((time.time() - start) * 1e9)
-
     return {
         "model": req.model,
-        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
         "response": text,
-        "done": True,
-        "done_reason": "stop",
-        "total_duration": duration,
-        "load_duration": 0,
-        "prompt_eval_count": len(req.prompt.split()),
-        "prompt_eval_duration": 0,
-        "eval_count": len(text.split()),
-        "eval_duration": duration
+        "done": True
     }
 
 
 # ---------------------------
-# CHAT STREAM (OLLAMA STYLE)
+# CHAT (Ollama streaming)
 # ---------------------------
 
 @app.post("/api/chat")
@@ -237,16 +198,30 @@ def chat(req: ChatRequest):
         "http://localhost:8080/completion",
         json={
             "prompt": prompt,
-            "stream": True,
-            "n_predict": 1024,
+            "stream": req.stream,
+            "n_predict": 256,
             "temperature": 0.7,
             "top_p": 0.9,
             "stop": ["User:", "</s>"]
         },
-        stream=True
+        stream=req.stream
     )
 
-    def stream():
+    if not req.stream:
+
+        data = r.json()
+        text = data.get("content", "")
+
+        return JSONResponse({
+            "model": req.model,
+            "message": {
+                "role": "assistant",
+                "content": text
+            },
+            "done": True
+        })
+
+    def stream_generator():
 
         for line in r.iter_lines():
 
@@ -255,30 +230,24 @@ def chat(req: ChatRequest):
 
             line = line.decode("utf-8").strip()
 
-            if not line:
-                continue
-            
             if line.startswith("data:"):
                 line = line[5:].strip()
-            
+
             try:
                 data = json.loads(line)
-            except Exception:
+            except:
                 continue
 
             token = data.get("content", "")
 
-            chunk = {
+            yield json.dumps({
                 "model": req.model,
-                "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
                 "message": {
                     "role": "assistant",
                     "content": token
                 },
                 "done": False
-            }
-
-            yield json.dumps(chunk) + "\n"
+            }) + "\n"
 
         yield json.dumps({
             "model": req.model,
@@ -286,7 +255,11 @@ def chat(req: ChatRequest):
             "done_reason": "stop"
         }) + "\n"
 
-    return StreamingResponse(stream(), media_type="application/x-ndjson")
+    return StreamingResponse(
+        stream_generator(),
+        media_type="application/x-ndjson",
+        headers={"Cache-Control": "no-cache"}
+    )
 
 
 # ---------------------------