import litserve as ls
from typing import List, Dict, Any
from llama_cpp import Llama
from fastapi import Depends, HTTPException
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer


class SimpleLitAPI(ls.LitAPI):
    def setup(self, device):
        self.llm = Llama(
        model_path="Llama-3.2-1B-Instruct-Q4_K_M.gguf",
        verbose=False 
    )
    def decode_request(self, request):
        return request

    def predict(self, input):
        temperature = input.get("temperature", 1.0) 
        max_tokens = input.get("max_tokens", 100)
        top_k = input.get("top_k", 50)  
        top_p = input.get("top_p", 0.9)       
        repeat_penalty = input.get("repeat_penalty", 1.0) 
        return self.llm.create_chat_completion(
            messages=input["messages"],
            temperature=temperature,
            max_tokens=max_tokens,
            top_k=top_k,
            top_p=top_p,
            repeat_penalty=repeat_penalty,
        )

    def encode_response(self, output):
        return {"output": output}
    
    def authorize(self, auth: HTTPAuthorizationCredentials = Depends(HTTPBearer())):
        if auth.scheme != "Bearer" or auth.credentials != "1234":
            raise HTTPException(status_code=401, detail="Bad token")
        
if __name__ == "__main__":
    api = SimpleLitAPI()
    server = ls.LitServer(api, accelerator="cpu")
    server.run(port=7860)