File size: 5,622 Bytes
e8ab4b6
 
 
3ced951
 
e8ab4b6
 
 
 
 
 
ed84950
e8ab4b6
 
 
 
 
 
 
 
3ced951
 
 
 
ed84950
3ced951
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e8ab4b6
67d45db
 
e8ab4b6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3ced951
e8ab4b6
 
 
3ced951
 
 
 
 
 
ed84950
3ced951
 
 
 
 
 
 
 
 
 
ed84950
3ced951
 
 
 
 
 
d653078
3ced951
 
d653078
3ced951
 
 
 
d653078
 
 
 
 
 
 
 
 
 
 
3ced951
 
 
 
 
 
e8ab4b6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3ced951
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
from fastapi import FastAPI, Request, HTTPException, Depends
from fastapi.security import APIKeyHeader
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, HTMLResponse
from fastapi.openapi.utils import get_openapi
import os
import json
import time
import uuid
from llama_cpp import Llama

app = FastAPI(title="KAI API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

def custom_openapi():
    if app.openapi_schema:
        return app.openapi_schema
    schema = get_openapi(
        title="KAI API",
        version="0.1.0",
        routes=app.routes,
    )
    schema["components"]["securitySchemes"] = {
        "ApiKeyAuth": {
            "type": "apiKey",
            "in": "header",
            "name": "X-API-Key"
        }
    }
    schema["security"] = [{"ApiKeyAuth": []}]
    app.openapi_schema = schema
    return schema

app.openapi = custom_openapi

API_KEY = os.environ.get("QWEN_API_KEY")
MODEL_PATH = "/models/Qwen2.5-7B-Instruct-Q4_K_M.gguf"

print(f"Loading model from {MODEL_PATH}...")
llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=4096,
    n_threads=8,
    n_gpu_layers=-1,
    verbose=False
)
print("Model loaded successfully.")

api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)

async def verify_key(api_key: str = Depends(api_key_header)):
    if not API_KEY:
        raise HTTPException(status_code=500, detail="QWEN_API_KEY not set in HF Secrets")
    if api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key")
    return api_key

@app.get("/", response_class=HTMLResponse)
async def ui():
    return """
<!DOCTYPE html>
<html>
<head>
    <title>KAI API Tester</title>
    <style>
        body { font-family: Arial; max-width: 800px; margin: 40px auto; padding: 20px; background: #1a1a1a; color: #fff; }
        input, textarea { width: 100%; padding: 10px; margin: 8px 0; background: #2a2a2a; color: #fff; border: 1px solid #444; border-radius: 6px; box-sizing: border-box; }
        button { background: #ff6b35; color: white; padding: 12px 24px; border: none; border-radius: 6px; cursor: pointer; font-size: 16px; }
        pre { background: #2a2a2a; padding: 15px; border-radius: 6px; overflow-x: auto; white-space: pre-wrap; }
        label { color: #aaa; font-size: 14px; }
        h1 { color: #ff6b35; }
    </style>
</head>
<body>
    <h1>🤖 KAI API Tester</h1>
    <label>API Key</label>
    <input type="password" id="apikey" placeholder="Twój klucz API">
    <label>System prompt</label>
    <input type="text" id="system" value="You are a helpful assistant.">
    <label>Wiadomość</label>
    <textarea id="message" rows="4" placeholder="Napisz coś..."></textarea>
    <button onclick="sendMessage()">Wyślij 🚀</button>
    <pre id="output">Odpowiedź pojawi się tutaj...</pre>
    <script>
        async function sendMessage() {
            const key = document.getElementById('apikey').value;
            const system = document.getElementById('system').value;
            const msg = document.getElementById('message').value;
            document.getElementById('output').textContent = 'Ładowanie...';
            try {
                const res = await fetch('/v1/chat/completions', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json', 'X-API-Key': key },
                    body: JSON.stringify({ messages: [{ role: 'system', content: system }, { role: 'user', content: msg }], max_tokens: 256 })
                });
                const data = await res.json();
                document.getElementById('output').textContent = data.choices?.[0]?.message?.content || JSON.stringify(data, null, 2);
            } catch(e) {
                document.getElementById('output').textContent = 'Błąd: ' + e.message;
            }
        }
    </script>
</body>
</html>
"""

@app.get("/health")
async def health():
    return {
        "status": "ok",
        "model": MODEL_PATH,
        "api_key_loaded": API_KEY is not None
    }

def format_prompt(messages):
    prompt = ""
    for msg in messages:
        role = msg["role"]
        content = msg["content"]
        prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"
    prompt += "<|im_start|>assistant\n"
    return prompt

@app.post("/v1/chat/completions")
async def chat_completions(request: Request, api_key: str = Depends(verify_key)):
    body = await request.json()
    messages = body.get("messages", [])
    max_tokens = body.get("max_tokens", 512)
    temperature = body.get("temperature", 0.7)
    top_p = body.get("top_p", 0.95)

    if not messages:
        raise HTTPException(status_code=400, detail="messages is required")

    prompt = format_prompt(messages)
    output = llm(
        prompt,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        stop=["<|im_end|>"]
    )

    response_text = output["choices"][0]["text"].strip()

    return JSONResponse({
        "id": f"chatcmpl-{uuid.uuid4()}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": MODEL_PATH,
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": response_text},
            "finish_reason": output["choices"][0]["finish_reason"]
        }],
        "usage": {
            "prompt_tokens": output["usage"]["prompt_tokens"],
            "completion_tokens": output["usage"]["completion_tokens"],
            "total_tokens": output["usage"]["total_tokens"]
        }
    })