isekai-toolkit/test/embedding_search.py

import base
import utils.local as local
from service.database import DatabaseService

from service.embedding_search import EmbeddingSearchService

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

    async with EmbeddingSearchService(dbs, "代号:曙光的世界/黄昏的阿瓦隆") as embedding_search:
        await embedding_search.prepare_update_index()

        async def on_index_progress(current, length):
            print("\r索引进度：%.1f%%" % (current / length * 100), end="", flush=True)

        print("")
        await embedding_search.update_page_index(on_index_progress)
        print("")

        while True:
            query = input("请输入要搜索的问题 (.exit 退出)：")
            if query == ".exit":
                break
            res, token_usage = await embedding_search.search(query, 5)
            total_length = 0
            if res:
                for one in res:
                    total_length += len(one["markdown"])
                    print("%s, distance=%.4f" % (one["markdown"], one["distance"]))
            else:
                print("未搜索到相关内容")

            print("总长度：%d" % total_length)
    
    await local.noawait.end()

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