isekai-toolkit/test/chatcomplete.py

import traceback
import base
import utils.local as local
from service.chat_complete import ChatCompleteService
from service.database import DatabaseService

from service.tiktoken import TikTokenService

async def main():
    dbs = await DatabaseService.create()
    tiktoken = await TikTokenService.create()

    async with ChatCompleteService(dbs, "代号:曙光的世界/黄昏的阿瓦隆") as chat_complete:
        question = "你是谁？"
        question_tokens = await tiktoken.get_tokens(question)

        try:
            prepare_res = await chat_complete.prepare_chat_complete(question, None, 1, question_tokens, {
                "distance_limit": 20,
                "limit": 10
            })

            print(prepare_res)
            async def on_message(message: str):
                # print(message)
                pass

            res = await chat_complete.finish_chat_complete(on_message)
            print(res)
        except Exception as err:
            print(err)
            traceback.print_exc()

    await local.noawait.end()

if __name__ == '__main__':
    local.loop.run_until_complete(main())