@xynehq/jaf
Version:
Juspay Agent Framework - A purely functional agent framework with immutable state and composable tools
405 lines • 14.2 kB
JavaScript
/**
* JAF ADK Layer - Artifact Storage System
*
* Provides persistent key-value storage for agent conversations
* with support for multiple storage backends
*/
// ========== In-Memory Storage ==========
export const createMemoryArtifactStorage = (config) => {
const storage = new Map();
const maxSize = config?.maxSize || 10 * 1024 * 1024; // 10MB default
const ttl = config?.ttl; // Optional TTL
// Helper to get session storage
const getSessionStorage = (sessionId) => {
if (!storage.has(sessionId)) {
storage.set(sessionId, new Map());
}
return storage.get(sessionId);
};
// Helper to check TTL
const isExpired = (artifact) => {
if (!ttl)
return false;
const age = Date.now() - artifact.metadata.lastModified.getTime();
return age > ttl * 1000;
};
// Helper to estimate size
const estimateSize = (value) => {
try {
return JSON.stringify(value).length;
}
catch {
return 0;
}
};
return {
get: async (sessionId, key) => {
const sessionStorage = getSessionStorage(sessionId);
const artifact = sessionStorage.get(key);
if (!artifact)
return null;
if (isExpired(artifact)) {
sessionStorage.delete(key);
return null;
}
return artifact;
},
set: async (sessionId, key, value, metadata) => {
const sessionStorage = getSessionStorage(sessionId);
const size = estimateSize(value);
if (size > maxSize) {
throw new Error(`Artifact size (${size} bytes) exceeds maximum allowed size (${maxSize} bytes)`);
}
const now = new Date();
const artifact = {
key,
value,
metadata: {
created: metadata?.created || sessionStorage.get(key)?.metadata.created || now,
lastModified: now,
contentType: metadata?.contentType,
size,
tags: metadata?.tags
}
};
sessionStorage.set(key, artifact);
return artifact;
},
delete: async (sessionId, key) => {
const sessionStorage = getSessionStorage(sessionId);
return sessionStorage.delete(key);
},
list: async (sessionId) => {
const sessionStorage = getSessionStorage(sessionId);
const artifacts = [];
for (const [key, artifact] of sessionStorage.entries()) {
if (!isExpired(artifact)) {
artifacts.push(artifact);
}
else {
sessionStorage.delete(key);
}
}
return artifacts;
},
clear: async (sessionId) => {
storage.delete(sessionId);
},
exists: async (sessionId, key) => {
const sessionStorage = getSessionStorage(sessionId);
const artifact = sessionStorage.get(key);
if (!artifact)
return false;
if (isExpired(artifact)) {
sessionStorage.delete(key);
return false;
}
return true;
}
};
};
// ========== Redis Storage ==========
export const createRedisArtifactStorage = (config) => {
let redis;
const keyPrefix = config.keyPrefix || 'jaf:artifacts:';
const maxSize = config.maxSize || 10 * 1024 * 1024; // 10MB default
const ttl = config.ttl; // Optional TTL in seconds
try {
// eslint-disable-next-line @typescript-eslint/no-var-requires
const Redis = require('ioredis');
redis = new Redis({
host: config.host,
port: config.port,
password: config.password,
db: config.database || 0,
connectTimeout: 5000, // 5 second timeout
enableOfflineQueue: false, // Fail fast if offline
maxRetriesPerRequest: 3,
retryStrategy: (times) => {
if (times > 3)
return null; // Stop retrying after 3 attempts
return Math.min(times * 100, 3000);
}
});
}
catch (error) {
throw new Error('Redis artifact storage requires ioredis to be installed');
}
const getKey = (sessionId, key) => `${keyPrefix}${sessionId}:${key}`;
const getSessionPattern = (sessionId) => `${keyPrefix}${sessionId}:*`;
return {
get: async (sessionId, key) => {
const redisKey = getKey(sessionId, key);
const data = await redis.get(redisKey);
if (!data)
return null;
try {
return JSON.parse(data);
}
catch {
return null;
}
},
set: async (sessionId, key, value, metadata) => {
const now = new Date();
const artifact = {
key,
value,
metadata: {
created: metadata?.created || now,
lastModified: now,
contentType: metadata?.contentType,
size: JSON.stringify(value).length,
tags: metadata?.tags
}
};
const serialized = JSON.stringify(artifact);
if (serialized.length > maxSize) {
throw new Error(`Artifact size exceeds maximum allowed size`);
}
const redisKey = getKey(sessionId, key);
if (ttl) {
await redis.setex(redisKey, ttl, serialized);
}
else {
await redis.set(redisKey, serialized);
}
return artifact;
},
delete: async (sessionId, key) => {
const redisKey = getKey(sessionId, key);
const result = await redis.del(redisKey);
return result > 0;
},
list: async (sessionId) => {
const pattern = getSessionPattern(sessionId);
const keys = await redis.keys(pattern);
if (keys.length === 0)
return [];
const values = await redis.mget(...keys);
const artifacts = [];
for (const value of values) {
if (value) {
try {
artifacts.push(JSON.parse(value));
}
catch {
// Skip invalid artifacts
}
}
}
return artifacts;
},
clear: async (sessionId) => {
const pattern = getSessionPattern(sessionId);
const keys = await redis.keys(pattern);
if (keys.length > 0) {
await redis.del(...keys);
}
},
exists: async (sessionId, key) => {
const redisKey = getKey(sessionId, key);
const exists = await redis.exists(redisKey);
return exists > 0;
}
};
};
// ========== PostgreSQL Storage ==========
export const createPostgresArtifactStorage = (config) => {
let pool;
const tableName = config.tableName || 'jaf_artifacts';
// const maxSize = config.maxSize || 10 * 1024 * 1024; // TODO: Implement size limit checking
const ttl = config.ttl;
try {
// eslint-disable-next-line @typescript-eslint/no-var-requires
const { Pool } = require('pg');
pool = new Pool({
connectionString: config.connectionString,
connectionTimeoutMillis: 5000, // 5 second timeout
query_timeout: 5000, // 5 second query timeout
statement_timeout: 5000 // 5 second statement timeout
});
}
catch (error) {
throw new Error('PostgreSQL artifact storage requires pg to be installed');
}
async function initializeTable() {
const query = `
CREATE TABLE IF NOT EXISTS ${tableName} (
session_id VARCHAR(255) NOT NULL,
key VARCHAR(255) NOT NULL,
value JSONB NOT NULL,
metadata JSONB NOT NULL,
created_at TIMESTAMP WITH TIME ZONE NOT NULL,
modified_at TIMESTAMP WITH TIME ZONE NOT NULL,
expires_at TIMESTAMP WITH TIME ZONE,
PRIMARY KEY (session_id, key)
);
CREATE INDEX IF NOT EXISTS idx_session_id ON ${tableName} (session_id);
CREATE INDEX IF NOT EXISTS idx_expires_at ON ${tableName} (expires_at);
`;
await pool.query(query);
}
// Initialize table asynchronously
const initPromise = initializeTable().catch(error => {
console.error('Failed to initialize PostgreSQL artifact storage table:', error);
throw error;
});
return {
get: async (sessionId, key) => {
await initPromise; // Ensure table exists
const query = `
SELECT * FROM ${tableName}
WHERE session_id = $1 AND key = $2
AND (expires_at IS NULL OR expires_at > NOW())
`;
const result = await pool.query(query, [sessionId, key]);
if (result.rows.length === 0)
return null;
const row = result.rows[0];
return {
key: row.key,
value: row.value,
metadata: {
...row.metadata,
created: new Date(row.created_at),
lastModified: new Date(row.modified_at)
}
};
},
set: async (sessionId, key, value, metadata) => {
await initPromise; // Ensure table exists
const now = new Date();
const expiresAt = ttl ? new Date(now.getTime() + ttl * 1000) : null;
const artifact = {
key,
value,
metadata: {
created: metadata?.created || now,
lastModified: now,
contentType: metadata?.contentType,
size: JSON.stringify(value).length,
tags: metadata?.tags
}
};
const query = `
INSERT INTO ${tableName} (session_id, key, value, metadata, created_at, modified_at, expires_at)
VALUES ($1, $2, $3, $4, $5, $6, $7)
ON CONFLICT (session_id, key)
DO UPDATE SET
value = $3,
metadata = $4,
modified_at = $6,
expires_at = $7
RETURNING *
`;
await pool.query(query, [
sessionId,
key,
JSON.stringify(value),
JSON.stringify(artifact.metadata),
artifact.metadata.created,
now,
expiresAt
]);
return artifact;
},
delete: async (sessionId, key) => {
await initPromise; // Ensure table exists
const query = `
DELETE FROM ${tableName}
WHERE session_id = $1 AND key = $2
`;
const result = await pool.query(query, [sessionId, key]);
return result.rowCount > 0;
},
list: async (sessionId) => {
await initPromise; // Ensure table exists
const query = `
SELECT * FROM ${tableName}
WHERE session_id = $1
AND (expires_at IS NULL OR expires_at > NOW())
ORDER BY modified_at DESC
`;
const result = await pool.query(query, [sessionId]);
return result.rows.map((row) => ({
key: row.key,
value: row.value,
metadata: {
...row.metadata,
created: new Date(row.created_at),
lastModified: new Date(row.modified_at)
}
}));
},
clear: async (sessionId) => {
await initPromise; // Ensure table exists
const query = `
DELETE FROM ${tableName}
WHERE session_id = $1
`;
await pool.query(query, [sessionId]);
},
exists: async (sessionId, key) => {
await initPromise; // Ensure table exists
const query = `
SELECT 1 FROM ${tableName}
WHERE session_id = $1 AND key = $2
AND (expires_at IS NULL OR expires_at > NOW())
LIMIT 1
`;
const result = await pool.query(query, [sessionId, key]);
return result.rows.length > 0;
}
};
};
// ========== Session Integration ==========
/**
* Helper functions to integrate artifacts with sessions
*/
export const getSessionArtifact = (session, key) => {
return session.artifacts[key] || null;
};
export const setSessionArtifact = (session, key, value) => {
return {
...session,
artifacts: {
...session.artifacts,
[key]: value
}
};
};
export const deleteSessionArtifact = (session, key) => {
const { [key]: _, ...rest } = session.artifacts;
return {
...session,
artifacts: rest
};
};
export const clearSessionArtifacts = (session) => {
return {
...session,
artifacts: {}
};
};
export const listSessionArtifacts = (session) => {
return Object.keys(session.artifacts);
};
// ========== Factory Function ==========
export const createArtifactStorage = (config) => {
switch (config.type) {
case 'memory':
return createMemoryArtifactStorage(config.config);
case 'redis':
return createRedisArtifactStorage(config.config);
case 'postgres':
return createPostgresArtifactStorage(config.config);
case 's3':
throw new Error('S3 artifact storage not yet implemented');
case 'gcs':
throw new Error('GCS artifact storage not yet implemented');
default:
throw new Error(`Unknown artifact storage type: ${config.type}`);
}
};
//# sourceMappingURL=index.js.map