File size: 4,472 Bytes
6111b2b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | /**
* Memory Injection — prepend retrieved memories into the request message list.
*
* Injection strategy:
* 1. If the provider supports system messages (most providers), inject as a
* leading system message so it takes effect without disrupting user turns.
* 2. Otherwise (fallback for providers that reject system role), inject as the
* first user message prefixed with the memory context label.
*
* Format: "Memory context: <content>"
*/
import { Memory } from "./types";
import { logger } from "../../../open-sse/utils/logger.ts";
const log = logger("MEMORY_INJECTION");
export interface ChatMessage {
role: "system" | "user" | "assistant";
content: string;
name?: string;
}
export interface ChatRequest {
model: string;
messages: ChatMessage[];
system?: string;
temperature?: number;
max_tokens?: number;
stream?: boolean;
[key: string]: unknown;
}
/**
* Providers known NOT to support a top-level system-role message.
* These receive memories injected as the first user message instead.
*/
const PROVIDERS_WITHOUT_SYSTEM_MESSAGE = new Set([
"o1",
"o1-mini",
"o1-preview",
"glm", // GLM/ZhipuAI rejects system role (#1701)
"glmt", // GLM Thinking variant
"glm-cn", // GLM China variant
"zai", // Z.AI uses same GLM backend
"qianfan", // Baidu ERNIE rejects system role
]);
/**
* Returns true when the given provider accepts a system-role message.
* Falls back to true for unknown/null providers (safe default).
*/
export function providerSupportsSystemMessage(provider: string | null | undefined): boolean {
if (!provider) return true;
const normalized = provider.toLowerCase().trim();
return !PROVIDERS_WITHOUT_SYSTEM_MESSAGE.has(normalized);
}
/**
* Format memories into a single labeled context string.
* Format: "Memory context: <content1>\n<content2>..."
*/
export function formatMemoryContext(memories: Memory[]): string {
if (!memories || memories.length === 0) return "";
const content = memories
.map((m) => m.content.trim())
.filter(Boolean)
.join("\n");
return content ? `Memory context: ${content}` : "";
}
/**
* Inject retrieved memories into the request message array.
*
* @param request - The chat completion request body
* @param memories - Memories retrieved for the current API key / session
* @param provider - Provider identifier used to choose injection strategy
* @returns A new request body with memories prepended to messages
*/
export function injectMemory(
request: ChatRequest,
memories: Memory[],
provider: string | null | undefined
): ChatRequest {
if (!memories || memories.length === 0) {
log.info("memory.injection.skipped", { reason: "no_memories", model: request.model });
return request;
}
const memoryText = formatMemoryContext(memories);
if (!memoryText) {
log.info("memory.injection.skipped", { reason: "empty_context", model: request.model });
return request;
}
const messages: ChatMessage[] = Array.isArray(request.messages) ? [...request.messages] : [];
if (providerSupportsSystemMessage(provider)) {
// Strategy 1: inject as a leading system message.
// Prepending before any existing system messages keeps memory context
// accessible without overriding the caller's own system instructions.
const memorySystemMessage: ChatMessage = { role: "system", content: memoryText };
log.info("memory.injection.injected", {
count: memories.length,
strategy: "system",
model: request.model,
});
return { ...request, messages: [memorySystemMessage, ...messages] };
} else {
// Strategy 2 (fallback): inject as the first user message.
// Used for providers like o1-mini that reject the system role.
const memoryUserMessage: ChatMessage = { role: "user", content: memoryText };
log.info("memory.injection.injected", {
count: memories.length,
strategy: "user",
model: request.model,
});
return { ...request, messages: [memoryUserMessage, ...messages] };
}
}
/**
* Returns true when memory injection should be attempted for this request.
*/
export function shouldInjectMemory(request: ChatRequest, config?: { enabled?: boolean }): boolean {
if (config?.enabled === false) return false;
return Array.isArray(request.messages) && request.messages.length > 0;
}
|