from llama_index.core.base.llms.types import ChatMessage
from llama_index.llms.dashscope import DashScope
import asyncio

llm = DashScope(model_name="qwen-max") # 设置检索引擎生成回答时调用的大模型。

def test1():
    response = llm.complete("William Shakespeare is ")
    print(response)

async def test2():
    response = await llm.acomplete("William Shakespeare is ")
    print(response)


def test3():
    response = llm.stream_complete("William Shakespeare is ")
    for chunk in response:
        print(chunk)

def test4():
    handle = llm.stream_complete("William Shakespeare is ")

    for token in handle:
        print(token.delta, end="", flush=True)


def test5():
    messages = [
        ChatMessage(role="system", content="You are a helpful assistant."),
        ChatMessage(role="user", content="Tell me a joke."),
    ]
    chat_response = llm.chat(messages)
    print(chat_response)

if __name__ == '__main__':
    # test1()
    # asyncio.run(test2())
    # test3()
    # test4()
    test5()