'commit'
This commit is contained in:
195
dsLightRag/Volcengine/VikingDBMemoryService.py
Normal file
195
dsLightRag/Volcengine/VikingDBMemoryService.py
Normal file
@@ -0,0 +1,195 @@
|
||||
"""
|
||||
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)
|
||||
Reference in New Issue
Block a user