from openai import OpenAI
# import traceback


def chat_oai_stream(
    base_url="http://127.0.0.1:8000/v1",
    api_key="dummy_key",
    model="/data/models/Qwen-7B-Chat-Int4",
    prompt="Who are you?",
):
    """Chat with OpenAI's GPT-3 model using the specified parameters.
    
    Args:
        base_url (str): The base URL for the OpenAI API. Default is "http://127.0.0.1:8000/v1".
        api_key (str): The API key for accessing the OpenAI API. Default is "dummy_key".
        model (str): The model ID to use for the chat. Default is "/data/models/Qwen-7B-Chat-Int4".
        prompt (str): The initial prompt for the chat conversation.
    
    Yields:
        str: The generated content from the chat conversation.
    
    """
    client = OpenAI(
        base_url=base_url,
        api_key=api_key,
    )
    response = client.chat.completions.create(
        model=model,
        messages=[{
            "role": "user",
            "content": prompt
        }],
        stream=True
    )
    
    for chunk in response:
        content = chunk.choices[0].delta.content
        yield content


def chat_oai_invoke(
    base_url="http://127.0.0.1:8000/v1",
    api_key="dummy_key",
    model="/data/models/Qwen-7B-Chat-Int4",
    prompt="Who are you?",
):
    """Invoke OpenAI chat API to generate a response based on the given prompt.
    
    Args:
        base_url (str): The base URL of the OpenAI API. Default is "http://127.0.0.1:8000/v1".
        api_key (str): The API key for accessing the OpenAI API. Default is "dummy_key".
        model (str): The model to use for generating the response. Default is "/data/models/Qwen-7B-Chat-Int4".
        prompt (str): The prompt message to start the conversation. Default is "Who are you?".
    
    Returns:
        str: The response generated by the OpenAI chat API based on the prompt.
    """
    client = OpenAI(
        base_url=base_url,
        api_key=api_key,
    )
    response = client.chat.completions.create(
        model=model,
        messages=[{
            "role": "user",
            "content": prompt
        }],
        stream=False
    )
    
    return response.choices[0].message.content

# client = OpenAI(
#     base_url="http://127.0.0.1:8000/v1",
#     api_key="dummy_key"  # 使用虚拟的API Key
# )


# response = client.chat.completions.create(
#     model="/data/models/Qwen-7B-Chat-Int4",
#     messages=[{"role": "user", "content": "给我讲个鬼故事"}],
#     stream=True
# )
# for chunk in response:
#     content = chunk.choices[0].delta.content
#     if content:
#         print(content, end='', flush=True)
# print()