UNPKG

@xynehq/jaf

Version:

Juspay Agent Framework - A purely functional agent framework with immutable state and composable tools

495 lines (494 loc) 21.1 kB
import { createSuccess, createFailure, createMemoryConnectionError, createMemoryNotFoundError, createMemoryStorageError } from '../types'; /** * PostgreSQL memory provider - fully persistent with advanced querying * Best for production environments requiring complex queries and full persistence */ export async function createPostgresProvider(config, postgresClient) { const fullConfig = { ...config, type: 'postgres', host: config.host ?? 'localhost', port: config.port ?? 5432, database: config.database ?? 'jaf_memory', username: config.username ?? 'postgres', ssl: config.ssl ?? false, tableName: config.tableName ?? 'conversations', maxConnections: config.maxConnections ?? 10 }; try { // Test connection and create table if needed await initializeSchema(postgresClient, fullConfig); console.log(`[MEMORY:Postgres] Connected to PostgreSQL at ${fullConfig.host}:${fullConfig.port}/${fullConfig.database}`); } catch (error) { throw createMemoryConnectionError('PostgreSQL', error); } const ensureConnected = () => { if (!postgresClient) { throw createMemoryConnectionError('PostgreSQL', new Error('PostgreSQL client not initialized')); } return postgresClient; }; const storeMessages = async (conversationId, messages, metadata) => { const client = ensureConnected(); try { const now = new Date(); const conversationMetadata = { totalMessages: messages.length, traceId: metadata?.traceId, ...metadata }; const sql = ` INSERT INTO ${fullConfig.tableName} (conversation_id, user_id, messages, metadata, created_at, updated_at, last_activity) VALUES ($1, $2, $3, $4, $5, $6, $7) ON CONFLICT (conversation_id) DO UPDATE SET messages = $3, metadata = $4, updated_at = $6, last_activity = $7 `; await client.query(sql, [ conversationId, metadata?.userId || null, JSON.stringify(messages), JSON.stringify(conversationMetadata), now, now, now ]); console.log(`[MEMORY:Postgres] Stored ${messages.length} messages for conversation ${conversationId}`); return createSuccess(undefined); } catch (error) { return createFailure(createMemoryStorageError('store messages', 'PostgreSQL', error)); } }; const getConversation = async (conversationId) => { const client = ensureConnected(); try { console.log(`[MEMORY:Postgres] Getting conversation ${conversationId}`); const sql = ` SELECT conversation_id, user_id, messages, metadata, created_at, updated_at, last_activity FROM ${fullConfig.tableName} WHERE conversation_id = $1 `; console.log(`[MEMORY:Postgres] Executing SQL: ${sql}`); console.log(`[MEMORY:Postgres] Parameters:`, [conversationId]); const result = await client.query(sql, [conversationId]); console.log(`[MEMORY:Postgres] Query result: ${result.rows.length} rows found`); if (result.rows.length === 0) { console.log(`[MEMORY:Postgres] No conversation found for ${conversationId}`); return createSuccess(null); } const row = result.rows[0]; console.log(`[MEMORY:Postgres] Raw row data:`, { conversation_id: row.conversation_id, user_id: row.user_id, messages_length: row.messages?.length || 0, metadata_length: row.metadata?.length || 0, created_at: row.created_at, updated_at: row.updated_at, last_activity: row.last_activity }); // Update last activity const updateSQL = ` UPDATE ${fullConfig.tableName} SET last_activity = NOW() WHERE conversation_id = $1 `; await client.query(updateSQL, [conversationId]); console.log(`[MEMORY:Postgres] Parsing messages JSON:`, row.messages); console.log(`[MEMORY:Postgres] Messages type:`, typeof row.messages); // Handle both string and object messages (PostgreSQL JSONB can return either) let parsedMessages; if (typeof row.messages === 'string') { parsedMessages = JSON.parse(row.messages); } else if (Array.isArray(row.messages)) { parsedMessages = row.messages; } else { throw new Error('Invalid messages format: expected string or array'); } console.log(`[MEMORY:Postgres] Parsed ${parsedMessages.length} messages`); console.log(`[MEMORY:Postgres] Parsing metadata JSON:`, row.metadata); console.log(`[MEMORY:Postgres] Metadata type:`, typeof row.metadata); // Handle both string and object metadata (PostgreSQL JSONB can return either) let parsedMetadata; if (typeof row.metadata === 'string') { parsedMetadata = JSON.parse(row.metadata || '{}'); } else if (typeof row.metadata === 'object' && row.metadata !== null) { parsedMetadata = row.metadata; } else { parsedMetadata = {}; } console.log(`[MEMORY:Postgres] Parsed metadata:`, parsedMetadata); const conversation = { conversationId: row.conversation_id, userId: row.user_id, messages: parsedMessages, metadata: { ...parsedMetadata, createdAt: new Date(row.created_at), updatedAt: new Date(row.updated_at), lastActivity: new Date(), // Use current time since we just updated it totalMessages: parsedMessages.length } }; console.log(`[MEMORY:Postgres] Successfully created conversation object with ${conversation.messages.length} messages`); console.log(`[MEMORY:Postgres] Retrieved conversation ${conversationId} with ${conversation.messages.length} messages`); return createSuccess(conversation); } catch (error) { console.error(`[MEMORY:Postgres] Error in getConversation:`, error); console.error(`[MEMORY:Postgres] Error stack:`, error.stack); return createFailure(createMemoryStorageError('get conversation', 'PostgreSQL', error)); } }; const appendMessages = async (conversationId, messages, metadata) => { const client = ensureConnected(); try { // Get existing conversation const existingResult = await getConversation(conversationId); if (!existingResult.success) { return existingResult; } if (!existingResult.data) { return createFailure(createMemoryNotFoundError(conversationId, 'PostgreSQL')); } const existing = existingResult.data; const updatedMessages = [...existing.messages, ...messages]; const now = new Date(); const updatedMetadata = { ...existing.metadata, totalMessages: updatedMessages.length, traceId: metadata?.traceId || existing.metadata?.traceId, ...metadata }; const sql = ` UPDATE ${fullConfig.tableName} SET messages = $1, metadata = $2, updated_at = $3, last_activity = $3 WHERE conversation_id = $4 `; await client.query(sql, [ JSON.stringify(updatedMessages), JSON.stringify(updatedMetadata), now, conversationId ]); console.log(`[MEMORY:Postgres] Appended ${messages.length} messages to conversation ${conversationId} (total: ${updatedMessages.length})`); return createSuccess(undefined); } catch (error) { return createFailure(createMemoryStorageError('append messages', 'PostgreSQL', error)); } }; const findConversations = async (query) => { const client = ensureConnected(); try { let sql = ` SELECT conversation_id, user_id, messages, metadata, created_at, updated_at, last_activity FROM ${fullConfig.tableName} WHERE 1=1 `; const params = []; let paramIndex = 1; if (query.conversationId) { sql += ` AND conversation_id = $${paramIndex}`; params.push(query.conversationId); paramIndex++; } if (query.userId) { sql += ` AND user_id = $${paramIndex}`; params.push(query.userId); paramIndex++; } if (query.traceId) { sql += ` AND metadata->>'traceId' = $${paramIndex}`; params.push(query.traceId); paramIndex++; } if (query.since) { sql += ` AND created_at >= $${paramIndex}`; params.push(query.since); paramIndex++; } if (query.until) { sql += ` AND created_at <= $${paramIndex}`; params.push(query.until); paramIndex++; } // Sort by last activity (most recent first) sql += ` ORDER BY last_activity DESC`; // Add pagination if (query.limit) { sql += ` LIMIT $${paramIndex}`; params.push(query.limit); paramIndex++; } if (query.offset) { sql += ` OFFSET $${paramIndex}`; params.push(query.offset); paramIndex++; } const result = await client.query(sql, params); const conversations = result.rows.map(row => { // Handle both string and object messages (PostgreSQL JSONB can return either) let parsedMessages; if (typeof row.messages === 'string') { parsedMessages = JSON.parse(row.messages); } else if (Array.isArray(row.messages)) { parsedMessages = row.messages; } else { throw new Error('Invalid messages format: expected string or array'); } // Handle both string and object metadata (PostgreSQL JSONB can return either) let parsedMetadata; if (typeof row.metadata === 'string') { parsedMetadata = JSON.parse(row.metadata || '{}'); } else if (typeof row.metadata === 'object' && row.metadata !== null) { parsedMetadata = row.metadata; } else { parsedMetadata = {}; } return { conversationId: row.conversation_id, userId: row.user_id, messages: parsedMessages, metadata: { ...parsedMetadata, createdAt: new Date(row.created_at), updatedAt: new Date(row.updated_at), lastActivity: new Date(row.last_activity), totalMessages: parsedMessages.length } }; }); console.log(`[MEMORY:Postgres] Found ${conversations.length} conversations matching query`); return createSuccess(conversations); } catch (error) { return createFailure(createMemoryStorageError('find conversations', 'PostgreSQL', error)); } }; const getRecentMessages = async (conversationId, limit = 50) => { const conversationResult = await getConversation(conversationId); if (!conversationResult.success) { return conversationResult; } if (!conversationResult.data) { return createSuccess([]); } const messages = conversationResult.data.messages.slice(-limit); console.log(`[MEMORY:Postgres] Retrieved ${messages.length} recent messages for conversation ${conversationId}`); return createSuccess(messages); }; const deleteConversation = async (conversationId) => { const client = ensureConnected(); try { const sql = `DELETE FROM ${fullConfig.tableName} WHERE conversation_id = $1`; const result = await client.query(sql, [conversationId]); const deleted = result.rowCount > 0; console.log(`[MEMORY:Postgres] ${deleted ? 'Deleted' : 'Attempted to delete non-existent'} conversation ${conversationId}`); return createSuccess(deleted); } catch (error) { return createFailure(createMemoryStorageError('delete conversation', 'PostgreSQL', error)); } }; const clearUserConversations = async (userId) => { const client = ensureConnected(); try { const sql = `DELETE FROM ${fullConfig.tableName} WHERE user_id = $1`; const result = await client.query(sql, [userId]); const deletedCount = result.rowCount; console.log(`[MEMORY:Postgres] Cleared ${deletedCount} conversations for user ${userId}`); return createSuccess(deletedCount); } catch (error) { return createFailure(createMemoryStorageError('clear user conversations', 'PostgreSQL', error)); } }; const getStats = async (userId) => { const client = ensureConnected(); try { let sql = ` SELECT COUNT(*) as total_conversations, MIN(created_at) as oldest_conversation, MAX(created_at) as newest_conversation, SUM((metadata->>'totalMessages')::int) as total_messages FROM ${fullConfig.tableName} `; const params = []; if (userId) { sql += ' WHERE user_id = $1'; params.push(userId); } const result = await client.query(sql, params); const row = result.rows[0]; return createSuccess({ totalConversations: parseInt(row.total_conversations) || 0, totalMessages: parseInt(row.total_messages) || 0, oldestConversation: row.oldest_conversation ? new Date(row.oldest_conversation) : undefined, newestConversation: row.newest_conversation ? new Date(row.newest_conversation) : undefined }); } catch (error) { return createFailure(createMemoryStorageError('get stats', 'PostgreSQL', error)); } }; const healthCheck = async () => { const start = Date.now(); try { const client = ensureConnected(); // Test basic connectivity await client.query('SELECT 1'); // Test table operations const testId = `health-check-${Date.now()}`; const storeResult = await storeMessages(testId, [{ role: 'user', content: 'health check' }]); if (!storeResult.success) { return createSuccess({ healthy: false, latencyMs: Date.now() - start, error: storeResult.error.message }); } const getResult = await getConversation(testId); if (!getResult.success) { return createSuccess({ healthy: false, latencyMs: Date.now() - start, error: getResult.error.message }); } const deleteResult = await deleteConversation(testId); if (!deleteResult.success) { return createSuccess({ healthy: false, latencyMs: Date.now() - start, error: deleteResult.error.message }); } const latencyMs = Date.now() - start; return createSuccess({ healthy: true, latencyMs }); } catch (error) { return createSuccess({ healthy: false, latencyMs: Date.now() - start, error: error instanceof Error ? error.message : 'Unknown error' }); } }; const close = async () => { try { if (postgresClient) { console.log('[MEMORY:Postgres] Closing PostgreSQL connection'); await postgresClient.end(); } return createSuccess(undefined); } catch (error) { return createFailure(createMemoryStorageError('close connection', 'PostgreSQL', error)); } }; const cleanupOldConversations = async (olderThanDays) => { const client = ensureConnected(); try { const sql = ` DELETE FROM ${fullConfig.tableName} WHERE last_activity < NOW() - INTERVAL '${olderThanDays} days' `; const result = await client.query(sql); const deletedCount = result.rowCount; console.log(`[MEMORY:Postgres] Cleaned up ${deletedCount} conversations older than ${olderThanDays} days`); return createSuccess(deletedCount); } catch (error) { return createFailure(createMemoryStorageError('cleanup old conversations', 'PostgreSQL', error)); } }; const getAnalytics = async (userId) => { const client = ensureConnected(); try { let whereClause = ''; const params = []; if (userId) { whereClause = 'WHERE user_id = $1'; params.push(userId); } const sql = ` SELECT AVG((metadata->>'totalMessages')::int) as avg_messages, COUNT(CASE WHEN created_at >= NOW() - INTERVAL '7 days' THEN 1 END) as last_week, COUNT(CASE WHEN created_at >= NOW() - INTERVAL '30 days' THEN 1 END) as last_month, EXTRACT(hour FROM last_activity) as activity_hour, TO_CHAR(last_activity, 'Day') as activity_day FROM ${fullConfig.tableName} ${whereClause} GROUP BY EXTRACT(hour FROM last_activity), TO_CHAR(last_activity, 'Day') ORDER BY COUNT(*) DESC LIMIT 1 `; const result = await client.query(sql, params); const row = result.rows[0] || {}; return createSuccess({ averageMessagesPerConversation: parseFloat(row.avg_messages) || 0, conversationsLastWeek: parseInt(row.last_week) || 0, conversationsLastMonth: parseInt(row.last_month) || 0, mostActiveHour: parseInt(row.activity_hour) || 0, mostActiveDay: (row.activity_day || '').trim() }); } catch (error) { return createFailure(createMemoryStorageError('get analytics', 'PostgreSQL', error)); } }; return { storeMessages, getConversation, appendMessages, findConversations, getRecentMessages, deleteConversation, clearUserConversations, getStats, healthCheck, close }; } async function initializeSchema(client, config) { const createTableSQL = ` CREATE TABLE IF NOT EXISTS ${config.tableName} ( conversation_id VARCHAR(255) PRIMARY KEY, user_id VARCHAR(255), messages JSONB NOT NULL, metadata JSONB, created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(), updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(), last_activity TIMESTAMP WITH TIME ZONE DEFAULT NOW() ); CREATE INDEX IF NOT EXISTS idx_${config.tableName}_user_id ON ${config.tableName} (user_id); CREATE INDEX IF NOT EXISTS idx_${config.tableName}_created_at ON ${config.tableName} (created_at); CREATE INDEX IF NOT EXISTS idx_${config.tableName}_last_activity ON ${config.tableName} (last_activity); CREATE INDEX IF NOT EXISTS idx_${config.tableName}_metadata_gin ON ${config.tableName} USING GIN (metadata); CREATE INDEX IF NOT EXISTS idx_${config.tableName}_trace_id ON ${config.tableName} ((metadata->>'traceId')); `; await client.query(createTableSQL); console.log(`[MEMORY:Postgres] Schema initialized for table ${config.tableName}`); } //# sourceMappingURL=postgres.js.map