Skip to content

Commit 2d396c9

Browse files
author
linyuan.yang
committed
saver
1 parent d238e47 commit 2d396c9

6 files changed

Lines changed: 215 additions & 82 deletions

File tree

‎packages/scorpio.ai/src/Saver/AgentFileSaver.ts‎

Lines changed: 54 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,13 @@ import { IAgentSaverService, ChatMessage, StoredMessage, ChatMessageOptions } fr
55
import { ILoggerService, ILogger } from "../Logger";
66
import { inject } from "scorpio.di";
77
import { T_DBPath } from "../Core/tokens";
8-
import { applyTokenLimit } from "./messageSerializer";
98

109
interface ThreadFile {
1110
messages: StoredMessage[];
1211
thinks?: Record<string, StoredMessage[]>;
1312
metadata?: Record<string, string>;
13+
nextId?: number;
14+
compactedIds?: number[];
1415
}
1516

1617
// ─────────────────────────────────────────────────────────────────────────────
@@ -35,14 +36,20 @@ export class AgentFileSaver implements IAgentSaverService {
3536
if (this.cache) return this.cache;
3637
try {
3738
if (!existsSync(this.filePath)) {
38-
this.cache = { messages: [], thinks: {} };
39+
this.cache = { messages: [], thinks: {}, nextId: 1, compactedIds: [] };
3940
} else {
4041
const content = await readFile(this.filePath, "utf-8");
4142
this.cache = JSON.parse(content) as ThreadFile;
43+
let nextId = this.cache.nextId ?? 1;
44+
for (const m of this.cache.messages) {
45+
if (m.id == null) m.id = nextId++;
46+
}
47+
this.cache.nextId = nextId;
48+
if (!this.cache.compactedIds) this.cache.compactedIds = [];
4249
}
4350
} catch (error: any) {
4451
this.logger?.warn(`读取文件失败: ${error.message}`);
45-
this.cache = { messages: [], thinks: {} };
52+
this.cache = { messages: [], thinks: {}, nextId: 1, compactedIds: [] };
4653
}
4754
return this.cache!;
4855
}
@@ -55,37 +62,67 @@ export class AgentFileSaver implements IAgentSaverService {
5562

5663
async pushMessage(message: ChatMessage, options?: ChatMessageOptions): Promise<void> {
5764
const file = await this.getFile();
58-
const thinks = file.thinks ?? {};
59-
60-
file.messages.push({ message, createdAt: Math.floor(Date.now() / 1000), thinkId: options?.thinkId });
61-
62-
// trim: keep last 1000, clean up orphaned thinks
63-
if (file.messages.length > 1000) {
64-
const removed = file.messages.splice(0, file.messages.length - 1000);
65-
for (const row of removed) {
66-
if (row.thinkId) delete thinks[row.thinkId];
65+
const id = file.nextId ?? 1;
66+
file.nextId = id + 1;
67+
file.messages.push({ id, message, createdAt: Math.floor(Date.now() / 1000), thinkId: options?.thinkId });
68+
69+
const compactedSet = new Set(file.compactedIds);
70+
const nonCompacted = file.messages.filter(m => !compactedSet.has(m.id!));
71+
if (nonCompacted.length > 1000) {
72+
for (const m of nonCompacted.slice(0, nonCompacted.length - 1000)) {
73+
compactedSet.add(m.id!);
6774
}
75+
file.compactedIds = [...compactedSet];
6876
}
6977

70-
file.thinks = thinks;
7178
await this.writeThreadFile(file);
7279
}
7380

7481
async getAllMessages(): Promise<StoredMessage[]> {
75-
return (await this.getFile()).messages;
82+
const file = await this.getFile();
83+
const compactedSet = new Set(file.compactedIds);
84+
return file.messages.filter(m => !compactedSet.has(m.id!));
7685
}
7786

78-
async getMessages(maxTokens: number): Promise<ChatMessage[]> {
79-
return applyTokenLimit((await this.getAllMessages()).map((r) => r.message), maxTokens);
87+
async getMessages(): Promise<ChatMessage[]> {
88+
return (await this.getAllMessages()).map((r) => r.message);
8089
}
8190

8291
async replaceAllMessages(messages: StoredMessage[]): Promise<void> {
8392
const file = await this.getFile();
84-
file.messages = [...messages];
93+
const compactedSet = new Set(file.compactedIds);
94+
const compacted = file.messages.filter(m => compactedSet.has(m.id!));
95+
const newMessages = messages.map(m => {
96+
if (m.id != null) return m;
97+
const id = file.nextId ?? 1;
98+
file.nextId = id + 1;
99+
return { ...m, id };
100+
});
101+
file.messages = [...compacted, ...newMessages];
85102
file.thinks = {};
86103
await this.writeThreadFile(file);
87104
}
88105

106+
async markMessagesAsCompacted(ids: number[]): Promise<void> {
107+
if (ids.length === 0) return;
108+
const file = await this.getFile();
109+
const compactedSet = new Set(file.compactedIds);
110+
for (const id of ids) compactedSet.add(id);
111+
file.compactedIds = [...compactedSet];
112+
await this.writeThreadFile(file);
113+
}
114+
115+
async searchMessages(query: string, limit: number = 20): Promise<StoredMessage[]> {
116+
const file = await this.getFile();
117+
const lower = query.toLowerCase();
118+
return file.messages
119+
.filter(m => {
120+
const c = m.message.content;
121+
return (typeof c === 'string' ? c : JSON.stringify(c)).toLowerCase().includes(lower);
122+
})
123+
.slice(-limit);
124+
}
125+
89126
async clearMessages(): Promise<void> {
90127
this.cache = undefined;
91128
if (existsSync(this.filePath)) {

‎packages/scorpio.ai/src/Saver/AgentMemorySaver.ts‎

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,37 +1,59 @@
11
import { IAgentSaverService, ChatMessage, StoredMessage, ChatMessageOptions } from "./IAgentSaverService";
2-
import { applyTokenLimit } from "./messageSerializer";
32

43
/**
54
* 纯内存实现的 AgentSaver，不持久化。
65
* 适用于临时会话、单次任务或测试场景。
76
*/
87
export class AgentMemorySaver implements IAgentSaverService {
98
private messages: StoredMessage[] = [];
9+
private compactedIds = new Set<number>();
1010
private thinks: Record<string, StoredMessage[]> = {};
1111
private metadata: Record<string, string> = {};
12+
private nextId = 1;
1213

1314
async getAllMessages(): Promise<StoredMessage[]> {
14-
return [...this.messages];
15+
return this.messages.filter(m => !this.compactedIds.has(m.id!));
1516
}
1617

17-
async getMessages(maxTokens: number): Promise<ChatMessage[]> {
18-
return applyTokenLimit(this.messages.map((r) => r.message), maxTokens);
18+
async getMessages(): Promise<ChatMessage[]> {
19+
return (await this.getAllMessages()).map(r => r.message);
1920
}
2021

2122
async pushMessage(message: ChatMessage, options?: ChatMessageOptions): Promise<void> {
22-
this.messages.push({ message, createdAt: Math.floor(Date.now() / 1000), thinkId: options?.thinkId });
23+
this.messages.push({ id: this.nextId++, message, createdAt: Math.floor(Date.now() / 1000), thinkId: options?.thinkId });
24+
const nonCompacted = this.messages.filter(m => !this.compactedIds.has(m.id!));
25+
if (nonCompacted.length > 1000) {
26+
for (const m of nonCompacted.slice(0, nonCompacted.length - 1000)) {
27+
this.compactedIds.add(m.id!);
28+
}
29+
}
2330
}
2431

2532
async replaceAllMessages(messages: StoredMessage[]): Promise<void> {
26-
this.messages = [...messages];
27-
this.thinks = {};
33+
const compacted = this.messages.filter(m => this.compactedIds.has(m.id!));
34+
this.messages = [...compacted, ...messages.map(m => ({ ...m, id: m.id ?? this.nextId++ }))];
2835
}
2936

3037
async clearMessages(): Promise<void> {
3138
this.messages = [];
39+
this.compactedIds.clear();
3240
this.thinks = {};
3341
}
3442

43+
async markMessagesAsCompacted(ids: number[]): Promise<void> {
44+
for (const id of ids) this.compactedIds.add(id);
45+
}
46+
47+
async searchMessages(query: string, limit: number = 20): Promise<StoredMessage[]> {
48+
const lower = query.toLowerCase();
49+
return this.messages
50+
.filter(m => {
51+
const c = m.message.content;
52+
return (typeof c === 'string' ? c : JSON.stringify(c)).toLowerCase().includes(lower);
53+
})
54+
.slice(-limit);
55+
}
56+
3557
async getThink(thinkId: string): Promise<StoredMessage[]> {
3658
return this.thinks[thinkId] ?? [];
3759
}

‎packages/scorpio.ai/src/Saver/AgentPostgresSaver.ts‎

Lines changed: 48 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ import { IAgentSaverService, ChatMessage, StoredMessage, ChatMessageOptions } fr
33
import { ILoggerService, ILogger } from "../Logger";
44
import { inject } from "scorpio.di";
55
import { T_DBUrl, T_DBTable } from "../Core/tokens";
6-
import { applyTokenLimit } from "./messageSerializer";
76

87
// ─────────────────────────────────────────────────────────────────────────────
98
// AgentPostgresSaver
@@ -17,6 +16,7 @@ export class AgentPostgresSaver implements IAgentSaverService {
1716

1817
private readonly table: string;
1918
private get thinksTable() { return `${this.table}_thinks`; }
19+
private get metadataTable() { return `${this.table}_metadata`; }
2020

2121
constructor(
2222
@inject(T_DBTable) table: string,
@@ -33,8 +33,6 @@ export class AgentPostgresSaver implements IAgentSaverService {
3333
return this.setupPromise;
3434
}
3535

36-
private get metadataTable() { return `${this.table}_metadata`; }
37-
3836
private async initTables(): Promise<void> {
3937
await this.pool.query(`
4038
CREATE TABLE IF NOT EXISTS ${this.table} (
@@ -57,6 +55,12 @@ export class AgentPostgresSaver implements IAgentSaverService {
5755
CREATE INDEX IF NOT EXISTS idx_${this.table}_thinks_think_id
5856
ON ${this.thinksTable} (think_id);
5957
`);
58+
try {
59+
await this.pool.query(`ALTER TABLE ${this.table} ADD COLUMN compacted INTEGER NOT NULL DEFAULT 0`);
60+
} catch { /* column already exists */ }
61+
await this.pool.query(
62+
`CREATE INDEX IF NOT EXISTS idx_${this.table}_fts ON ${this.table} USING gin(to_tsvector('simple', data))`
63+
);
6064
}
6165

6266
async pushMessage(message: ChatMessage, options?: ChatMessageOptions): Promise<void> {
@@ -66,9 +70,9 @@ export class AgentPostgresSaver implements IAgentSaverService {
6670
[JSON.stringify(message), Math.floor(Date.now() / 1000), options?.thinkId ?? null]
6771
);
6872
await this.pool.query(`
69-
DELETE FROM ${this.table}
70-
WHERE id < (
71-
SELECT id FROM ${this.table}
73+
UPDATE ${this.table} SET compacted = 1
74+
WHERE compacted = 0 AND id < (
75+
SELECT id FROM ${this.table} WHERE compacted = 0
7276
ORDER BY id DESC
7377
LIMIT 1 OFFSET 999
7478
)
@@ -79,9 +83,10 @@ export class AgentPostgresSaver implements IAgentSaverService {
7983
await this.ensureSetup();
8084
try {
8185
const result = await this.pool.query(
82-
`SELECT data, created_at, think_id FROM ${this.table} ORDER BY id`
86+
`SELECT id, data, created_at, think_id FROM ${this.table} WHERE compacted = 0 ORDER BY id`
8387
);
84-
return result.rows.map((r: { data: string; created_at: number; think_id: string | null }) => ({
88+
return result.rows.map((r: any) => ({
89+
id: parseInt(r.id),
8590
message: JSON.parse(r.data) as ChatMessage,
8691
createdAt: r.created_at,
8792
thinkId: r.think_id ?? undefined,
@@ -92,15 +97,15 @@ export class AgentPostgresSaver implements IAgentSaverService {
9297
}
9398
}
9499

95-
async getMessages(maxTokens: number): Promise<ChatMessage[]> {
96-
return applyTokenLimit((await this.getAllMessages()).map((r) => r.message), maxTokens);
100+
async getMessages(): Promise<ChatMessage[]> {
101+
return (await this.getAllMessages()).map((r) => r.message);
97102
}
98103

99104
async replaceAllMessages(messages: StoredMessage[]): Promise<void> {
100105
await this.ensureSetup();
101106
await this.pool.query(`BEGIN`);
102107
try {
103-
await this.pool.query(`DELETE FROM ${this.table}`);
108+
await this.pool.query(`DELETE FROM ${this.table} WHERE compacted = 0`);
104109
await this.pool.query(`DELETE FROM ${this.thinksTable}`);
105110
for (const stored of messages) {
106111
await this.pool.query(
@@ -115,6 +120,38 @@ export class AgentPostgresSaver implements IAgentSaverService {
115120
}
116121
}
117122

123+
async markMessagesAsCompacted(ids: number[]): Promise<void> {
124+
if (ids.length === 0) return;
125+
await this.ensureSetup();
126+
const placeholders = ids.map((_, i) => `$${i + 1}`).join(',');
127+
await this.pool.query(
128+
`UPDATE ${this.table} SET compacted = 1 WHERE id IN (${placeholders})`,
129+
ids,
130+
);
131+
}
132+
133+
async searchMessages(query: string, limit: number = 20): Promise<StoredMessage[]> {
134+
await this.ensureSetup();
135+
try {
136+
const result = await this.pool.query(
137+
`SELECT id, data, created_at, think_id FROM ${this.table}
138+
WHERE to_tsvector('simple', data) @@ plainto_tsquery('simple', $1)
139+
ORDER BY ts_rank(to_tsvector('simple', data), plainto_tsquery('simple', $1)) DESC
140+
LIMIT $2`,
141+
[query, limit],
142+
);
143+
return result.rows.map((r: any) => ({
144+
id: parseInt(r.id),
145+
message: JSON.parse(r.data) as ChatMessage,
146+
createdAt: r.created_at,
147+
thinkId: r.think_id ?? undefined,
148+
}));
149+
} catch (error: any) {
150+
this.logger?.warn(`Postgres FTS 搜索失败: ${error.message}`);
151+
return [];
152+
}
153+
}
154+
118155
async clearMessages(): Promise<void> {
119156
await this.ensureSetup();
120157
await this.pool.query(`DELETE FROM ${this.table}`);

0 commit comments

Comments
 (0)