UNPKG

@juspay/neurolink

Version:

The pipe layer of an AI nervous system: one interface connecting provider neurons to your application, across three inference types — generate, stream and decide. `decide` returns typed, calibrated judgements from a non-generative model (~400ms, ~$0.00002

489 lines (488 loc) • 20.7 kB
/** * Conversation Memory Manager for NeuroLink * Handles in-memory conversation storage, session management, and context injection */ import { randomUUID } from "crypto"; import { DEFAULT_MAX_SESSIONS, MEMORY_THRESHOLD_PERCENTAGE, MESSAGES_PER_TURN, } from "../config/conversationMemory.js"; import { TokenUtils } from "../constants/tokens.js"; import { SummarizationEngine } from "../context/summarizationEngine.js"; import { context } from "@opentelemetry/api"; import { tracers, withSpan } from "../telemetry/index.js"; import { ConversationMemoryError } from "../types/index.js"; import { buildContextFromPointer, getEffectiveTokenThreshold, } from "../utils/conversationMemory.js"; import { logger } from "../utils/logger.js"; import { assistantTurnIdentity, isTurnAlreadyStored, userTurnIdentity, } from "../utils/conversationTurnIdentity.js"; import { extractIsErrorText } from "../utils/toolResultStatus.js"; export class ConversationMemoryManager { sessions = new Map(); config; isInitialized = false; summarizationEngine = new SummarizationEngine(); /** * Track sessions currently being summarized to prevent race conditions */ summarizationInProgress = new Set(); constructor(config) { this.config = config; } /** * Initialize the memory manager */ async initialize() { if (this.isInitialized) { return; } try { this.isInitialized = true; logger.info("ConversationMemoryManager initialized", { storage: "in-memory", maxSessions: this.config.maxSessions, maxTurnsPerSession: this.config.maxTurnsPerSession, }); } catch (error) { throw new ConversationMemoryError("Failed to initialize conversation memory", "CONFIG_ERROR", { error: error instanceof Error ? error.message : String(error), }); } } /** Whether this memory manager can persist data (always true for in-memory within process) */ get canPersist() { return true; } /** Whether Redis client is configured (always false for in-memory) */ get isRedisConfigured() { return false; } /** Get health status for monitoring */ getHealthStatus() { return { initialized: this.isInitialized, connected: false, }; } /** * Store a conversation turn for a session * TOKEN-BASED: Validates message size and triggers summarization based on tokens */ async storeConversationTurn(options) { return withSpan({ name: "memory.storeTurn", tracer: tracers.memory, attributes: { "memory.type": "in-memory", "session.id": options.sessionId, "memory.operation": "store_turn", }, }, async (span) => { await this.ensureInitialized(); try { // Get or create session let session = this.sessions.get(options.sessionId); if (!session) { session = this.createNewSession(options.sessionId, options.userId); this.sessions.set(options.sessionId, session); } if (isTurnAlreadyStored(session.messages, options.turnId)) { return; } const tokenThreshold = options.providerDetails ? getEffectiveTokenThreshold(options.providerDetails.provider, options.providerDetails.model, this.config.tokenThreshold, session.tokenThreshold) : this.config.tokenThreshold || 50000; const userMsg = await this.validateAndPrepareMessage(options.userMessage, "user", tokenThreshold, userTurnIdentity(options)); const assistantMsg = await this.validateAndPrepareMessage(options.aiResponse, "assistant", tokenThreshold, assistantTurnIdentity(options)); if (options.events && options.events.length > 0) { assistantMsg.events = options.events; } session.messages.push(userMsg); // Pinned skill activations ride between ask and answer, mirroring // the actual order (ask → skill loaded → answer). Stored verbatim — // skill instructions are never truncated. if (options.skillMessages && options.skillMessages.length > 0) { session.messages.push(...options.skillMessages); } session.messages.push(assistantMsg); session.lastActivity = Date.now(); // Store API-reported token counts if available if (options.tokenUsage) { session.lastApiTokenCount = options.tokenUsage; } span.setAttribute("memory.message_count", session.messages.length); const shouldSummarize = options.enableSummarization !== undefined ? options.enableSummarization : this.config.enableSummarization; if (shouldSummarize) { // Only trigger summarization if not already in progress for this session if (!this.summarizationInProgress.has(options.sessionId)) { const summarizeWithContext = context.bind(context.active(), async () => { try { await this.checkAndSummarize(session, tokenThreshold, options.requestId); } catch (error) { logger.error("Background summarization failed", { sessionId: session.sessionId, requestId: options.requestId, error: error instanceof Error ? error.message : String(error), }); } }); setImmediate(summarizeWithContext); } else { logger.debug("[ConversationMemoryManager] Summarization already in progress, skipping", { sessionId: options.sessionId, }); } } this.enforceSessionLimit(); } catch (error) { throw new ConversationMemoryError(`Failed to store conversation turn for session ${options.sessionId}`, "STORAGE_ERROR", { sessionId: options.sessionId, error: error instanceof Error ? error.message : String(error), }); } }); } /** * Validate and prepare a message before adding to session * Truncates if message exceeds token limit */ async validateAndPrepareMessage(content, role, threshold, identity) { const id = identity?.id ?? randomUUID(); const tokenCount = TokenUtils.estimateTokenCount(content); const maxMessageSize = Math.floor(threshold * MEMORY_THRESHOLD_PERCENTAGE); if (tokenCount > maxMessageSize) { const truncated = TokenUtils.truncateToTokenLimit(content, maxMessageSize); logger.warn("Message truncated due to token limit", { id, role, originalTokens: tokenCount, threshold, truncatedTo: maxMessageSize, }); return { id, role, content: truncated, timestamp: new Date().toISOString(), metadata: { ...identity?.metadata, truncated: true, }, }; } return { id, role, content, timestamp: new Date().toISOString(), ...(identity?.metadata ? { metadata: identity.metadata } : {}), }; } /** * Check if summarization is needed based on token count */ async checkAndSummarize(session, threshold, requestId) { // Acquire lock - if already in progress, skip if (this.summarizationInProgress.has(session.sessionId)) { logger.debug("[ConversationMemoryManager] Summarization already in progress, skipping", { sessionId: session.sessionId, }); return; } this.summarizationInProgress.add(session.sessionId); try { await this.summarizationEngine.checkAndSummarize(session, threshold, this.config, "[ConversationMemory]", requestId); } catch (error) { logger.error("Token counting or summarization failed", { sessionId: session.sessionId, error: error instanceof Error ? error.message : String(error), }); } finally { // Release lock when done this.summarizationInProgress.delete(session.sessionId); } } /** * Estimate total tokens for a list of messages */ estimateTokens(messages) { return messages.reduce((total, msg) => { let msgTokens = TokenUtils.estimateTokenCount(msg.content); if (msg.events && Array.isArray(msg.events) && msg.events.length > 0) { const eventsJson = JSON.stringify(msg.events); msgTokens += TokenUtils.estimateTokenCount(eventsJson); } return total + msgTokens; }, 0); } /** * Build context messages for AI prompt injection (TOKEN-BASED) * Returns messages from pointer onwards (or all if no pointer) * Now consistently async to match Redis implementation */ async buildContextMessages(sessionId, _userId, _enableSummarization, requestId) { return withSpan({ name: "memory.buildContext", tracer: tracers.memory, attributes: { "memory.type": "in-memory", "session.id": sessionId, "memory.operation": "build_context", }, }, async (span) => { const session = this.sessions.get(sessionId); const messages = session ? buildContextFromPointer(session, requestId) : []; span.setAttribute("memory.message_count", messages.length); return messages; }); } getSession(sessionId, _userId) { return this.sessions.get(sessionId); } createSummarySystemMessage(content, summarizesFrom, summarizesTo) { return { id: `summary-${randomUUID()}`, role: "system", content: `Summary of previous conversation turns:\n\n${content}`, timestamp: new Date().toISOString(), metadata: { isSummary: true, summarizesFrom, summarizesTo, }, }; } async ensureInitialized() { if (!this.isInitialized) { await this.initialize(); } } createNewSession(sessionId, userId) { return { sessionId, userId, messages: [], createdAt: Date.now(), lastActivity: Date.now(), }; } enforceSessionLimit() { const maxSessions = this.config.maxSessions || DEFAULT_MAX_SESSIONS; if (this.sessions.size <= maxSessions) { return; } const sessions = Array.from(this.sessions.entries()).sort(([, a], [, b]) => a.lastActivity - b.lastActivity); const sessionsToRemove = sessions.slice(0, this.sessions.size - maxSessions); for (const [sessionId] of sessionsToRemove) { this.sessions.delete(sessionId); } } async getStats() { await this.ensureInitialized(); const sessions = Array.from(this.sessions.values()); // Pinned skill messages are extra rows inside a turn — exclude them so // a skill-activating turn still counts as one turn. const totalTurns = sessions.reduce((sum, session) => sum + session.messages.filter((msg) => !msg.metadata?.isSkill).length / MESSAGES_PER_TURN, 0); return { totalSessions: sessions.length, totalTurns, }; } /** * List all sessions with metadata * @param userId - Optional user ID to filter sessions * @returns Array of session list items with metadata */ async listSessions(userId) { await this.ensureInitialized(); const sessions = Array.from(this.sessions.values()); const now = Date.now(); return sessions .filter((session) => !userId || session.userId === userId) .map((session) => { const lastActive = this.formatTimeAgo(now - session.lastActivity); return { id: session.sessionId, title: session.sessionId, // In-memory doesn't store title, use sessionId createdAt: new Date(session.createdAt).toISOString(), updatedAt: new Date(session.lastActivity).toISOString(), userId: session.userId, messageCount: session.messages.length, lastActive, }; }) .sort((a, b) => new Date(b.updatedAt).getTime() - new Date(a.updatedAt).getTime()); } /** * Format milliseconds into human-readable time ago string */ formatTimeAgo(ms) { const seconds = Math.floor(ms / 1000); const minutes = Math.floor(seconds / 60); const hours = Math.floor(minutes / 60); const days = Math.floor(hours / 24); if (days > 0) { return `${days} day${days > 1 ? "s" : ""} ago`; } if (hours > 0) { return `${hours} hour${hours > 1 ? "s" : ""} ago`; } if (minutes > 0) { return `${minutes} minute${minutes > 1 ? "s" : ""} ago`; } return "just now"; } async clearSession(sessionId) { return withSpan({ name: "memory.clear", tracer: tracers.memory, attributes: { "memory.type": "in-memory", "session.id": sessionId, "memory.operation": "clear_session", }, }, async (span) => { const session = this.sessions.get(sessionId); if (!session) { span.setAttribute("memory.session_found", false); return false; } this.sessions.delete(sessionId); span.setAttribute("memory.session_found", true); logger.info("Session cleared", { sessionId }); return true; }); } async clearAllSessions() { const sessionIds = Array.from(this.sessions.keys()); this.sessions.clear(); logger.info("All sessions cleared", { clearedCount: sessionIds.length }); } /** * Get the raw messages array for a session. * Returns the full messages list without context filtering or summarization. * Returns a deep copy to prevent external mutation of internal state. */ async getSessionMessages(sessionId, _userId) { await this.ensureInitialized(); const session = this.sessions.get(sessionId); return session ? session.messages.map((msg) => ({ ...msg })) : []; } /** * Replace the entire messages array for a session. * Creates the session if it does not exist. * Resets summary pointers since old pointers may reference messages that no longer exist. */ async setSessionMessages(sessionId, messages, userId) { await this.ensureInitialized(); let session = this.sessions.get(sessionId); if (!session) { session = this.createNewSession(sessionId, userId); this.sessions.set(sessionId, session); this.enforceSessionLimit(); } session.messages = [...messages]; session.summarizedUpToMessageId = undefined; session.summarizedMessage = undefined; session.lastTokenCount = undefined; session.lastCountedAt = undefined; session.lastActivity = Date.now(); } /** * Persist a step's tool calls and results as `tool_call` / `tool_result` * messages, mirroring the Redis manager. * * Parity fix: this used to exist only on the Redis backend, so an in-memory * session never turned tool activity into messages and every downstream path * that reasons about tool batches (compaction, pruning, pair repair) saw a * different history shape depending on `STORAGE_TYPE`. * * Calls are written before results — the same order the Redis flush uses — * and `toolCallId` is carried on BOTH sides so `repairToolPairs` can match by * id rather than adjacency (a parallel batch has no positional pairing). */ async storeToolExecution(sessionId, userId, toolCalls, toolResults, currentTime) { await this.ensureInitialized(); let session = this.sessions.get(sessionId); if (!session) { session = this.createNewSession(sessionId, userId); this.sessions.set(sessionId, session); this.enforceSessionLimit(); } const timestamp = (currentTime ?? new Date()).toISOString(); const toolNameById = new Map(); for (const toolCall of toolCalls ?? []) { const toolCallId = toolCall.toolCallId ?? ""; const toolName = toolCall.toolName ?? "unknown"; if (toolCallId) { toolNameById.set(toolCallId, toolName); } const callStepIndex = toolCall.stepIndex; session.messages.push({ id: randomUUID(), role: "tool_call", content: "", // Tool calls carry their payload in `args`, not content. tool: toolName, ...(toolCallId ? { toolCallId } : {}), args: (toolCall.args ?? {}), // Zero-based loop step, so replay can keep sequential steps apart. ...(typeof callStepIndex === "number" ? { metadata: { stepIndex: callStepIndex } } : {}), timestamp, }); } for (const toolResult of toolResults ?? []) { const toolCallId = toolResult.toolCallId ?? ""; const toolName = (toolCallId ? toolNameById.get(toolCallId) : undefined) ?? String(toolResult.toolName ?? "unknown"); const rawOutput = "output" in toolResult ? toolResult.output : toolResult.result; let content; if (typeof rawOutput === "string") { content = rawOutput; } else { try { content = JSON.stringify(rawOutput ?? null) ?? "null"; } catch (error) { content = `[Serialization failed: ${error instanceof Error ? error.message : String(error)}]`; } } // A result the tool RETURNED with `isError: true` (a thrown error that // neurolink.executeTool converted, a breaker refusal, an MCP failure) // is a failure too; recording it as success made the next turn's // replay say "→ ok" for a call that failed. const error = toolResult.error !== undefined && toolResult.error !== null ? String(toolResult.error) : extractIsErrorText(rawOutput); const stepIndex = toolResult.stepIndex; session.messages.push({ id: randomUUID(), role: "tool_result", content, tool: toolName, ...(toolCallId ? { toolCallId } : {}), result: { success: error === undefined, ...(error !== undefined ? { error } : {}), }, ...(typeof stepIndex === "number" ? { metadata: { stepIndex } } : {}), timestamp, }); } session.lastActivity = Date.now(); } /** Close/shutdown — no-op for in-memory manager (no external connections to release) */ async close() { // In-memory manager has nothing to close } }