Skip to content

Commit

Permalink
Merge branch 'main' into new_tt_dev
Browse files Browse the repository at this point in the history
  • Loading branch information
csunny committed Nov 28, 2023
2 parents bca701e + d9cc102 commit 04e780c
Show file tree
Hide file tree
Showing 15 changed files with 262 additions and 83 deletions.
4 changes: 4 additions & 0 deletions assets/schema/knowledge_management.sql
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ CREATE TABLE `connect_config` (
`db_user` varchar(255) DEFAULT NULL COMMENT 'db user',
`db_pwd` varchar(255) DEFAULT NULL COMMENT 'db password',
`comment` text COMMENT 'db comment',
`sys_code` varchar(128) DEFAULT NULL COMMENT 'System code',
PRIMARY KEY (`id`),
UNIQUE KEY `uk_db` (`db_name`),
KEY `idx_q_db_type` (`db_type`)
Expand All @@ -78,6 +79,7 @@ CREATE TABLE `chat_history` (
`summary` varchar(255) COLLATE utf8mb4_unicode_ci NOT NULL COMMENT 'Conversation record summary',
`user_name` varchar(255) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'interlocutor',
`messages` text COLLATE utf8mb4_unicode_ci COMMENT 'Conversation details',
`sys_code` varchar(128) DEFAULT NULL COMMENT 'System code',
PRIMARY KEY (`id`)
) ENGINE=InnoDB AUTO_INCREMENT=2 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT 'Chat history';

Expand Down Expand Up @@ -110,6 +112,7 @@ CREATE TABLE `my_plugin` (
`version` varchar(255) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'plugin version',
`use_count` int DEFAULT NULL COMMENT 'plugin total use count',
`succ_count` int DEFAULT NULL COMMENT 'plugin total success count',
`sys_code` varchar(128) DEFAULT NULL COMMENT 'System code',
`gmt_created` TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT 'plugin install time',
PRIMARY KEY (`id`),
UNIQUE KEY `name` (`name`)
Expand Down Expand Up @@ -141,6 +144,7 @@ CREATE TABLE `prompt_manage` (
`prompt_name` varchar(512) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'prompt name',
`content` longtext COLLATE utf8mb4_unicode_ci COMMENT 'Prompt content',
`user_name` varchar(128) COLLATE utf8mb4_unicode_ci DEFAULT NULL COMMENT 'User name',
`sys_code` varchar(128) DEFAULT NULL COMMENT 'System code',
`gmt_created` timestamp NULL DEFAULT CURRENT_TIMESTAMP COMMENT 'created time',
`gmt_modified` timestamp NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT 'update time',
PRIMARY KEY (`id`),
Expand Down
6 changes: 6 additions & 0 deletions pilot/base_modules/agent/db/my_plugin_db.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ class MyPluginEntity(Base):
succ_count = Column(
Integer, nullable=True, default=0, comment="plugin total success count"
)
sys_code = Column(String(128), index=True, nullable=True, comment="System code")
gmt_created = Column(
DateTime, default=datetime.utcnow, comment="plugin install time"
)
Expand All @@ -58,6 +59,7 @@ def add(self, engity: MyPluginEntity):
version=engity.version,
use_count=engity.use_count or 0,
succ_count=engity.succ_count or 0,
sys_code=engity.sys_code,
gmt_created=datetime.now(),
)
session.add(my_plugin)
Expand Down Expand Up @@ -107,6 +109,8 @@ def list(self, query: MyPluginEntity, page=1, page_size=20) -> list[MyPluginEnti
my_plugins = my_plugins.filter(MyPluginEntity.user_code == query.user_code)
if query.user_name is not None:
my_plugins = my_plugins.filter(MyPluginEntity.user_name == query.user_name)
if query.sys_code is not None:
my_plugins = my_plugins.filter(MyPluginEntity.sys_code == query.sys_code)

my_plugins = my_plugins.order_by(MyPluginEntity.id.desc())
my_plugins = my_plugins.offset((page - 1) * page_size).limit(page_size)
Expand All @@ -133,6 +137,8 @@ def count(self, query: MyPluginEntity):
my_plugins = my_plugins.filter(MyPluginEntity.user_code == query.user_code)
if query.user_name is not None:
my_plugins = my_plugins.filter(MyPluginEntity.user_name == query.user_name)
if query.sys_code is not None:
my_plugins = my_plugins.filter(MyPluginEntity.sys_code == query.sys_code)
count = my_plugins.scalar()
session.close()
return count
Expand Down
4 changes: 4 additions & 0 deletions pilot/configs/model_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,10 @@ def get_device() -> str:
"xwin-lm-70b-v0.1": os.path.join(MODEL_PATH, "Xwin-LM-70B-V0.1"),
# https://huggingface.co/01-ai/Yi-34B-Chat
"yi-34b-chat": os.path.join(MODEL_PATH, "Yi-34B-Chat"),
# https://huggingface.co/01-ai/Yi-34B-Chat-8bits
"yi-34b-chat-8bits": os.path.join(MODEL_PATH, "Yi-34B-Chat-8bits"),
# https://huggingface.co/01-ai/Yi-34B-Chat-4bits
"yi-34b-chat-4bits": os.path.join(MODEL_PATH, "Yi-34B-Chat-4bits"),
"yi-6b-chat": os.path.join(MODEL_PATH, "Yi-6B-Chat"),
}

Expand Down
1 change: 1 addition & 0 deletions pilot/connections/manages/connect_config_db.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ class ConnectConfigEntity(Base):
db_user = Column(String(255), nullable=True, comment="db user")
db_pwd = Column(String(255), nullable=True, comment="db password")
comment = Column(Text, nullable=True, comment="db comment")
sys_code = Column(String(128), index=True, nullable=True, comment="System code")

__table_args__ = (
UniqueConstraint("db_name", name="uk_db"),
Expand Down
15 changes: 6 additions & 9 deletions pilot/memory/chat_history/base.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from __future__ import annotations

from abc import ABC, abstractmethod
from typing import List
from typing import List, Optional, Dict
from enum import Enum
from pilot.scene.message import OnceConversation

Expand Down Expand Up @@ -35,11 +35,6 @@ def append(self, message: OnceConversation) -> None:
# def clear(self) -> None:
# """Clear session memory from the local file"""

@abstractmethod
def conv_list(self, user_name: str = None) -> None:
"""get user's conversation list"""
pass

@abstractmethod
def update(self, messages: List[OnceConversation]) -> None:
pass
Expand All @@ -49,13 +44,15 @@ def delete(self) -> bool:
pass

@abstractmethod
def conv_info(self, conv_uid: str = None) -> None:
def conv_info(self, conv_uid: Optional[str] = None) -> None:
pass

@abstractmethod
def get_messages(self) -> List[OnceConversation]:
pass

@staticmethod
def conv_list(cls, user_name: str = None) -> None:
pass
def conv_list(
user_name: Optional[str] = None, sys_code: Optional[str] = None
) -> List[Dict]:
"""get user's conversation list"""
3 changes: 2 additions & 1 deletion pilot/memory/chat_history/chat_hisotry_factory.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from typing import Type
from .base import MemoryStoreType
from pilot.configs.config import Config
from pilot.memory.chat_history.base import BaseChatHistoryMemory
Expand Down Expand Up @@ -32,5 +33,5 @@ def get_store_instance(self, chat_session_id: str) -> BaseChatHistoryMemory:
chat_session_id
)

def get_store_cls(self):
def get_store_cls(self) -> Type[BaseChatHistoryMemory]:
return self.mem_store_class_map.get(CFG.CHAT_HISTORY_STORE_TYPE)
10 changes: 7 additions & 3 deletions pilot/memory/chat_history/chat_history_db.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import List
from typing import List, Optional
from sqlalchemy import Column, Integer, String, Index, DateTime, func, Boolean, Text
from sqlalchemy import UniqueConstraint

Expand Down Expand Up @@ -32,7 +32,7 @@ class ChatHistoryEntity(Base):
messages = Column(
Text(length=2**31 - 1), nullable=True, comment="Conversation details"
)

sys_code = Column(String(128), index=True, nullable=True, comment="System code")
UniqueConstraint("conv_uid", name="uk_conversation")
Index("idx_q_user", "user_name")
Index("idx_q_mode", "chat_mode")
Expand All @@ -48,11 +48,15 @@ def __init__(self):
session=session,
)

def list_last_20(self, user_name: str = None):
def list_last_20(
self, user_name: Optional[str] = None, sys_code: Optional[str] = None
):
session = self.get_session()
chat_history = session.query(ChatHistoryEntity)
if user_name:
chat_history = chat_history.filter(ChatHistoryEntity.user_name == user_name)
if sys_code:
chat_history = chat_history.filter(ChatHistoryEntity.sys_code == sys_code)

chat_history = chat_history.order_by(ChatHistoryEntity.id.desc())

Expand Down
38 changes: 24 additions & 14 deletions pilot/memory/chat_history/store_type/duckdb_history.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import json
import os
import duckdb
from typing import List
from typing import List, Dict, Optional

from pilot.configs.config import Config
from pilot.memory.chat_history.base import BaseChatHistoryMemory
Expand Down Expand Up @@ -37,7 +37,7 @@ def __init_chat_history_tables(self):
if not result:
# 如果表不存在,则创建新表
self.connect.execute(
"CREATE TABLE chat_history (id integer primary key, conv_uid VARCHAR(100) UNIQUE, chat_mode VARCHAR(50), summary VARCHAR(255), user_name VARCHAR(100), messages TEXT)"
"CREATE TABLE chat_history (id integer primary key, conv_uid VARCHAR(100) UNIQUE, chat_mode VARCHAR(50), summary VARCHAR(255), user_name VARCHAR(100), sys_code VARCHAR(128), messages TEXT)"
)
self.connect.execute("CREATE SEQUENCE seq_id START 1;")

Expand All @@ -61,8 +61,8 @@ def create(self, chat_mode, summary: str, user_name: str) -> None:
try:
cursor = self.connect.cursor()
cursor.execute(
"INSERT INTO chat_history(id, conv_uid, chat_mode summary, user_name, messages)VALUES(nextval('seq_id'),?,?,?,?,?)",
[self.chat_seesion_id, chat_mode, summary, user_name, ""],
"INSERT INTO chat_history(id, conv_uid, chat_mode summary, user_name, sys_code, messages)VALUES(nextval('seq_id'),?,?,?,?,?,?)",
[self.chat_seesion_id, chat_mode, summary, user_name, "", ""],
)
cursor.commit()
self.connect.commit()
Expand All @@ -83,12 +83,13 @@ def append(self, once_message: OnceConversation) -> None:
)
else:
cursor.execute(
"INSERT INTO chat_history(id, conv_uid, chat_mode, summary, user_name, messages)VALUES(nextval('seq_id'),?,?,?,?,?)",
"INSERT INTO chat_history(id, conv_uid, chat_mode, summary, user_name, sys_code, messages)VALUES(nextval('seq_id'),?,?,?,?,?,?)",
[
self.chat_seesion_id,
once_message.chat_mode,
once_message.get_user_conv().content,
"",
once_message.user_name,
once_message.sys_code,
json.dumps(conversations, ensure_ascii=False),
],
)
Expand Down Expand Up @@ -149,17 +150,26 @@ def get_messages(self) -> List[OnceConversation]:
return None

@staticmethod
def conv_list(cls, user_name: str = None) -> None:
def conv_list(
user_name: Optional[str] = None, sys_code: Optional[str] = None
) -> List[Dict]:
if os.path.isfile(duckdb_path):
cursor = duckdb.connect(duckdb_path).cursor()
query = "SELECT * FROM chat_history"
params = []
conditions = []
if user_name:
cursor.execute(
"SELECT * FROM chat_history where user_name=? order by id desc limit 20",
[user_name],
)
else:
cursor.execute("SELECT * FROM chat_history order by id desc limit 20")
# 获取查询结果字段名
conditions.append("user_name = ?")
params.append(user_name)
if sys_code:
conditions.append("sys_code = ?")
params.append(sys_code)

if conditions:
query += " WHERE " + " AND ".join(conditions)

query += " ORDER BY id DESC LIMIT 20"
cursor.execute(query, params)
fields = [field[0] for field in cursor.description]
data = []
for row in cursor.fetchall():
Expand Down
11 changes: 7 additions & 4 deletions pilot/memory/chat_history/store_type/meta_db_history.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import json
import logging
from typing import List
from typing import List, Dict, Optional
from sqlalchemy import Column, Integer, String, Index, DateTime, func, Boolean, Text
from sqlalchemy import UniqueConstraint
from pilot.configs.config import Config
Expand Down Expand Up @@ -62,7 +62,8 @@ def append(self, once_message: OnceConversation) -> None:
chat_history: ChatHistoryEntity = ChatHistoryEntity()
chat_history.conv_uid = self.chat_seesion_id
chat_history.chat_mode = once_message.chat_mode
chat_history.user_name = "default"
chat_history.user_name = once_message.user_name
chat_history.sys_code = once_message.sys_code
chat_history.summary = once_message.get_user_conv().content

conversations.append(_conversation_to_dic(once_message))
Expand Down Expand Up @@ -92,9 +93,11 @@ def get_messages(self) -> List[OnceConversation]:
return []

@staticmethod
def conv_list(cls, user_name: str = None) -> None:
def conv_list(
user_name: Optional[str] = None, sys_code: Optional[str] = None
) -> List[Dict]:
chat_history_dao = ChatHistoryDao()
history_list = chat_history_dao.list_last_20()
history_list = chat_history_dao.list_last_20(user_name, sys_code)
result = []
for history in history_list:
result.append(history.__dict__)
Expand Down
Loading

0 comments on commit 04e780c

Please sign in to comment.