import os 
from zhipuai import ZhipuAI
from typing import List, Optional, Tuple, Dict
from http import HTTPStatus



History = List[Tuple[str, str]]
Messages = List[Dict[str, str]]

default_system = 'You are a helpful assistant.'

YOUR_API_TOKEN = os.getenv('YOUR_API_TOKEN')


SYSTEM = "assistant"
USER = "user"
ASSISTANT = "assistant"

def clear_session() -> History:
    return '', []

def modify_system_session(system: str) -> str:
    if system is None or len(system) == 0:
        system = default_system
    return system, system, []

def history_to_messages(history: History, system: str) -> Messages:
    messages = [{'role': SYSTEM, 'content': system}]
    for h in history:
        messages.append({'role': USER, 'content': h[0]})
        messages.append({'role': ASSISTANT, 'content': h[1]})
    return messages


def messages_to_history(messages: Messages) -> Tuple[str, History]:
    assert messages[0]['role'] == SYSTEM
    system = messages[0]['content']
    history = []
    for q, r in zip(messages[1::2], messages[2::2]):
        history.append([q['content'], r['content']])
    return system, history


def model_chat(query: Optional[str], history: Optional[History], system: str
) -> Tuple[str, str, History]:
    if query is None:
        query = ''
    if history is None:
        history = []
    messages = history_to_messages(history, system)
    messages.append({'role': USER, 'content': query})

    client = ZhipuAI(api_key=YOUR_API_TOKEN) # 填写您自己的APIKey
    gen = client.chat.completions.create(
        model="glm-4",  # 填写需要调用的模型名称
        messages=messages,
        stream=True
    )

    content = ""
    for response in gen:
        # print(response)
        role = response.choices[0].delta.role
        content += response.choices[0].delta.content

        system, history = messages_to_history(messages + [{'role': role, 'content': content}])
        yield '', history, system


if __name__ == '__main__':
    output = model_chat("who are you?", [], "You are a helpful assistant.")
    for o in output:
        print(o)