Update server.py

server.py

@@ -1,6 +1,7 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from huggingface_hub import hf_hub_download
+import requests
 import subprocess
 import uvicorn
 import os
@@ -9,8 +10,7 @@ import json
 app = FastAPI()
 
 MODELS = {
-    "tinyllama": "models/tinyllama.gguf",
-    "qwen": "models/qwen1.5b.gguf"
+    "tinyllama": "models/tinyllama.gguf"
 }
 
 class ChatRequest(BaseModel):
@@ -23,7 +23,7 @@ class GenerateRequest(BaseModel):
 
 
 # ---------------------------
-# Utility Logging
+# logging
 # ---------------------------
 
 def log(title, data):
@@ -34,177 +34,153 @@ def log(title, data):
 
 
 # ---------------------------
-# Prompt builder
+# prompt builder
 # ---------------------------
 
 def build_prompt(messages):
 
-    log("CHAT HISTORY", json.dumps(messages, indent=2))
-
     prompt = ""
 
     for m in messages:
         role = m.get("role", "user")
         content = m.get("content", "")
 
+        if content.strip() == "":
+            continue
+
         prompt += f"{role}: {content}\n"
 
     prompt += "assistant:"
 
-    log("FINAL PROMPT", prompt)
+    log("PROMPT", prompt)
 
     return prompt
 
 
 # ---------------------------
-# Run llama.cpp
+# download model
 # ---------------------------
 
-def run_model(model_path, prompt):
+os.makedirs("models", exist_ok=True)
+
+MODEL_FILES = {
+    "tinyllama": (
+        "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF",
+        "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
+    )
+}
 
-    log("MODEL PATH", model_path)
+for name, (repo, file) in MODEL_FILES.items():
 
-    if not os.path.exists(model_path):
-        log("ERROR", f"Model file missing: {model_path}")
-        return "Model file not found"
+    path = f"models/{name}.gguf"
 
-    command = [
-        "./llama.cpp/build/bin/llama-cli",
-        "-m", model_path,
-        "-p", prompt,
-        "-n", "200",
-        "--no-display-prompt"
-    ]
+    if not os.path.exists(path):
 
-    log("EXEC COMMAND", command)
+        print(f"Downloading model {name}")
 
-    result = subprocess.run(
-        command,
-        capture_output=True,
-        text=True
-    )
+        downloaded = hf_hub_download(
+            repo_id=repo,
+            filename=file
+        )
 
-    log("LLAMA STDOUT", result.stdout)
-    log("LLAMA STDERR", result.stderr)
+        os.system(f"cp {downloaded} {path}")
+
+        print(f"Model ready: {path}")
 
-    output = result.stdout.strip()
 
-    if "assistant:" in output:
-        output = output.split("assistant:")[-1].strip()
+# ---------------------------
+# start llama-server
+# ---------------------------
 
-    log("FINAL OUTPUT", output)
+print("Starting llama-server...")
 
-    return output
+subprocess.Popen([
+    "./llama.cpp/build/bin/llama-server",
+    "-m", "models/tinyllama.gguf",
+    "--host", "0.0.0.0",
+    "--port", "8080",
+    "-c", "2048"
+])
 
 
 # ---------------------------
-# Root endpoint
+# root
 # ---------------------------
 
 @app.get("/")
 def root():
-    log("SERVER STATUS", "Server running")
     return {"status": "running"}
 
 
 # ---------------------------
-# Model list (Ollama compatible)
+# model list
 # ---------------------------
 
 @app.get("/api/tags")
 def list_models():
 
-    models = []
-
-    for name in MODELS.keys():
-        models.append({
-            "name": name,
-            "model": name
-        })
-
-    log("MODEL LIST REQUEST", models)
-
-    return {"models": models}
-
-
-# ---------------------------
-# Generate endpoint
-# ---------------------------
-
-@app.post("/api/generate")
-def generate(req: GenerateRequest):
-
-    log("GENERATE REQUEST", req.dict())
-
-    if req.model not in MODELS:
-        return {"error": "model not found"}
-
-    model_path = MODELS[req.model]
-
-    response = run_model(model_path, req.prompt)
-
     return {
-        "model": req.model,
-        "response": response,
-        "done": True
+        "models": [
+            {"name": "tinyllama"}
+        ]
     }
 
 
 # ---------------------------
-# Chat endpoint
+# chat endpoint
 # ---------------------------
 
 @app.post("/api/chat")
 def chat(req: ChatRequest):
 
-    log("CHAT REQUEST", req.dict())
-
-    if req.model not in MODELS:
-        return {"error": "model not found"}
-
-    model_path = MODELS[req.model]
-
     prompt = build_prompt(req.messages)
 
-    response = run_model(model_path, prompt)
+    response = requests.post(
+        "http://localhost:8080/completion",
+        json={
+            "prompt": prompt,
+            "n_predict": 200
+        }
+    )
+
+    data = response.json()
 
     return {
         "model": req.model,
         "message": {
             "role": "assistant",
-            "content": response
+            "content": data["content"]
         },
         "done": True
     }
 
-os.makedirs("models", exist_ok=True)
-
-MODEL_FILES = {
-    "tinyllama": (
-        "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF",
-        "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
-    )
-}
-
-for name, (repo, file) in MODEL_FILES.items():
 
-    path = f"models/{name}.gguf"
+# ---------------------------
+# generate endpoint
+# ---------------------------
 
-    if not os.path.exists(path):
+@app.post("/api/generate")
+def generate(req: GenerateRequest):
 
-        print(f"Downloading model {name} from {repo}")
+    response = requests.post(
+        "http://localhost:8080/completion",
+        json={
+            "prompt": req.prompt,
+            "n_predict": 200
+        }
+    )
 
-        downloaded = hf_hub_download(
-            repo_id=repo,
-            filename=file
-        )
+    data = response.json()
 
-        os.system(f"cp {downloaded} {path}")
+    return {
+        "model": req.model,
+        "response": data["content"],
+        "done": True
+    }
 
-        print(f"Model ready: {path}")
 
 # ---------------------------
-# Start server
+# start API
 # ---------------------------
 
 if __name__ == "__main__":