昇腾社区首页
中文
注册

类功能

功能描述

主要提供用户问题和答案的缓存存储、缓存更新、缓存落盘功能。

函数原型

from mx_rag.cache import MxRAGCache
MxRAGCache(cache_name, config)

输入参数说明

参数名

数据类型

可选/必选

说明

cache_name

str

必选

缓存名字,该命名会体现在落盘文件名中,字符串长度范围:(0, 64)。

只能是字母、数字或下划线。

config

CacheConfig/SimilarityCacheConfig

必选

缓存配置,参考缓存配置

调用示例

import json
from pymilvus import MilvusClient
from mx_rag.cache import CacheConfig, SimilarityCacheConfig
from mx_rag.cache import EvictPolicy
from mx_rag.cache import MxRAGCache
from mx_rag.utils import ClientParam

dim = 1024

cache_config = CacheConfig(
    cache_size=100,
    eviction_policy=EvictPolicy.LRU,
    data_save_folder="path_to_cache_save_folder"
)
mxrag_l1_cache = MxRAGCache("memory_cache", cache_config)
cache = MxRAGCache("memory_cache", cache_config)
# 检查cache实例是否初始化成功
cache_obj = cache.get_obj()
if cache_obj is None:
    print(f"cache init failed")
similarity_config = SimilarityCacheConfig(
    vector_config={
        "vector_type": "milvus_db",
        "x_dim": dim,
        "client": MilvusClient(https://<ip>:<port>, user="xxx", password=getpass.getpass(), token="xxx"),
        "collection_name": "mxrag_cache_123",  # milvus db的标签
        "use_http": False,
        "param": None
    },
    cache_config="sqlite",
    emb_config={
        "embedding_type": "tei_embedding",
        "url": "https://<ip>:<port>/embed",  # tei_embedding 服务的IP地址和侦听端口
        "client_param": ClientParam(ca_file="/path/to/ca.crt")
    },
    similarity_config={
        "similarity_type": "tei_reranker",
        "url": "https://<ip>:<port>/rerank",  # tei_reranker 服务的IP地址和侦听端口
        "client_param": ClientParam(ca_file="/path/to/ca.crt")
    },
    retrieval_top_k=1,
    cache_size=100,
    auto_flush=100,
    similarity_threshold=0.70,
    data_save_folder="path_to_cache_save_folder",
    disable_report=True,
    eviction_policy=EvictPolicy.FIFO
)
similarity_cache = MxRAGCache("similarity_cache", similarity_config)
# 设置缓存级联
cache.join(similarity_cache)
# 设置缓存每条的字符限制为4000个字符
cache.set_cache_limit(4000)
# 设置是否详细显示缓存过程
cache.set_verbose(False)
# 手动更新缓存
cache.update("小明的爸爸是谁?", json.dumps({"小明的爸爸是谁?": "小明的爸爸名字是大明"}))
# 精确匹配结果
res = cache.search("小明的爸爸是谁?")
print(f"memory match res: {res}")
# 语义近似匹配结果
res = cache.search("小明的爸爸叫什么名字")
print(f"similarity match res: {res}")
# 手动调用flush 将缓存落盘,也会按照auto_flush配置进行自动落盘
cache.flush()
# 删除已落盘的文件和数据
cache.clear()