|
|
| import os |
| import json |
| import time |
| import uuid |
| import threading |
| from typing import Any |
|
|
| from fastapi import FastAPI, HTTPException |
| from fastapi.responses import StreamingResponse |
| from huggingface_hub import HfApi, hf_hub_download |
| from llama_cpp import Llama |
|
|
|
|
| MODEL_ID = os.getenv("MODEL_ID", "zelapi/lilo-cyber") |
| MODEL_NAME = os.getenv("MODEL_NAME", "lilo-cyber") |
| MODEL_FILE = os.getenv("MODEL_FILE", "").strip() |
|
|
| N_CTX = int(os.getenv("N_CTX", "1024")) |
| N_THREADS = int(os.getenv("N_THREADS", str(os.cpu_count() or 2))) |
| MAX_TOKENS = int(os.getenv("MAX_TOKENS", "256")) |
|
|
| SYSTEM_PROMPT = ( |
| "Kamu adalah Lilo Cyber, asisten AI yang membantu pemrograman, " |
| "keamanan siber defensif, analisis log, dan DevOps. " |
| "Jawab dengan jelas, akurat, dan gunakan bahasa pengguna." |
| ) |
|
|
| api = FastAPI( |
| title="Lilo Cyber API", |
| description="OpenAI-compatible API for Lilo Cyber", |
| version="1.0.0", |
| ) |
|
|
| generation_lock = threading.RLock() |
| llm = None |
|
|
|
|
| def get_model(): |
| global llm |
|
|
| if llm is not None: |
| return llm |
|
|
| with generation_lock: |
| if llm is not None: |
| return llm |
|
|
| filename = MODEL_FILE |
|
|
| if not filename: |
| files = HfApi().list_repo_files(MODEL_ID) |
|
|
| candidates = [ |
| name for name in files |
| if name.lower().endswith(".gguf") |
| ] |
|
|
| if not candidates: |
| raise RuntimeError( |
| f"Tidak ditemukan file GGUF di {MODEL_ID}" |
| ) |
|
|
| preferred = [ |
| name for name in candidates |
| if "Q4_K_M" in name.upper() |
| ] |
|
|
| filename = preferred[0] if preferred else candidates[0] |
|
|
| model_path = hf_hub_download( |
| repo_id=MODEL_ID, |
| filename=filename, |
| ) |
|
|
| llm = Llama( |
| model_path=model_path, |
| n_ctx=N_CTX, |
| n_threads=max(1, N_THREADS), |
| n_batch=128, |
| verbose=False, |
| ) |
|
|
| return llm |
|
|
|
|
| def prepare_messages(messages): |
| if not isinstance(messages, list) or not messages: |
| raise HTTPException( |
| status_code=400, |
| detail="messages must be a non-empty array", |
| ) |
|
|
| result = [ |
| { |
| "role": "system", |
| "content": SYSTEM_PROMPT, |
| } |
| ] |
|
|
| for message in messages: |
| if not isinstance(message, dict): |
| raise HTTPException( |
| status_code=400, |
| detail="Each message must be an object", |
| ) |
|
|
| role = message.get("role") |
| content = message.get("content") |
|
|
| if role == "system": |
| continue |
|
|
| if role not in ("user", "assistant"): |
| raise HTTPException( |
| status_code=400, |
| detail=f"Unsupported role: {role}", |
| ) |
|
|
| if not isinstance(content, str): |
| raise HTTPException( |
| status_code=400, |
| detail="Only string message content is supported", |
| ) |
|
|
| result.append({ |
| "role": role, |
| "content": content, |
| }) |
|
|
| if len(result) == 1: |
| raise HTTPException( |
| status_code=400, |
| detail="At least one user or assistant message is required", |
| ) |
|
|
| return result |
|
|
|
|
| def completion_parameters(body: dict[str, Any]): |
| model = body.get("model", MODEL_NAME) |
|
|
| if model != MODEL_NAME: |
| raise HTTPException( |
| status_code=404, |
| detail=f"Model not found: {model}", |
| ) |
|
|
| messages = prepare_messages(body.get("messages")) |
| max_tokens = body.get("max_tokens", MAX_TOKENS) |
| temperature = body.get("temperature", 0.7) |
| top_p = body.get("top_p", 0.9) |
|
|
| if ( |
| not isinstance(max_tokens, int) |
| or isinstance(max_tokens, bool) |
| ): |
| raise HTTPException( |
| status_code=400, |
| detail="max_tokens must be an integer", |
| ) |
|
|
| if not isinstance(temperature, (int, float)): |
| raise HTTPException( |
| status_code=400, |
| detail="temperature must be a number", |
| ) |
|
|
| if not isinstance(top_p, (int, float)): |
| raise HTTPException( |
| status_code=400, |
| detail="top_p must be a number", |
| ) |
|
|
| return { |
| "messages": messages, |
| "max_tokens": max(1, min(max_tokens, 512)), |
| "temperature": max(0.0, min(float(temperature), 2.0)), |
| "top_p": max(0.0, min(float(top_p), 1.0)), |
| "repeat_penalty": 1.1, |
| } |
|
|
|
|
| @api.get("/") |
| def root(): |
| return { |
| "name": "Lilo Cyber API", |
| "model": MODEL_NAME, |
| "endpoints": [ |
| "/health", |
| "/v1/models", |
| "/v1/chat/completions", |
| ], |
| } |
|
|
|
|
| @api.get("/health") |
| def health(): |
| return { |
| "status": "ok", |
| "model": MODEL_NAME, |
| "runtime": "cpu", |
| "loaded": llm is not None, |
| } |
|
|
|
|
| @api.get("/v1/models") |
| def list_models(): |
| return { |
| "object": "list", |
| "data": [ |
| { |
| "id": MODEL_NAME, |
| "object": "model", |
| "created": int(time.time()), |
| "owned_by": "zelapi", |
| } |
| ], |
| } |
|
|
|
|
| @api.post("/v1/chat/completions") |
| def chat_completions(body: dict[str, Any]): |
| params = completion_parameters(body) |
| request_id = "chatcmpl-" + uuid.uuid4().hex |
| created = int(time.time()) |
|
|
| if body.get("stream", False): |
| def generate(): |
| try: |
| with generation_lock: |
| model = get_model() |
|
|
| yield "data: " + json.dumps({ |
| "id": request_id, |
| "object": "chat.completion.chunk", |
| "created": created, |
| "model": MODEL_NAME, |
| "choices": [{ |
| "index": 0, |
| "delta": {"role": "assistant"}, |
| "finish_reason": None, |
| }], |
| }) + "\n\n" |
|
|
| for chunk in model.create_chat_completion( |
| **params, |
| stream=True, |
| ): |
| choice = chunk["choices"][0] |
|
|
| payload = { |
| "id": request_id, |
| "object": "chat.completion.chunk", |
| "created": created, |
| "model": MODEL_NAME, |
| "choices": [{ |
| "index": 0, |
| "delta": choice.get("delta", {}), |
| "finish_reason": choice.get("finish_reason"), |
| }], |
| } |
|
|
| yield ( |
| "data: " |
| + json.dumps(payload, ensure_ascii=False) |
| + "\n\n" |
| ) |
|
|
| yield "data: [DONE]\n\n" |
|
|
| except GeneratorExit: |
| raise |
|
|
| except Exception: |
| error = { |
| "error": { |
| "message": "Model inference failed", |
| "type": "server_error", |
| } |
| } |
|
|
| yield "data: " + json.dumps(error) + "\n\n" |
|
|
| return StreamingResponse( |
| generate(), |
| media_type="text/event-stream", |
| headers={ |
| "Cache-Control": "no-cache", |
| "X-Accel-Buffering": "no", |
| }, |
| ) |
|
|
| with generation_lock: |
| result = get_model().create_chat_completion(**params) |
|
|
| result["id"] = request_id |
| result["created"] = created |
| result["model"] = MODEL_NAME |
|
|
| return result |
|
|
