Spaces:
Sleeping
Sleeping
| 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 | |
| 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> | |
| """ | |
| 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 | |
| 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"] | |
| } | |
| }) | |