""" pip install volcengine pip install --upgrade "volcengine-python-sdk[ark]" """ import json import threading from volcengine.ApiInfo import ApiInfo from volcengine.Credentials import Credentials from volcengine.base.Service import Service from volcengine.ServiceInfo import ServiceInfo from volcengine.auth.SignerV4 import SignerV4 from volcengine.base.Request import Request 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)