mirror of
https://github.com/1Panel-dev/MaxKB.git
synced 2025-12-26 01:33:05 +00:00
53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
# coding=utf-8
|
||
"""
|
||
@project: maxkb
|
||
@Author：虎
|
||
@file： embedding_config.py
|
||
@date：2023/10/23 16:03
|
||
@desc:
|
||
"""
|
||
import time
|
||
|
||
from common.cache.mem_cache import MemCache
|
||
|
||
|
||
class EmbeddingModelManage:
|
||
cache = MemCache('model', {})
|
||
up_clear_time = time.time()
|
||
|
||
@staticmethod
|
||
def get_model(_id, get_model):
|
||
model_instance = EmbeddingModelManage.cache.get(_id)
|
||
if model_instance is None:
|
||
model_instance = get_model(_id)
|
||
EmbeddingModelManage.cache.set(_id, model_instance, timeout=60 * 30)
|
||
return model_instance
|
||
# 续期
|
||
EmbeddingModelManage.cache.touch(_id, timeout=60 * 30)
|
||
EmbeddingModelManage.clear_timeout_cache()
|
||
return model_instance
|
||
|
||
@staticmethod
|
||
def clear_timeout_cache():
|
||
if time.time() - EmbeddingModelManage.up_clear_time > 60:
|
||
EmbeddingModelManage.cache.clear_timeout_data()
|
||
|
||
|
||
class VectorStore:
|
||
from embedding.vector.pg_vector import PGVector
|
||
from embedding.vector.base_vector import BaseVectorStore
|
||
instance_map = {
|
||
'pg_vector': PGVector,
|
||
}
|
||
instance = None
|
||
|
||
@staticmethod
|
||
def get_embedding_vector() -> BaseVectorStore:
|
||
from embedding.vector.pg_vector import PGVector
|
||
if VectorStore.instance is None:
|
||
from smartdoc.const import CONFIG
|
||
vector_store_class = VectorStore.instance_map.get(CONFIG.get("VECTOR_STORE_NAME"),
|
||
PGVector)
|
||
VectorStore.instance = vector_store_class()
|
||
return VectorStore.instance
|