xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-08-15 06:50:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: 数据库接口支持

Browse Source
This commit is contained in:
Rock Chin
2023-03-18 12:57:36 +00:00
parent 0490ad9207
commit d056cb6769
1 changed files with 31 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/database/manager.py
+31 -16
View File
@@ -54,20 +54,27 @@ class DatabaseManager:
`last_interact_timestamp` bigint not null, `last_interact_timestamp` bigint not null,
`status` varchar(255) not null default 'on_going', `status` varchar(255) not null default 'on_going',
`default_prompt` text not null default '', `default_prompt` text not null default '',
`prompt` text not null `prompt` text not null,
`token_counts` text not null default '[]',
) )
""") """)
# 检查sessions表是否存在`default_prompt`字段 # 检查sessions表是否存在`default_prompt`字段, 检查是否存在`token_counts`字段
self.__execute__("PRAGMA table_info('sessions')") self.__execute__("PRAGMA table_info('sessions')")
columns = self.cursor.fetchall() columns = self.cursor.fetchall()
has_default_prompt = False has_default_prompt = False
has_token_counts = False
for field in columns: for field in columns:
if field[1] == 'default_prompt': if field[1] == 'default_prompt':
has_default_prompt = True has_default_prompt = True
if field[1] == 'token_counts':
has_token_counts = True
if has_default_prompt and has_token_counts:
break break
if not has_default_prompt: if not has_default_prompt:
self.__execute__("alter table `sessions` add column `default_prompt` text not null default ''") self.__execute__("alter table `sessions` add column `default_prompt` text not null default ''")
if not has_token_counts:
self.__execute__("alter table `sessions` add column `token_counts` text not null default '[]'")
self.__execute__(""" self.__execute__("""
@@ -89,7 +96,7 @@ class DatabaseManager:
# session持久化 # session持久化
def persistence_session(self, subject_type: str, subject_number: int, create_timestamp: int, def persistence_session(self, subject_type: str, subject_number: int, create_timestamp: int,
last_interact_timestamp: int, prompt: str, default_prompt: str = ''): last_interact_timestamp: int, prompt: str, default_prompt: str = '', token_counts: list = []):
"""持久化指定session""" """持久化指定session"""
# 检查是否已经有了此name和create_timestamp的session # 检查是否已经有了此name和create_timestamp的session
@@ -102,20 +109,20 @@ class DatabaseManager:
if count == 0: if count == 0:
sql = """ sql = """
insert into `sessions` (`name`, `type`, `number`, `create_timestamp`, `last_interact_timestamp`, `prompt`, `default_prompt`) insert into `sessions` (`name`, `type`, `number`, `create_timestamp`, `last_interact_timestamp`, `prompt`, `default_prompt`, `token_counts`)
values (?, ?, ?, ?, ?, ?, ?) values (?, ?, ?, ?, ?, ?, ?, ?)
""" """
self.__execute__(sql, self.__execute__(sql,
("{}_{}".format(subject_type, subject_number), subject_type, subject_number, create_timestamp, ("{}_{}".format(subject_type, subject_number), subject_type, subject_number, create_timestamp,
last_interact_timestamp, prompt, default_prompt)) last_interact_timestamp, prompt, default_prompt, json.dumps(token_counts)))
else: else:
sql = """ sql = """
update `sessions` set `last_interact_timestamp` = ?, `prompt` = ? update `sessions` set `last_interact_timestamp` = ?, `prompt` = ?, `token_counts` = ?
where `type` = ? and `number` = ? and `create_timestamp` = ? where `type` = ? and `number` = ? and `create_timestamp` = ?
""" """
self.__execute__(sql, (last_interact_timestamp, prompt, subject_type, self.__execute__(sql, (last_interact_timestamp, prompt, json.dumps(token_counts), subject_type,
subject_number, create_timestamp)) subject_number, create_timestamp))
# 显式关闭一个session # 显式关闭一个session
@@ -140,7 +147,7 @@ class DatabaseManager:
# 从数据库中加载所有还没过期的session # 从数据库中加载所有还没过期的session
config = pkg.utils.context.get_config() config = pkg.utils.context.get_config()
self.__execute__(""" self.__execute__("""
select `name`, `type`, `number`, `create_timestamp`, `last_interact_timestamp`, `prompt`, `status`, `default_prompt` select `name`, `type`, `number`, `create_timestamp`, `last_interact_timestamp`, `prompt`, `status`, `default_prompt`, `token_counts`
from `sessions` where `last_interact_timestamp` > {} from `sessions` where `last_interact_timestamp` > {}
""".format(int(time.time()) - config.session_expire_time)) """.format(int(time.time()) - config.session_expire_time))
results = self.cursor.fetchall() results = self.cursor.fetchall()
@@ -154,6 +161,7 @@ class DatabaseManager:
prompt = result[5] prompt = result[5]
status = result[6] status = result[6]
default_prompt = result[7] default_prompt = result[7]
token_counts = result[8]
# 当且仅当最后一个该对象的会话是on_going状态时，才会被加载 # 当且仅当最后一个该对象的会话是on_going状态时，才会被加载
if status == 'on_going': if status == 'on_going':
@@ -163,7 +171,8 @@ class DatabaseManager:
'create_timestamp': create_timestamp, 'create_timestamp': create_timestamp,
'last_interact_timestamp': last_interact_timestamp, 'last_interact_timestamp': last_interact_timestamp,
'prompt': prompt, 'prompt': prompt,
'default_prompt': default_prompt 'default_prompt': default_prompt,
'token_counts': json.loads(token_counts)
} }
else: else:
if session_name in sessions: if session_name in sessions:
@@ -175,7 +184,7 @@ class DatabaseManager:
def last_session(self, session_name: str, cursor_timestamp: int): def last_session(self, session_name: str, cursor_timestamp: int):
self.__execute__(""" self.__execute__("""
select `name`, `type`, `number`, `create_timestamp`, `last_interact_timestamp`, `prompt`, `status`, `default_prompt` select `name`, `type`, `number`, `create_timestamp`, `last_interact_timestamp`, `prompt`, `status`, `default_prompt`, `token_counts`
from `sessions` where `name` = '{}' and `last_interact_timestamp` < {} order by `last_interact_timestamp` desc from `sessions` where `name` = '{}' and `last_interact_timestamp` < {} order by `last_interact_timestamp` desc
limit 1 limit 1
""".format(session_name, cursor_timestamp)) """.format(session_name, cursor_timestamp))
@@ -192,6 +201,7 @@ class DatabaseManager:
prompt = result[5] prompt = result[5]
status = result[6] status = result[6]
default_prompt = result[7] default_prompt = result[7]
token_counts = result[8]
return { return {
'subject_type': subject_type, 'subject_type': subject_type,
@@ -199,14 +209,15 @@ class DatabaseManager:
'create_timestamp': create_timestamp, 'create_timestamp': create_timestamp,
'last_interact_timestamp': last_interact_timestamp, 'last_interact_timestamp': last_interact_timestamp,
'prompt': prompt, 'prompt': prompt,
'default_prompt': default_prompt 'default_prompt': default_prompt,
'token_counts': json.loads(token_counts)
} }
# 获取此session_name后一个session的数据 # 获取此session_name后一个session的数据
def next_session(self, session_name: str, cursor_timestamp: int): def next_session(self, session_name: str, cursor_timestamp: int):
self.__execute__(""" self.__execute__("""
select `name`, `type`, `number`, `create_timestamp`, `last_interact_timestamp`, `prompt`, `status`, `default_prompt` select `name`, `type`, `number`, `create_timestamp`, `last_interact_timestamp`, `prompt`, `status`, `default_prompt`, `token_counts`
from `sessions` where `name` = '{}' and `last_interact_timestamp` > {} order by `last_interact_timestamp` asc from `sessions` where `name` = '{}' and `last_interact_timestamp` > {} order by `last_interact_timestamp` asc
limit 1 limit 1
""".format(session_name, cursor_timestamp)) """.format(session_name, cursor_timestamp))
@@ -223,6 +234,7 @@ class DatabaseManager:
prompt = result[5] prompt = result[5]
status = result[6] status = result[6]
default_prompt = result[7] default_prompt = result[7]
token_counts = result[8]
return { return {
'subject_type': subject_type, 'subject_type': subject_type,
@@ -230,13 +242,14 @@ class DatabaseManager:
'create_timestamp': create_timestamp, 'create_timestamp': create_timestamp,
'last_interact_timestamp': last_interact_timestamp, 'last_interact_timestamp': last_interact_timestamp,
'prompt': prompt, 'prompt': prompt,
'default_prompt': default_prompt 'default_prompt': default_prompt,
'token_counts': json.loads(token_counts)
} }
# 列出与某个对象的所有对话session # 列出与某个对象的所有对话session
def list_history(self, session_name: str, capacity: int, page: int): def list_history(self, session_name: str, capacity: int, page: int):
self.__execute__(""" self.__execute__("""
select `name`, `type`, `number`, `create_timestamp`, `last_interact_timestamp`, `prompt`, `status`, `default_prompt` select `name`, `type`, `number`, `create_timestamp`, `last_interact_timestamp`, `prompt`, `status`, `default_prompt`, `token_counts`
from `sessions` where `name` = '{}' order by `last_interact_timestamp` desc limit {} offset {} from `sessions` where `name` = '{}' order by `last_interact_timestamp` desc limit {} offset {}
""".format(session_name, capacity, capacity * page)) """.format(session_name, capacity, capacity * page))
results = self.cursor.fetchall() results = self.cursor.fetchall()
@@ -250,6 +263,7 @@ class DatabaseManager:
prompt = result[5] prompt = result[5]
status = result[6] status = result[6]
default_prompt = result[7] default_prompt = result[7]
token_counts = result[8]
sessions.append({ sessions.append({
'subject_type': subject_type, 'subject_type': subject_type,
@@ -257,7 +271,8 @@ class DatabaseManager:
'create_timestamp': create_timestamp, 'create_timestamp': create_timestamp,
'last_interact_timestamp': last_interact_timestamp, 'last_interact_timestamp': last_interact_timestamp,
'prompt': prompt, 'prompt': prompt,
'default_prompt': default_prompt 'default_prompt': default_prompt,
'token_counts': json.loads(token_counts)
}) })
return sessions return sessions