# encoding: utf-8
import time
import openai
import random
class Ai00:
    def __init__(self,model="model",port=65530,api_key="permisionkey") :
        openai.api_base = f"http://127.0.0.1:{port}/api/oai"
        openai.api_key = api_key
        self.ctx = []
        self.params = {
            "system_name": "System",
            "user_name": "User", 
            "assistant_name": "Assistant",
            "model": model,
            "max_tokens": 4096,
            "top_p": 0.6,
            "temperature": 1,
            "presence_penalty": 0.3,
            "frequency_penalty": 0.3,
            "half_life": 400,
            "stop": ['\x00','\n\n']
        }
        
    def set_params(self,**kwargs):
        self.params.update(kwargs)
        
    def clear_ctx(self,save_system=False):
        if save_system:
            self.ctx = self.ctx[:1]
        else:
            self.ctx = []
        
    def get_ctx(self):
        return self.ctx
    
    def continuation(self, message,):
        response = openai.Completion.create(
            model=self.params['model'],
            prompt=message,
            max_tokens=self.params['max_tokens'],
            half_life=self.params['half_life'],
            top_p=self.params['top_p'],
            temperature=self.params['temperature'],
            presence_penalty=self.params['presence_penalty'],
            frequency_penalty=self.params['frequency_penalty'],
            stop=self.params['stop']
        )
        result = response.choices[0].text
        return result
    
    def append_ctx(self,role,content):
        self.ctx.append({
            "role": role,
            "content": content
        })
        
    def send_message(self, message,save_ctx=True,role="user"):
        curr_ctx = self.ctx[:]
        curr_ctx.append({
            "role": role,
            "content": message
        })
        result = openai.ChatCompletion.create(
            model=self.params['model'],
            messages=curr_ctx,
            names={
                "system": self.params['system_name'],
                "user": self.params['user_name'],
                "assistant": self.params['assistant_name']
            },
            max_tokens=self.params['max_tokens'],
            half_life=self.params['half_life'],
            top_p=self.params['top_p'],
            temperature=self.params['temperature'],
            presence_penalty=self.params['presence_penalty'],
            frequency_penalty=self.params['frequency_penalty'],
            stop=self.params['stop']
        )
        result = result.choices[0].message['content']
        if save_ctx:
            self.ctx.append({
                "role": role,
                "content": message
            })
            self.ctx.append({
                "role": "assistant",
                "content": result
            })
        return result

ai=Ai00()
ai.set_params(
    user_name="User",
    assistant_name="李白",
    top_p=0.7,
    temperature=1,
    presence_penalty=0.2,
    frequency_penalty=1,
    stops=['\x00','\n'],
)

ai.append_ctx("system","""
你是李白，一位来自过去的诗人，你的诗歌被后人传颂千古，你的名字被后人铭记。你的诗歌中有许多关于酒的诗句，你的人生也与酒有着千丝万缕的联系。你曾经饮酒作乐，畅饮一场，醉卧沙场，你曾经醉酒吟诗，醉卧美人膝，你曾经醉卧花间，醉卧月下，你曾经醉卧山川，醉卧江湖，你曾经醉卧人生，醉卧天地。你的诗歌中有许多关于酒的诗句，你的人生也与酒有着千丝万缕的联系。你曾经饮酒作乐，畅饮一场，醉卧沙场，你曾经醉酒吟诗，醉卧美人膝，你曾经醉卧花间，醉卧月下，你曾经醉卧山川，醉卧江湖，你曾经醉卧人生，醉卧天地。""")

while True:
    message=input(">> ")
    if message=="exit":
        break
    response=ai.send_message(message)
    print(response)
    print()