huanghai/dsProject
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

'commit'

Browse Source
This commit is contained in:
HuangHai
2025-09-07 12:58:25 +08:00
parent abca226568
commit d9c954a06e
2 changed files with 2 additions and 333 deletions
Download Patch File Download Diff File Expand all files Collapse all files

228
dsLightRag/Volcengine/T2_CreateIndex.py
View File

@@ -14,7 +14,8 @@ from dotenv import load_dotenv
from volcenginesdkarkruntime import Ark
from Config.Config import VOLC_SECRETKEY, VOLC_API_KEY, VOLC_ACCESSKEY
from Volcengine.VikingDBMemoryService import MEMORY_COLLECTION_NAME
from Volcengine.VikingDBMemoryService import MEMORY_COLLECTION_NAME, VikingDBMemoryException, VikingDBMemoryService, \
initialize_services, ensure_collection_exists
# 控制日志输出
logger = logging.getLogger('CollectionMemory')
@@ -26,231 +27,6 @@ if not logger.handlers:
handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
logger.addHandler(handler)
class VikingDBMemoryException(Exception):
def __init__(self, code, request_id, message=None):
self.code = code
self.request_id = request_id
self.message = "{}, code:{}，request_id:{}".format(message, self.code, self.request_id)
def __str__(self):
return self.message
class VikingDBMemoryService(Service):
_instance_lock = threading.Lock()
def __new__(cls, *args, **kwargs):
if not hasattr(VikingDBMemoryService, "_instance"):
with VikingDBMemoryService._instance_lock:
if not hasattr(VikingDBMemoryService, "_instance"):
VikingDBMemoryService._instance = object.__new__(cls)
return VikingDBMemoryService._instance
def __init__(self, host="api-knowledgebase.mlp.cn-beijing.volces.com", region="cn-beijing", ak="", sk="",
sts_token="", scheme='https',
connection_timeout=30, socket_timeout=30):
self.service_info = VikingDBMemoryService.get_service_info(host, region, scheme, connection_timeout,
socket_timeout)
self.api_info = VikingDBMemoryService.get_api_info()
super(VikingDBMemoryService, self).__init__(self.service_info, self.api_info)
if ak:
self.set_ak(ak)
if sk:
self.set_sk(sk)
if sts_token:
self.set_session_token(session_token=sts_token)
try:
self.get_body("Ping", {}, json.dumps({}))
except Exception as e:
raise VikingDBMemoryException(1000028, "missed", "host or region is incorrect".format(str(e))) from None
def setHeader(self, header):
api_info = VikingDBMemoryService.get_api_info()
for key in api_info:
for item in header:
api_info[key].header[item] = header[item]
self.api_info = api_info
@staticmethod
def get_service_info(host, region, scheme, connection_timeout, socket_timeout):
service_info = ServiceInfo(host, {"Host": host},
Credentials('', '', 'air', region), connection_timeout, socket_timeout,
scheme=scheme)
return service_info
@staticmethod
def get_api_info():
api_info = {
"CreateCollection": ApiInfo("POST", "/api/memory/collection/create", {}, {},
{'Accept': 'application/json', 'Content-Type': 'application/json'}),
"GetCollection": ApiInfo("POST", "/api/memory/collection/info", {}, {},
{'Accept': 'application/json', 'Content-Type': 'application/json'}),
"DropCollection": ApiInfo("POST", "/api/memory/collection/delete", {}, {},
{'Accept': 'application/json', 'Content-Type': 'application/json'}),
"UpdateCollection": ApiInfo("POST", "/api/memory/collection/update", {}, {},
{'Accept': 'application/json', 'Content-Type': 'application/json'}),
"SearchMemory": ApiInfo("POST", "/api/memory/search", {}, {},
{'Accept': 'application/json', 'Content-Type': 'application/json'}),
"AddSession": ApiInfo("POST", "/api/memory/session/add", {}, {},
{'Accept': 'application/json', 'Content-Type': 'application/json'}),
"Ping": ApiInfo("GET", "/api/memory/ping", {}, {},
{'Accept': 'application/json', 'Content-Type': 'application/json'}),
}
return api_info
def get_body(self, api, params, body):
if not (api in self.api_info):
raise Exception("no such api")
api_info = self.api_info[api]
r = self.prepare_request(api_info, params)
r.headers['Content-Type'] = 'application/json'
r.headers['Traffic-Source'] = 'SDK'
r.body = body
SignerV4.sign(r, self.service_info.credentials)
url = r.build()
resp = self.session.get(url, headers=r.headers, data=r.body,
timeout=(self.service_info.connection_timeout, self.service_info.socket_timeout))
if resp.status_code == 200:
return json.dumps(resp.json())
else:
raise Exception(resp.text.encode("utf-8"))
def get_body_exception(self, api, params, body):
try:
res = self.get_body(api, params, body)
except Exception as e:
try:
res_json = json.loads(e.args[0].decode("utf-8"))
except:
raise VikingDBMemoryException(1000028, "missed", "json load res error, res:{}".format(str(e))) from None
code = res_json.get("code", 1000028)
request_id = res_json.get("request_id", 1000028)
message = res_json.get("message", None)
raise VikingDBMemoryException(code, request_id, message)
if res == '':
raise VikingDBMemoryException(1000028, "missed",
"empty response due to unknown error, please contact customer service") from None
return res
def get_exception(self, api, params):
try:
res = self.get(api, params)
except Exception as e:
try:
res_json = json.loads(e.args[0].decode("utf-8"))
except:
raise VikingDBMemoryException(1000028, "missed", "json load res error, res:{}".format(str(e))) from None
code = res_json.get("code", 1000028)
request_id = res_json.get("request_id", 1000028)
message = res_json.get("message", None)
raise VikingDBMemoryException(code, request_id, message)
if res == '':
raise VikingDBMemoryException(1000028, "missed",
"empty response due to unknown error, please contact customer service") from None
return res
def create_collection(self, collection_name, description="", custom_event_type_schemas=[],
custom_entity_type_schemas=[], builtin_event_types=[], builtin_entity_types=[]):
params = {
"CollectionName": collection_name, "Description": description,
"CustomEventTypeSchemas": custom_event_type_schemas, "CustomEntityTypeSchemas": custom_entity_type_schemas,
"BuiltinEventTypes": builtin_event_types, "BuiltinEntityTypes": builtin_entity_types,
}
res = self.json("CreateCollection", {}, json.dumps(params))
return json.loads(res)
def get_collection(self, collection_name):
params = {"CollectionName": collection_name}
res = self.json("GetCollection", {}, json.dumps(params))
return json.loads(res)
def drop_collection(self, collection_name):
params = {"CollectionName": collection_name}
res = self.json("DropCollection", {}, json.dumps(params))
return json.loads(res)
def update_collection(self, collection_name, custom_event_type_schemas=[], custom_entity_type_schemas=[],
builtin_event_types=[], builtin_entity_types=[]):
params = {
"CollectionName": collection_name,
"CustomEventTypeSchemas": custom_event_type_schemas, "CustomEntityTypeSchemas": custom_entity_type_schemas,
"BuiltinEventTypes": builtin_event_types, "BuiltinEntityTypes": builtin_entity_types,
}
res = self.json("UpdateCollection", {}, json.dumps(params))
return json.loads(res)
def search_memory(self, collection_name, query, filter, limit=10):
params = {
"collection_name": collection_name,
"query": query,
"limit": limit,
"filter": filter,
}
res = self.json("SearchMemory", {}, json.dumps(params))
return json.loads(res)
def add_session(self, collection_name, session_id, messages, metadata, entities=None):
params = {
"collection_name": collection_name,
"session_id": session_id,
"messages": messages,
"metadata": metadata,
}
if entities is not None:
params["entities"] = entities
res = self.json("AddSession", {}, json.dumps(params))
return json.loads(res)
def initialize_services():
load_dotenv()
ak = VOLC_ACCESSKEY
sk = VOLC_SECRETKEY
ark_api_key = VOLC_API_KEY
if not all([ak, sk, ark_api_key]):
raise ValueError("必须在环境变量中设置 VOLC_ACCESSKEY, VOLC_SECRETKEY, 和 ARK_API_KEY。")
memory_service = VikingDBMemoryService(ak=ak, sk=sk)
llm_client = Ark(
base_url="https://ark.cn-beijing.volces.com/api/v3",
api_key=ark_api_key,
)
return memory_service, llm_client
def ensure_collection_exists(memory_service, collection_name):
"""检查记忆集合是否存在，如果不存在则创建。"""
try:
memory_service.get_collection(collection_name)
logger.info(f"记忆集合 '{collection_name}' 已存在。")
except Exception as e:
error_message = str(e)
if "collection not exist" in error_message:
logger.info(f"记忆集合 '{collection_name}' 未找到，正在创建...")
try:
memory_service.create_collection(
collection_name=collection_name,
description="中文情感陪伴场景测试",
builtin_event_types=["sys_event_v1", "sys_profile_collect_v1"],
builtin_entity_types=["sys_profile_v1"]
)
logger.info(f"记忆集合 '{collection_name}' 创建成功。")
logger.info("等待集合准备就绪...")
except Exception as create_e:
logger.info(f"创建集合失败: {create_e}")
raise
else:
logger.info(f"检查集合时出错: {e}")
raise
def search_relevant_memories(memory_service, collection_name, user_id, query):