♻️ refactor: refactor the db model (#1567)

* ♻️ refactor: refactor the db code

*  test: fix test

* ♻️ refactor: refactor the  user model
This commit is contained in:
Arvin Xu
2024-03-14 20:20:09 +08:00
committed by GitHub
parent 2cf2abd029
commit 3d56dd6a1a
13 changed files with 388 additions and 319 deletions

1
.gitignore vendored
View File

@@ -59,3 +59,4 @@ bun.lockb
sitemap*.xml
robots.txt
*.patch

View File

@@ -148,7 +148,7 @@
"uuid": "^9",
"yaml": "^2",
"zod": "^3",
"zustand": "^4.4",
"zustand": "^4.5.2",
"zustand-utils": "^1.3.2"
},
"devDependencies": {

View File

@@ -11,7 +11,7 @@ export const GET = async (req: Request) => {
let res: Response;
res = await fetch(pluginStore.getPluginIndexUrl(locale as any));
res = await fetch(pluginStore.getPluginIndexUrl(locale as any), { next: { revalidate: 3600 } });
if (res.status === 404) {
res = await fetch(pluginStore.getPluginIndexUrl(DEFAULT_LANG));

View File

@@ -26,19 +26,8 @@ class _MessageModel extends BaseModel {
constructor() {
super('messages', DB_MessageSchema);
}
async create(data: CreateMessageParams) {
const id = nanoid();
const messageData: DB_Message = this.mapChatMessageToDBMessage(data as ChatMessage);
return this._add(messageData, id);
}
async batchCreate(messages: ChatMessage[]) {
const data: DB_Message[] = messages.map((m) => this.mapChatMessageToDBMessage(m));
return this._batchAdd(data);
}
// **************** Query *************** //
async query({
sessionId,
@@ -91,6 +80,73 @@ class _MessageModel extends BaseModel {
return this.table.get(id);
}
async queryAll() {
const data: DBModel<DB_Message>[] = await this.table.orderBy('updatedAt').toArray();
return data.map((element) => this.mapToChatMessage(element));
}
async queryBySessionId(sessionId: string) {
return this.table.where('sessionId').equals(sessionId).toArray();
}
queryByTopicId = async (topicId: string) => {
const dbMessages = await this.table.where('topicId').equals(topicId).toArray();
return dbMessages.map((message) => this.mapToChatMessage(message));
};
async count() {
return this.table.count();
}
// **************** Create *************** //
async create(data: CreateMessageParams) {
const id = nanoid();
const messageData: DB_Message = this.mapChatMessageToDBMessage(data as ChatMessage);
return this._add(messageData, id);
}
async batchCreate(messages: ChatMessage[]) {
const data: DB_Message[] = messages.map((m) => this.mapChatMessageToDBMessage(m));
return this._batchAdd(data);
}
async duplicateMessages(messages: ChatMessage[]): Promise<ChatMessage[]> {
const duplicatedMessages = await this.createDuplicateMessages(messages);
// 批量添加复制后的消息到数据库
await this.batchCreate(duplicatedMessages);
return duplicatedMessages;
}
async createDuplicateMessages(messages: ChatMessage[]): Promise<ChatMessage[]> {
// 创建一个映射来存储原始消息ID和复制消息ID之间的关系
const idMapping = new Map<string, string>();
// 首先复制所有消息并为每个复制的消息生成新的ID
const duplicatedMessages = messages.map((originalMessage) => {
const newId = nanoid();
idMapping.set(originalMessage.id, newId);
return { ...originalMessage, id: newId };
});
// 更新 parentId 为复制后的新ID
for (const duplicatedMessage of duplicatedMessages) {
if (duplicatedMessage.parentId && idMapping.has(duplicatedMessage.parentId)) {
duplicatedMessage.parentId = idMapping.get(duplicatedMessage.parentId);
}
}
return duplicatedMessages;
}
// **************** Delete *************** //
async delete(id: string) {
return this.table.delete(id);
}
@@ -99,8 +155,36 @@ class _MessageModel extends BaseModel {
return this.table.clear();
}
/**
* Deletes multiple messages based on the assistantId and optionally the topicId.
* If topicId is not provided, it deletes messages where topicId is undefined or null.
* If topicId is provided, it deletes messages with that specific topicId.
*
* @param {string} sessionId - The identifier of the assistant associated with the messages.
* @param {string | undefined} topicId - The identifier of the topic associated with the messages (optional).
* @returns {Promise<void>}
*/
async batchDelete(sessionId: string, topicId: string | undefined): Promise<void> {
// If topicId is specified, use both assistantId and topicId as the filter criteria in the query.
// Otherwise, filter by assistantId and require that topicId is undefined.
const query = !!topicId
? this.table.where({ sessionId, topicId }) // Use a compound index
: this.table
.where('sessionId')
.equals(sessionId)
.and((message) => !message.topicId);
// Retrieve a collection of message IDs that satisfy the criteria
const messageIds = await query.primaryKeys();
// Use the bulkDelete method to delete all selected messages in bulk
return this.table.bulkDelete(messageIds);
}
// **************** Update *************** //
async update(id: string, data: DeepPartial<DB_Message>) {
return super._update(id, data);
return this._update(id, data);
}
async updatePluginState(id: string, key: string, value: any) {
@@ -132,80 +216,7 @@ class _MessageModel extends BaseModel {
return updatedMessages.length;
}
/**
* Deletes multiple messages based on the assistantId and optionally the topicId.
* If topicId is not provided, it deletes messages where topicId is undefined or null.
* If topicId is provided, it deletes messages with that specific topicId.
*
* @param {string} sessionId - The identifier of the assistant associated with the messages.
* @param {string | undefined} topicId - The identifier of the topic associated with the messages (optional).
* @returns {Promise<void>}
*/
async batchDelete(sessionId: string, topicId: string | undefined): Promise<void> {
// If topicId is specified, use both assistantId and topicId as the filter criteria in the query.
// Otherwise, filter by assistantId and require that topicId is undefined.
const query = !!topicId
? this.table.where({ sessionId, topicId }) // Use a compound index
: this.table
.where('sessionId')
.equals(sessionId)
.and((message) => !message.topicId);
// Retrieve a collection of message IDs that satisfy the criteria
const messageIds = await query.primaryKeys();
// Use the bulkDelete method to delete all selected messages in bulk
return this.table.bulkDelete(messageIds);
}
async queryAll() {
const data: DBModel<DB_Message>[] = await this.table.orderBy('updatedAt').toArray();
return data.map((element) => this.mapToChatMessage(element));
}
async count() {
return this.table.count();
}
async queryBySessionId(sessionId: string) {
return this.table.where('sessionId').equals(sessionId).toArray();
}
queryByTopicId = async (topicId: string) => {
const dbMessages = await this.table.where('topicId').equals(topicId).toArray();
return dbMessages.map((message) => this.mapToChatMessage(message));
};
async duplicateMessages(messages: ChatMessage[]): Promise<ChatMessage[]> {
const duplicatedMessages = await this.createDuplicateMessages(messages);
// 批量添加复制后的消息到数据库
await this.batchCreate(duplicatedMessages);
return duplicatedMessages;
}
async createDuplicateMessages(messages: ChatMessage[]): Promise<ChatMessage[]> {
// 创建一个映射来存储原始消息ID和复制消息ID之间的关系
const idMapping = new Map<string, string>();
// 首先复制所有消息并为每个复制的消息生成新的ID
const duplicatedMessages = messages.map((originalMessage) => {
const newId = nanoid();
idMapping.set(originalMessage.id, newId);
return { ...originalMessage, id: newId };
});
// 更新 parentId 为复制后的新ID
for (const duplicatedMessage of duplicatedMessages) {
if (duplicatedMessage.parentId && idMapping.has(duplicatedMessage.parentId)) {
duplicatedMessage.parentId = idMapping.get(duplicatedMessage.parentId);
}
}
return duplicatedMessages;
}
// **************** Helper *************** //
private mapChatMessageToDBMessage(message: ChatMessage): DB_Message {
const { extra, ...messageData } = message;

View File

@@ -7,11 +7,14 @@ class _PluginModel extends BaseModel {
constructor() {
super('plugins', DB_PluginSchema);
}
// **************** Query *************** //
getList = async (): Promise<DB_Plugin[]> => {
return this.table.toArray();
};
// **************** Create *************** //
create = async (plugin: DB_Plugin) => {
const old = await this.table.get(plugin.identifier);
@@ -21,18 +24,20 @@ class _PluginModel extends BaseModel {
batchCreate = async (plugins: DB_Plugin[]) => {
return this._batchAdd(plugins);
};
// **************** Delete *************** //
delete(id: string) {
return this.table.delete(id);
}
clear() {
return this.table.clear();
}
// **************** Update *************** //
update: (id: string, value: Partial<DB_Plugin>) => Promise<number> = async (id, value) => {
return this.table.update(id, value);
};
clear() {
return this.table.clear();
}
}
export const PluginModel = new _PluginModel();

View File

@@ -21,29 +21,7 @@ class _SessionModel extends BaseModel {
super('sessions', DB_SessionSchema);
}
async create(type: 'agent' | 'group', defaultValue: Partial<LobeAgentSession>, id = uuid()) {
const data = merge(DEFAULT_AGENT_LOBE_SESSION, { type, ...defaultValue });
const dataDB = this.mapToDB_Session(data);
return this._add(dataDB, id);
}
async batchCreate(sessions: LobeAgentSession[]) {
const DB_Sessions = await Promise.all(
sessions.map(async (s) => {
if (s.group && s.group !== SessionDefaultGroup.Default) {
// Check if the group exists in the SessionGroup table
const groupExists = await SessionGroupModel.findById(s.group);
// If the group does not exist, set it to default
if (!groupExists) {
s.group = SessionDefaultGroup.Default;
}
}
return this.mapToDB_Session(s);
}),
);
return this._batchAdd<DB_Session>(DB_Sessions, { idGenerator: uuid });
}
// **************** Query *************** //
async query({
pageSize = 9999,
@@ -103,59 +81,6 @@ class _SessionModel extends BaseModel {
return Object.fromEntries(groupItems);
}
async update(id: string, data: Partial<DB_Session>) {
return super._update(id, data);
}
async updatePinned(id: string, pinned: boolean) {
return this.update(id, { pinned: pinned ? 1 : 0 });
}
async updateConfig(id: string, data: DeepPartial<LobeAgentConfig>) {
const session = await this.findById(id);
if (!session) return;
const config = merge(session.config, data);
return this.update(id, { config });
}
/**
* Delete a session , also delete all messages and topic associated with it.
*/
async delete(id: string) {
return this.db.transaction('rw', [this.table, this.db.topics, this.db.messages], async () => {
// Delete all topics associated with the session
const topics = await this.db.topics.where('sessionId').equals(id).toArray();
const topicIds = topics.map((topic) => topic.id);
if (topicIds.length > 0) {
await this.db.topics.bulkDelete(topicIds);
}
// Delete all messages associated with the session
const messages = await this.db.messages.where('sessionId').equals(id).toArray();
const messageIds = messages.map((message) => message.id);
if (messageIds.length > 0) {
await this.db.messages.bulkDelete(messageIds);
}
// Finally, delete the session itself
await this.table.delete(id);
});
}
async clearTable() {
return this.table.clear();
}
async findById(id: string): Promise<DBModel<DB_Session>> {
return this.table.get(id);
}
async isEmpty() {
return (await this.table.count()) === 0;
}
/**
* Query sessions by keyword in title, description, content, or translated content
* @param keyword The keyword to search for
@@ -225,15 +150,6 @@ class _SessionModel extends BaseModel {
return this.mapToAgentSessions(items);
}
async duplicate(id: string, newTitle?: string) {
const session = await this.findById(id);
if (!session) return;
const newSession = merge(session, { meta: { title: newTitle } });
return this._add(newSession, uuid());
}
async getPinnedSessions(): Promise<LobeSessions> {
const items: DBModel<DB_Session>[] = await this.table
.where('pinned')
@@ -244,6 +160,100 @@ class _SessionModel extends BaseModel {
return this.mapToAgentSessions(items);
}
async findById(id: string): Promise<DBModel<DB_Session>> {
return this.table.get(id);
}
async isEmpty() {
return (await this.table.count()) === 0;
}
// **************** Create *************** //
async create(type: 'agent' | 'group', defaultValue: Partial<LobeAgentSession>, id = uuid()) {
const data = merge(DEFAULT_AGENT_LOBE_SESSION, { type, ...defaultValue });
const dataDB = this.mapToDB_Session(data);
return this._add(dataDB, id);
}
async batchCreate(sessions: LobeAgentSession[]) {
const DB_Sessions = await Promise.all(
sessions.map(async (s) => {
if (s.group && s.group !== SessionDefaultGroup.Default) {
// Check if the group exists in the SessionGroup table
const groupExists = await SessionGroupModel.findById(s.group);
// If the group does not exist, set it to default
if (!groupExists) {
s.group = SessionDefaultGroup.Default;
}
}
return this.mapToDB_Session(s);
}),
);
return this._batchAdd<DB_Session>(DB_Sessions, { idGenerator: uuid });
}
async duplicate(id: string, newTitle?: string) {
const session = await this.findById(id);
if (!session) return;
const newSession = merge(session, { meta: { title: newTitle } });
return this._add(newSession, uuid());
}
// **************** Delete *************** //
/**
* Delete a session , also delete all messages and topic associated with it.
*/
async delete(id: string) {
return this.db.transaction('rw', [this.table, this.db.topics, this.db.messages], async () => {
// Delete all topics associated with the session
const topics = await this.db.topics.where('sessionId').equals(id).toArray();
const topicIds = topics.map((topic) => topic.id);
if (topicIds.length > 0) {
await this.db.topics.bulkDelete(topicIds);
}
// Delete all messages associated with the session
const messages = await this.db.messages.where('sessionId').equals(id).toArray();
const messageIds = messages.map((message) => message.id);
if (messageIds.length > 0) {
await this.db.messages.bulkDelete(messageIds);
}
// Finally, delete the session itself
await this.table.delete(id);
});
}
async clearTable() {
return this.table.clear();
}
// **************** Update *************** //
async update(id: string, data: Partial<DB_Session>) {
return this._update(id, data);
}
async updatePinned(id: string, pinned: boolean) {
return this.update(id, { pinned: pinned ? 1 : 0 });
}
async updateConfig(id: string, data: DeepPartial<LobeAgentConfig>) {
const session = await this.findById(id);
if (!session) return;
const config = merge(session.config, data);
return this.update(id, { config });
}
// **************** Helper *************** //
private mapToDB_Session(session: LobeAgentSession): DBModel<DB_Session> {
return {
...session,

View File

@@ -8,32 +8,7 @@ class _SessionGroupModel extends BaseModel {
super('sessionGroups', DB_SessionGroupSchema);
}
async create(name: string, sort?: number, id = nanoid()) {
return this._add({ name, sort }, id);
}
async batchCreate(groups: SessionGroups) {
return this._batchAdd(groups, { idGenerator: nanoid });
}
async findById(id: string): Promise<DB_SessionGroup> {
return this.table.get(id);
}
async update(id: string, data: Partial<DB_SessionGroup>) {
return super._update(id, data);
}
async delete(id: string, removeGroupItem: boolean = false) {
this.db.sessions.toCollection().modify((session) => {
// update all session associated with the sessionGroup to default
if (session.group === id) session.group = 'default';
});
if (!removeGroupItem) {
return this.table.delete(id);
} else {
return this.db.sessions.where('group').equals(id).delete();
}
}
// **************** Query *************** //
async query(): Promise<SessionGroups> {
const allGroups = await this.table.toArray();
@@ -60,6 +35,43 @@ class _SessionGroupModel extends BaseModel {
});
}
async findById(id: string): Promise<DB_SessionGroup> {
return this.table.get(id);
}
// **************** Create *************** //
async create(name: string, sort?: number, id = nanoid()) {
return this._add({ name, sort }, id);
}
async batchCreate(groups: SessionGroups) {
return this._batchAdd(groups, { idGenerator: nanoid });
}
// **************** Delete *************** //
async delete(id: string, removeGroupItem: boolean = false) {
this.db.sessions.toCollection().modify((session) => {
// update all session associated with the sessionGroup to default
if (session.group === id) session.group = 'default';
});
if (!removeGroupItem) {
return this.table.delete(id);
} else {
return this.db.sessions.where('group').equals(id).delete();
}
}
async clear() {
this.table.clear();
}
// **************** Update *************** //
async update(id: string, data: Partial<DB_SessionGroup>) {
return super._update(id, data);
}
async updateOrder(sortMap: { id: string; sort: number }[]) {
return this.db.transaction('rw', this.table, async () => {
for (const { id, sort } of sortMap) {
@@ -67,10 +79,6 @@ class _SessionGroupModel extends BaseModel {
}
});
}
async clear() {
this.table.clear();
}
}
export const SessionGroupModel = new _SessionGroupModel();

View File

@@ -23,19 +23,7 @@ class _TopicModel extends BaseModel {
super('topics', DB_TopicSchema);
}
async create({ title, favorite, sessionId, messages }: CreateTopicParams, id = nanoid()) {
const topic = await this._add({ favorite: favorite ? 1 : 0, sessionId, title: title }, id);
// add topicId to these messages
if (messages) {
await this.db.messages.where('id').anyOf(messages).modify({ topicId: topic.id });
}
return topic;
}
async batchCreate(topics: CreateTopicParams[]) {
return this._batchAdd(topics.map((t) => ({ ...t, favorite: t.favorite ? 1 : 0 })));
}
// **************** Query *************** //
async query({ pageSize = 9999, current = 0, sessionId }: QueryTopicParams): Promise<ChatTopic[]> {
const offset = current * pageSize;
@@ -58,90 +46,6 @@ class _TopicModel extends BaseModel {
return pagedTopics.map((i) => this.mapToChatTopic(i));
}
async findBySessionId(sessionId: string) {
return this.table.where({ sessionId }).toArray();
}
async findById(id: string): Promise<DBModel<DB_Topic>> {
return this.table.get(id);
}
/**
* Deletes a topic and all messages associated with it.
*/
async delete(id: string) {
return this.db.transaction('rw', [this.table, this.db.messages], async () => {
// Delete all messages associated with the topic
const messages = await this.db.messages.where('topicId').equals(id).toArray();
if (messages.length > 0) {
const messageIds = messages.map((msg) => msg.id);
await this.db.messages.bulkDelete(messageIds);
}
await this.table.delete(id);
});
}
/**
* Deletes multiple topic based on the sessionId.
*
* @param {string} sessionId - The identifier of the assistant associated with the messages.
* @returns {Promise<void>}
*/
async batchDeleteBySessionId(sessionId: string): Promise<void> {
// use sessionId as the filter criteria in the query.
const query = this.table.where('sessionId').equals(sessionId);
// Retrieve a collection of message IDs that satisfy the criteria
const topicIds = await query.primaryKeys();
// Use the bulkDelete method to delete all selected messages in bulk
return this.table.bulkDelete(topicIds);
}
async clearTable() {
return this.table.clear();
}
async update(id: string, data: Partial<DB_Topic>) {
return super._update(id, { ...data, updatedAt: Date.now() });
}
async toggleFavorite(id: string, newState?: boolean) {
const topic = await this.findById(id);
if (!topic) {
throw new Error(`Topic with id ${id} not found`);
}
// Toggle the 'favorite' status
const nextState = typeof newState !== 'undefined' ? newState : !topic.favorite;
await this.update(id, { favorite: nextState ? 1 : 0 });
return nextState;
}
/**
* Deletes multiple topics and all messages associated with them in a transaction.
*/
async batchDelete(topicIds: string[]) {
return this.db.transaction('rw', [this.table, this.db.messages], async () => {
// Iterate over each topicId and delete related messages, then delete the topic itself
for (const topicId of topicIds) {
// Delete all messages associated with the topic
const messages = await this.db.messages.where('topicId').equals(topicId).toArray();
if (messages.length > 0) {
const messageIds = messages.map((msg) => msg.id);
await this.db.messages.bulkDelete(messageIds);
}
// Delete the topic
await this.table.delete(topicId);
}
});
}
queryAll() {
return this.table.orderBy('updatedAt').toArray();
}
@@ -201,6 +105,29 @@ class _TopicModel extends BaseModel {
return uniqueTopics.map((i) => ({ ...i, favorite: !!i.favorite }));
}
async findBySessionId(sessionId: string) {
return this.table.where({ sessionId }).toArray();
}
async findById(id: string): Promise<DBModel<DB_Topic>> {
return this.table.get(id);
}
// **************** Create *************** //
async create({ title, favorite, sessionId, messages }: CreateTopicParams, id = nanoid()) {
const topic = await this._add({ favorite: favorite ? 1 : 0, sessionId, title: title }, id);
// add topicId to these messages
if (messages) {
await this.db.messages.where('id').anyOf(messages).modify({ topicId: topic.id });
}
return topic;
}
async batchCreate(topics: CreateTopicParams[]) {
return this._batchAdd(topics.map((t) => ({ ...t, favorite: t.favorite ? 1 : 0 })));
}
async duplicateTopic(topicId: string, newTitle?: string) {
return this.db.transaction('rw', this.db.topics, this.db.messages, async () => {
// Step 1: get DB_Topic
@@ -226,6 +153,86 @@ class _TopicModel extends BaseModel {
});
}
// **************** Delete *************** //
/**
* Deletes a topic and all messages associated with it.
*/
async delete(id: string) {
return this.db.transaction('rw', [this.table, this.db.messages], async () => {
// Delete all messages associated with the topic
const messages = await this.db.messages.where('topicId').equals(id).toArray();
if (messages.length > 0) {
const messageIds = messages.map((msg) => msg.id);
await this.db.messages.bulkDelete(messageIds);
}
await this.table.delete(id);
});
}
/**
* Deletes multiple topic based on the sessionId.
*
* @param {string} sessionId - The identifier of the assistant associated with the messages.
* @returns {Promise<void>}
*/
async batchDeleteBySessionId(sessionId: string): Promise<void> {
// use sessionId as the filter criteria in the query.
const query = this.table.where('sessionId').equals(sessionId);
// Retrieve a collection of message IDs that satisfy the criteria
const topicIds = await query.primaryKeys();
// Use the bulkDelete method to delete all selected messages in bulk
return this.table.bulkDelete(topicIds);
}
/**
* Deletes multiple topics and all messages associated with them in a transaction.
*/
async batchDelete(topicIds: string[]) {
return this.db.transaction('rw', [this.table, this.db.messages], async () => {
// Iterate over each topicId and delete related messages, then delete the topic itself
for (const topicId of topicIds) {
// Delete all messages associated with the topic
const messages = await this.db.messages.where('topicId').equals(topicId).toArray();
if (messages.length > 0) {
const messageIds = messages.map((msg) => msg.id);
await this.db.messages.bulkDelete(messageIds);
}
// Delete the topic
await this.table.delete(topicId);
}
});
}
async clearTable() {
return this.table.clear();
}
// **************** Update *************** //
async update(id: string, data: Partial<DB_Topic>) {
return this._update(id, data);
}
async toggleFavorite(id: string, newState?: boolean) {
const topic = await this.findById(id);
if (!topic) {
throw new Error(`Topic with id ${id} not found`);
}
// Toggle the 'favorite' status
const nextState = typeof newState !== 'undefined' ? newState : !topic.favorite;
await this.update(id, { favorite: nextState ? 1 : 0 });
return nextState;
}
// **************** Helper *************** //
private mapToChatTopic = (dbTopic: DBModel<DB_Topic>): ChatTopic => ({
...dbTopic,
favorite: !!dbTopic.favorite,

View File

@@ -10,6 +10,7 @@ class _UserModel extends BaseModel {
constructor() {
super('users', DB_UserSchema);
}
// **************** Query *************** //
getUser = async (): Promise<DB_User & { id: number }> => {
const noUser = !(await this.table.count());
@@ -21,18 +22,20 @@ class _UserModel extends BaseModel {
return list[0];
};
// **************** Create *************** //
create = async (user: DB_User) => {
return this.table.put(user);
};
private update = async (id: number, value: DeepPartial<DB_User>) => {
return this.table.update(id, value);
};
// **************** Delete *************** //
clear() {
return this.table.clear();
}
// **************** Update *************** //
async updateSettings(settings: DeepPartial<GlobalSettings>) {
const user = await this.getUser();
@@ -50,6 +53,12 @@ class _UserModel extends BaseModel {
return this.update(user.id, { avatar });
}
// **************** Helper *************** //
private update = async (id: number, value: DeepPartial<DB_User>) => {
return this.table.update(id, value);
};
}
export const UserModel = new _UserModel();

18
src/libs/swr/index.ts Normal file
View File

@@ -0,0 +1,18 @@
import useSWR, { SWRHook } from 'swr';
/**
* 这一类请求方法是比较「死」的请求模式,只会在第一次请求时触发。不会自动刷新,刷新需要搭配 refreshXXX 这样的方法实现,
* 适用于 messages、topics、sessions 等由用户在客户端交互产生的数据。
*/
// @ts-ignore
export const useClientDataSWR: SWRHook = (key, fetch, config) =>
useSWR(key, fetch, {
// default is 2000ms ,it makes the user's quick switch don't work correctly.
// Cause issue like this: https://github.com/lobehub/lobe-chat/issues/532
// we need to set it to 0.
dedupingInterval: 0,
refreshWhenOffline: false,
revalidateOnFocus: false,
revalidateOnReconnect: false,
...config,
});

View File

@@ -559,7 +559,7 @@ describe('chatMessage actions', () => {
});
// 确保 mutate 调用了正确的参数
expect(mutate).toHaveBeenCalledWith([activeId, activeTopicId]);
expect(mutate).toHaveBeenCalledWith(['SWR_USE_FETCH_MESSAGES', activeId, activeTopicId]);
});
it('should handle errors during refreshing messages', async () => {
useChatStore.setState({ refreshMessages: realRefreshMessages });

View File

@@ -2,12 +2,13 @@
// Disable the auto sort key eslint rule to make the code more logic and readable
import { copyToClipboard } from '@lobehub/ui';
import { template } from 'lodash-es';
import useSWR, { SWRResponse, mutate } from 'swr';
import { SWRResponse, mutate } from 'swr';
import { StateCreator } from 'zustand/vanilla';
import { LOADING_FLAT, isFunctionMessageAtStart, testFunctionMessageAtEnd } from '@/const/message';
import { TraceEventType, TraceNameMap } from '@/const/trace';
import { CreateMessageParams } from '@/database/models/message';
import { useClientDataSWR } from '@/libs/swr';
import { chatService } from '@/services/chat';
import { messageService } from '@/services/message';
import { topicService } from '@/services/topic';
@@ -25,6 +26,8 @@ import { MessageDispatch, messagesReducer } from './reducer';
const n = setNamespace('message');
const SWR_USE_FETCH_MESSAGES = 'SWR_USE_FETCH_MESSAGES';
interface SendMessageParams {
message: string;
files?: { id: string; url: string }[];
@@ -274,9 +277,9 @@ export const chatMessage: StateCreator<
await get().internalUpdateMessageContent(id, content);
},
useFetchMessages: (sessionId, activeTopicId) =>
useSWR<ChatMessage[]>(
[sessionId, activeTopicId],
async ([sessionId, topicId]: [string, string | undefined]) =>
useClientDataSWR<ChatMessage[]>(
[SWR_USE_FETCH_MESSAGES, sessionId, activeTopicId],
async ([, sessionId, topicId]: [string, string, string | undefined]) =>
messageService.getMessages(sessionId, topicId),
{
onSuccess: (messages, key) => {
@@ -289,14 +292,10 @@ export const chatMessage: StateCreator<
}),
);
},
// default is 2000ms ,it makes the user's quick switch don't work correctly.
// Cause issue like this: https://github.com/lobehub/lobe-chat/issues/532
// we need to set it to 0.
dedupingInterval: 0,
},
),
refreshMessages: async () => {
await mutate([get().activeId, get().activeTopicId]);
await mutate([SWR_USE_FETCH_MESSAGES, get().activeId, get().activeTopicId]);
},
// the internal process method of the AI message

View File

@@ -6,6 +6,7 @@ import { StateCreator } from 'zustand/vanilla';
import { INBOX_SESSION_ID } from '@/const/session';
import { SESSION_CHAT_URL } from '@/const/url';
import { useClientDataSWR } from '@/libs/swr';
import { sessionService } from '@/services/session';
import { useGlobalStore } from '@/store/global';
import { settingsSelectors } from '@/store/global/selectors';
@@ -154,7 +155,7 @@ export const createSessionSlice: StateCreator<
},
useFetchSessions: () =>
useSWR<ChatSessionList>(FETCH_SESSIONS_KEY, sessionService.getSessionsWithGroup, {
useClientDataSWR<ChatSessionList>(FETCH_SESSIONS_KEY, sessionService.getSessionsWithGroup, {
onSuccess: (data) => {
// 由于 https://github.com/lobehub/lobe-chat/pull/541 的关系
// 只有触发了 refreshSessions 才会更新 sessions进而触发页面 rerender