#测试GPTCache，进行精确匹配缓存或基于语义相似性缓存结果
import langchain
import gptcache
import time
from langchain.llms import OpenAI
from gptcache.processor.pre import get_prompt
from gptcache.manager.factory import get_data_manager
from langchain.cache import GPTCache

llm = OpenAI(model_name="text-davinci-003", n=2, best_of=2)

# Avoid multiple caches using the same file, causing different llm model caches to affect each other
i = 0
file_prefix = "data_map"

def init_gptcache_map(cache_obj: gptcache.Cache):
    global i
    cache_path = f'{file_prefix}_{i}.txt'
    cache_obj.init(
        pre_embedding_func=get_prompt,
        data_manager=get_data_manager(data_path=cache_path),
    )
    i += 1

langchain.llm_cache = GPTCache(init_gptcache_map)

for i in range(20):
    start = time.perf_counter()
    prompt = "男生有2人，女生有{:d}人，一共多少人？".format(i)
    print("男生有2人，女生有{:d}人, {:s}。 suspend: {:0.4f}".format(i, llm(prompt), time.perf_counter() - start))