codefuse-chatbot/coagent/orm/commands/code_base_cds.py

83 lines
2.4 KiB
Python

# encoding: utf-8
'''
@author: 温进
@file: code_base_cds.py.py
@time: 2023/10/23 下午4:34
@desc:
'''
from loguru import logger
from coagent.orm.db import with_session
from coagent.orm.schemas.base_schema import CodeBaseSchema
@with_session
def add_cb_to_db(session, code_name, code_path, code_graph_node_num, code_file_num, do_interpret):
do_interpret = 'YES' if do_interpret else 'NO'
# 增:创建知识库实例
cb = session.query(CodeBaseSchema).filter_by(code_name=code_name).first()
if not cb:
cb = CodeBaseSchema(code_name=code_name, code_path=code_path, code_graph_node_num=code_graph_node_num,
code_file_num=code_file_num, do_interpret=do_interpret)
session.add(cb)
else:
cb.code_path = code_path
cb.code_graph_node_num = code_graph_node_num
return True
@with_session
def list_cbs_from_db(session):
'''
查:查询实例
'''
cbs = session.query(CodeBaseSchema.code_name).all()
cbs = [cb[0] for cb in cbs]
return cbs
@with_session
def cb_exists(session, code_name):
'''
判断是否存在
'''
cb = session.query(CodeBaseSchema).filter_by(code_name=code_name).first()
status = True if cb else False
return status
@with_session
def load_cb_from_db(session, code_name):
cb = session.query(CodeBaseSchema).filter_by(code_name=code_name).first()
if cb:
code_name, code_path, code_graph_node_num, do_interpret = cb.code_name, cb.code_path, cb.code_graph_node_num, cb.do_interpret
else:
code_name, code_path, code_graph_node_num = None, None, None, None
return code_name, code_path, code_graph_node_num, do_interpret
@with_session
def delete_cb_from_db(session, code_name):
cb = session.query(CodeBaseSchema).filter_by(code_name=code_name).first()
if cb:
session.delete(cb)
return True
@with_session
def get_cb_detail(session, code_name: str) -> dict:
cb: CodeBaseSchema = session.query(CodeBaseSchema).filter_by(code_name=code_name).first()
logger.info(cb)
logger.info('code_name={}'.format(cb.code_name))
if cb:
return {
"code_name": cb.code_name,
"code_path": cb.code_path,
"code_graph_node_num": cb.code_graph_node_num,
'code_file_num': cb.code_file_num,
'do_interpret': cb.do_interpret
}
else:
return {
}