Update server.py

server.py

@@ -2,55 +2,114 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 import subprocess
 import uvicorn
+import os
 
 app = FastAPI()
 
-MODELS={
- "tinyllama":"models/tinyllama.gguf",
- "qwen":"models/qwen1.5b.gguf",
- "gemma":"models/gemma2b.gguf"
+# ---------------------------
+# Models
+# ---------------------------
+
+MODELS = {
+    "tinyllama": "models/tinyllama.gguf",
+    "qwen": "models/qwen1.5b.gguf",
+    "gemma": "models/gemma2b.gguf"
 }
 
+# ---------------------------
+# Ensure model folder exists
+# ---------------------------
+
+os.makedirs("models", exist_ok=True)
+
+
+# ---------------------------
+# Request models
+# ---------------------------
+
 class ChatRequest(BaseModel):
-    model:str
-    messages:list
+    model: str
+    messages: list
+
 
 class GenerateRequest(BaseModel):
-    model:str
-    prompt:str
+    model: str
+    prompt: str
+
+
+# ---------------------------
+# Prompt builder
+# ---------------------------
 
 def build_prompt(messages):
-    prompt=""
+    prompt = ""
+
     for m in messages:
-        role=m["role"]
-        content=m["content"]
-        prompt += f"{role}: {content}\n"
-    prompt+="assistant:"
+        role = m.get("role", "user")
+        content = m.get("content", "")
+
+        if role == "assistant":
+            prompt += f"assistant: {content}\n"
+        else:
+            prompt += f"user: {content}\n"
+
+    prompt += "assistant:"
+
     return prompt
 
 
 # ---------------------------
-# Root
+# Run llama.cpp
+# ---------------------------
+
+def run_model(model_path, prompt):
+
+    result = subprocess.run(
+        [
+            "./llama.cpp/build/bin/llama-cli",
+            "-m", model_path,
+            "-p", prompt,
+            "-n", "200",
+            "--no-display-prompt"
+        ],
+        capture_output=True,
+        text=True
+    )
+
+    output = result.stdout.strip()
+
+    # Clean assistant prefix
+    if "assistant:" in output:
+        output = output.split("assistant:")[-1].strip()
+
+    return output
+
+
+# ---------------------------
+# Root endpoint
 # ---------------------------
 
 @app.get("/")
 def root():
-    return {"status":"running"}
+    return {"status": "running"}
 
 
 # ---------------------------
-# Model list (Ollama /api/tags)
+# Model list (Ollama compatible)
 # ---------------------------
 
 @app.get("/api/tags")
 def list_models():
-    return {
-        "models":[
-            {"name":"tinyllama"},
-            {"name":"qwen"},
-            {"name":"gemma"}
-        ]
-    }
+
+    models = []
+
+    for name in MODELS.keys():
+        models.append({
+            "name": name,
+            "model": name
+        })
+
+    return {"models": models}
 
 
 # ---------------------------
@@ -58,25 +117,19 @@ def list_models():
 # ---------------------------
 
 @app.post("/api/generate")
-def generate(req:GenerateRequest):
+def generate(req: GenerateRequest):
 
-    model_path=MODELS[req.model]
+    if req.model not in MODELS:
+        return {"error": "model not found"}
 
-    result=subprocess.run(
-        [
-            "./llama.cpp/build/bin/llama-cli",
-            "-m",model_path,
-            "-p",req.prompt,
-            "-n","200"
-        ],
-        capture_output=True,
-        text=True
-    )
+    model_path = MODELS[req.model]
+
+    response = run_model(model_path, req.prompt)
 
     return {
-        "model":req.model,
-        "response":result.stdout,
-        "done":True
+        "model": req.model,
+        "response": response,
+        "done": True
     }
 
 
@@ -85,31 +138,30 @@ def generate(req:GenerateRequest):
 # ---------------------------
 
 @app.post("/api/chat")
-def chat(req:ChatRequest):
+def chat(req: ChatRequest):
 
-    model_path=MODELS[req.model]
-    prompt=build_prompt(req.messages)
+    if req.model not in MODELS:
+        return {"error": "model not found"}
 
-    result=subprocess.run(
-        [
-            "./llama.cpp/build/bin/llama-cli",
-            "-m",model_path,
-            "-p",prompt,
-            "-n","200"
-        ],
-        capture_output=True,
-        text=True
-    )
+    model_path = MODELS[req.model]
+
+    prompt = build_prompt(req.messages)
+
+    response = run_model(model_path, prompt)
 
     return {
-        "model":req.model,
-        "message":{
-            "role":"assistant",
-            "content":result.stdout
+        "model": req.model,
+        "message": {
+            "role": "assistant",
+            "content": response
         },
-        "done":True
+        "done": True
     }
 
 
+# ---------------------------
+# Start server
+# ---------------------------
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)