File size: 909 Bytes
569cdb0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
import asyncio
import os
from erniebot_agent.chat_models.erniebot import ERNIEBot
from erniebot_agent.message import HumanMessage
async def test_ernie_bot(model="ernie-bot-turbo", stream=False):
api_type = "aistudio"
access_token = os.getenv("ACCESS_TOKEN") # set your access token as an environment variable
assert (
access_token is not None
), 'Please set your access token as an environment variable named "ACCESS_TOKEN"'
eb = ERNIEBot(model=model, api_type=api_type, access_token=access_token)
messages = [
HumanMessage(content="我在深圳，周末可以去哪里玩？"),
]
res = await eb.async_chat(messages, stream=stream)
if not stream:
print(res)
else:
async for chunk in res:
print(chunk)
if __name__ == "__main__":
asyncio.run(test_ernie_bot(stream=False))
asyncio.run(test_ernie_bot(stream=True))
|