103 lines
3.7 KiB
TypeScript
103 lines
3.7 KiB
TypeScript
import { Injectable } from "@nestjs/common";
|
||
import type { EmotionTag } from "@kokorone/shared";
|
||
import { HIGH_EMOTION_THRESHOLD } from "./constants.js";
|
||
import { getRedisClient } from "../cache/redis-client.js";
|
||
|
||
export interface WorkingMemoryEntry {
|
||
role: "user" | "character";
|
||
content: string;
|
||
timestamp: Date;
|
||
emotionTag?: EmotionTag;
|
||
emotionIntensity?: number; // 0~1，未標記情緒的訊息（如快速通道問候）可省略
|
||
}
|
||
|
||
interface SerializedEntry extends Omit<WorkingMemoryEntry, "timestamp"> {
|
||
timestamp: string;
|
||
}
|
||
|
||
const DEFAULT_TOKEN_LIMIT = 200;
|
||
const REDIS_KEY_PREFIX = "kokorone:working-memory:";
|
||
// session 閒置這麼久沒有新訊息就任由 Redis 自然過期，避免無限累積殭屍 session。
|
||
const SESSION_TTL_SECONDS = 60 * 60 * 6;
|
||
|
||
// 粗略估算：Mock 階段不需要真實 tokenizer，先以字元數／2 近似。
|
||
function estimateTokens(text: string): number {
|
||
return Math.ceil(text.length / 2);
|
||
}
|
||
|
||
function isHighEmotion(entry: WorkingMemoryEntry): boolean {
|
||
return (entry.emotionIntensity ?? 0) >= HIGH_EMOTION_THRESHOLD;
|
||
}
|
||
|
||
function evictOverflow(entries: WorkingMemoryEntry[], tokenLimit: number): WorkingMemoryEntry[] {
|
||
let result = entries;
|
||
let totalTokens = result.reduce((sum, e) => sum + estimateTokens(e.content), 0);
|
||
|
||
while (totalTokens > tokenLimit && result.length > 0) {
|
||
// 最舊的低情緒段落先被裁掉；若全部都是高情緒，最後才犧牲最舊的一則以確保不超出上限。
|
||
let evictIndex = result.findIndex((e) => !isHighEmotion(e));
|
||
if (evictIndex === -1) {
|
||
evictIndex = 0;
|
||
}
|
||
const [evicted] = result.splice(evictIndex, 1);
|
||
totalTokens -= estimateTokens(evicted.content);
|
||
}
|
||
|
||
return result;
|
||
}
|
||
|
||
function serialize(entries: WorkingMemoryEntry[]): string {
|
||
const payload: SerializedEntry[] = entries.map((e) => ({ ...e, timestamp: e.timestamp.toISOString() }));
|
||
return JSON.stringify(payload);
|
||
}
|
||
|
||
function deserialize(raw: string): WorkingMemoryEntry[] {
|
||
const payload = JSON.parse(raw) as SerializedEntry[];
|
||
return payload.map((e) => ({ ...e, timestamp: new Date(e.timestamp) }));
|
||
}
|
||
|
||
// C-1 工作記憶：session 對話上下文緩衝，如同海馬迴暫存，長期寫入只能由 C-4 睡眠固化觸發。
|
||
// R-2：改走 Redis（session 快取），Redis 不可用時退回行程記憶體 Map，行為（token 上限、高情緒優先保留）不變。
|
||
@Injectable()
|
||
export class WorkingMemoryService {
|
||
private readonly fallback = new Map<string, WorkingMemoryEntry[]>();
|
||
private readonly tokenLimit = DEFAULT_TOKEN_LIMIT;
|
||
|
||
async append(sessionId: string, entry: WorkingMemoryEntry): Promise<void> {
|
||
const entries = await this.getContext(sessionId);
|
||
entries.push(entry);
|
||
const evicted = evictOverflow(entries, this.tokenLimit);
|
||
|
||
const redis = getRedisClient();
|
||
try {
|
||
await redis.set(REDIS_KEY_PREFIX + sessionId, serialize(evicted), "EX", SESSION_TTL_SECONDS);
|
||
this.fallback.delete(sessionId);
|
||
return;
|
||
} catch {
|
||
this.fallback.set(sessionId, evicted);
|
||
}
|
||
}
|
||
|
||
async getContext(sessionId: string): Promise<WorkingMemoryEntry[]> {
|
||
const redis = getRedisClient();
|
||
try {
|
||
const raw = await redis.get(REDIS_KEY_PREFIX + sessionId);
|
||
if (raw !== null) return deserialize(raw);
|
||
if (this.fallback.has(sessionId)) return [...this.fallback.get(sessionId)!];
|
||
return [];
|
||
} catch {
|
||
return [...(this.fallback.get(sessionId) ?? [])];
|
||
}
|
||
}
|
||
|
||
async clear(sessionId: string): Promise<void> {
|
||
this.fallback.delete(sessionId);
|
||
const redis = getRedisClient();
|
||
try {
|
||
await redis.del(REDIS_KEY_PREFIX + sessionId);
|
||
} catch {
|
||
// Redis 不可用時 fallback 已經在上面清掉，沒有需要額外處理的狀態。
|
||
}
|
||
}
|
||
}
|