from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs
from llama_cpp import Llama

SYSTEM_PROMPT = "Ты — русскоязычный автоматический ассистент и профессиональный редактор. Ты выполняешь указания пользователя в точности, соблюдая все детали задания. Ты хорошо умеешь обобщать текст, выделяя только основной смысл."

def get_message_tokens(llm, role, content):
    content = f"{role}\n{content}\n</s>"
    content = content.encode("utf-8")
    message_tokens = llm.tokenize(content, special=True)
    return message_tokens

def get_system_tokens(llm):
    system_message = {
        "role": "system",
        "content": SYSTEM_PROMPT
    }
    return get_message_tokens(llm, **system_message)

llm = Llama(model_path="/home/oluser/olapp/model-q4_K.gguf", n_ctx=2048, n_parts=1)

system_tokens = get_system_tokens(llm)

class OlHandler(BaseHTTPRequestHandler):

    def do_GET(self):
        query_components = parse_qs(urlparse(self.path).query)
        q = query_components["q"][0]

        tokens = system_tokens
        llm.eval(tokens)

        # TODO: add few shot

        message_tokens = get_message_tokens(llm=llm, role="user",
            content="Напиши краткое изложение текста, представленного ниже, в одном предложении.\nПредложение должно быть лаконичным о отражать основной смысл события или новости.\n\n" + q )
        role_tokens = llm.tokenize("bot\n".encode("utf-8"), special=True)
        tokens += message_tokens + role_tokens
        # full_prompt = llm.detokenize(tokens)

        generator = llm.generate(
            tokens,
            top_k=30,
            top_p=.9,
            temp=.2,
            repeat_penalty=1.21
        )

        answ = ""

        self.send_response(200)
        self.send_header('Content-type','text/plain; charset=utf-8')
        self.end_headers()
        
        for tok in generator:
            token_str = llm.detokenize([tok]).decode("utf-8", errors="ignore")
            tokens.append(tok)
            if tok == llm.token_eos():
                break
            #self.wfile.write(token_str.encode('utf-8'))
            answ += token_str
            #print(token_str, end="", flush=True)
        
        #output = llm.create_completion(
        #  q,
        #  max_tokens=32,
        #  echo=False
        #)["choices"][0]["text"]

        #self.wfile.write(output.encode('utf-8'))
        self.wfile.write(answ.encode('utf-8'))
        return


if __name__ == '__main__':
    olserver = HTTPServer(('0.0.0.0', 7860), OlHandler)
    print('Starting server at http://0.0.0.0:7860')
    olserver.serve_forever()