376 lines
14 KiB
TypeScript
376 lines
14 KiB
TypeScript
import { z } from 'zod'
|
|
import { StructuredTool } from '@langchain/core/tools'
|
|
import { ICommonObject, INode, INodeData, INodeOptionsValue, INodeParams } from '../../../src/Interface'
|
|
import { getBaseClasses } from '../../../src/utils'
|
|
import { AWS_REGIONS, DEFAULT_AWS_REGION, AWSCredentials, getAWSCredentials } from '../../../src/awsToolsUtils'
|
|
import { DynamoDBClient, DescribeTableCommand, ListTablesCommand, PutItemCommand, QueryCommand } from '@aws-sdk/client-dynamodb'
|
|
|
|
// Operation enum
|
|
enum Operation {
|
|
STORE = 'store',
|
|
RETRIEVE = 'retrieve'
|
|
}
|
|
|
|
// Constants
|
|
const ERROR_PLACEHOLDER = 'error'
|
|
const KEY_SEPARATOR = '#'
|
|
const MAX_KEY_LENGTH = 2048 // DynamoDB limit for partition key
|
|
|
|
// Helper function to create DynamoDB client
|
|
function createDynamoDBClient(credentials: AWSCredentials, region: string): DynamoDBClient {
|
|
return new DynamoDBClient({
|
|
region,
|
|
credentials: {
|
|
accessKeyId: credentials.accessKeyId,
|
|
secretAccessKey: credentials.secretAccessKey,
|
|
...(credentials.sessionToken && { sessionToken: credentials.sessionToken })
|
|
}
|
|
})
|
|
}
|
|
|
|
// Helper function to build full key with optional prefix
|
|
function buildFullKey(key: string, keyPrefix: string): string {
|
|
const fullKey = keyPrefix ? `${keyPrefix}${KEY_SEPARATOR}${key}` : key
|
|
|
|
// Validate key length (DynamoDB limit)
|
|
if (fullKey.length > MAX_KEY_LENGTH) {
|
|
throw new Error(`Key too long. Maximum length is ${MAX_KEY_LENGTH} characters, got ${fullKey.length}`)
|
|
}
|
|
|
|
return fullKey
|
|
}
|
|
|
|
// Helper function to validate and sanitize input
|
|
function validateKey(key: string): void {
|
|
if (!key || key.trim().length === 0) {
|
|
throw new Error('Key must be a non-empty string')
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Tool for storing key-value pairs in DynamoDB with automatic versioning
|
|
*/
|
|
class DynamoDBStoreTool extends StructuredTool {
|
|
name = 'dynamodb_kv_store'
|
|
description = 'Store a text value with a key in DynamoDB. Input must be an object with "key" and "value" properties.'
|
|
schema = z.object({
|
|
key: z.string().min(1).describe('The key to store the value under'),
|
|
value: z.string().describe('The text value to store')
|
|
})
|
|
private readonly dynamoClient: DynamoDBClient
|
|
private readonly tableName: string
|
|
private readonly keyPrefix: string
|
|
|
|
constructor(dynamoClient: DynamoDBClient, tableName: string, keyPrefix: string = '') {
|
|
super()
|
|
this.dynamoClient = dynamoClient
|
|
this.tableName = tableName
|
|
this.keyPrefix = keyPrefix
|
|
}
|
|
|
|
async _call({ key, value }: z.infer<typeof this.schema>): Promise<string> {
|
|
try {
|
|
validateKey(key)
|
|
const fullKey = buildFullKey(key, this.keyPrefix)
|
|
const timestamp = Date.now()
|
|
const isoTimestamp = new Date(timestamp).toISOString()
|
|
|
|
const putCommand = new PutItemCommand({
|
|
TableName: this.tableName,
|
|
Item: {
|
|
pk: { S: fullKey },
|
|
sk: { S: timestamp.toString() },
|
|
value: { S: value },
|
|
timestamp: { S: isoTimestamp }
|
|
}
|
|
})
|
|
|
|
await this.dynamoClient.send(putCommand)
|
|
return `Successfully stored value with key "${key}" at ${isoTimestamp}`
|
|
} catch (error) {
|
|
const errorMessage = error instanceof Error ? error.message : String(error)
|
|
throw new Error(`Failed to store value: ${errorMessage}`)
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Tool for retrieving key-value pairs from DynamoDB with version control
|
|
*/
|
|
class DynamoDBRetrieveTool extends StructuredTool {
|
|
name = 'dynamodb_kv_retrieve'
|
|
description =
|
|
'Retrieve a value by key from DynamoDB. Returns JSON with value and timestamp. Specify which version to get (1=latest, 2=2nd latest, etc).'
|
|
schema = z.object({
|
|
key: z.string().min(1).describe('The key to retrieve the value for'),
|
|
nthLatest: z
|
|
.string()
|
|
.regex(/^\d+$/, 'Must be a positive number')
|
|
.describe(
|
|
'Which version to retrieve: "1" for latest, "2" for 2nd latest, "3" for 3rd latest, etc. Use "1" to get the most recent value.'
|
|
)
|
|
.optional()
|
|
.default('1')
|
|
})
|
|
private readonly dynamoClient: DynamoDBClient
|
|
private readonly tableName: string
|
|
private readonly keyPrefix: string
|
|
|
|
constructor(dynamoClient: DynamoDBClient, tableName: string, keyPrefix: string = '') {
|
|
super()
|
|
this.dynamoClient = dynamoClient
|
|
this.tableName = tableName
|
|
this.keyPrefix = keyPrefix
|
|
}
|
|
|
|
async _call(input: z.infer<typeof this.schema>): Promise<string> {
|
|
try {
|
|
const { key, nthLatest = '1' } = input
|
|
validateKey(key)
|
|
const fullKey = buildFullKey(key, this.keyPrefix)
|
|
|
|
// Convert string to number and validate
|
|
const nthLatestNum = parseInt(nthLatest, 10)
|
|
if (isNaN(nthLatestNum) || nthLatestNum < 1) {
|
|
throw new Error('nthLatest must be a positive number (1 or greater)')
|
|
}
|
|
|
|
const queryCommand = new QueryCommand({
|
|
TableName: this.tableName,
|
|
KeyConditionExpression: 'pk = :pk',
|
|
ExpressionAttributeValues: {
|
|
':pk': { S: fullKey }
|
|
},
|
|
ScanIndexForward: false, // Sort descending (newest first)
|
|
Limit: nthLatestNum
|
|
})
|
|
|
|
const result = await this.dynamoClient.send(queryCommand)
|
|
|
|
if (!result.Items || result.Items.length === 0) {
|
|
return JSON.stringify({
|
|
value: null,
|
|
timestamp: null
|
|
})
|
|
}
|
|
|
|
if (result.Items.length < nthLatestNum) {
|
|
return JSON.stringify({
|
|
value: null,
|
|
timestamp: null
|
|
})
|
|
}
|
|
|
|
const item = result.Items[nthLatestNum - 1]
|
|
const value = item.value?.S || null
|
|
const timestamp = item.timestamp?.S || item.sk?.S || null
|
|
|
|
// Return JSON with value and timestamp
|
|
return JSON.stringify({
|
|
value: value,
|
|
timestamp: timestamp
|
|
})
|
|
} catch (error) {
|
|
const errorMessage = error instanceof Error ? error.message : String(error)
|
|
throw new Error(`Failed to retrieve value: ${errorMessage}`)
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Node implementation for AWS DynamoDB KV Storage tools
|
|
*/
|
|
class AWSDynamoDBKVStorage_Tools implements INode {
|
|
label: string
|
|
name: string
|
|
version: number
|
|
type: string
|
|
icon: string
|
|
category: string
|
|
description: string
|
|
baseClasses: string[]
|
|
credential: INodeParams
|
|
inputs: INodeParams[]
|
|
|
|
constructor() {
|
|
this.label = 'AWS DynamoDB KV Storage'
|
|
this.name = 'awsDynamoDBKVStorage'
|
|
this.version = 1.0
|
|
this.type = 'AWSDynamoDBKVStorage'
|
|
this.icon = 'dynamodbkvstorage.svg'
|
|
this.category = 'Tools'
|
|
this.description = 'Store and retrieve versioned text values in AWS DynamoDB'
|
|
this.baseClasses = [this.type, ...getBaseClasses(DynamoDBStoreTool)]
|
|
this.credential = {
|
|
label: 'AWS Credentials',
|
|
name: 'credential',
|
|
type: 'credential',
|
|
credentialNames: ['awsApi']
|
|
}
|
|
this.inputs = [
|
|
{
|
|
label: 'AWS Region',
|
|
name: 'region',
|
|
type: 'options',
|
|
options: AWS_REGIONS,
|
|
default: DEFAULT_AWS_REGION,
|
|
description: 'AWS Region where your DynamoDB tables are located'
|
|
},
|
|
{
|
|
label: 'DynamoDB Table',
|
|
name: 'tableName',
|
|
type: 'asyncOptions',
|
|
loadMethod: 'listTables',
|
|
description: 'Select a DynamoDB table with partition key "pk" and sort key "sk"',
|
|
refresh: true
|
|
},
|
|
{
|
|
label: 'Key Prefix',
|
|
name: 'keyPrefix',
|
|
type: 'string',
|
|
description: 'Optional prefix to add to all keys (e.g., "myapp" would make keys like "myapp#userdata")',
|
|
optional: true,
|
|
additionalParams: true
|
|
},
|
|
{
|
|
label: 'Operation',
|
|
name: 'operation',
|
|
type: 'options',
|
|
options: [
|
|
{ label: 'Store', name: Operation.STORE },
|
|
{ label: 'Retrieve', name: Operation.RETRIEVE }
|
|
],
|
|
default: Operation.STORE,
|
|
description: 'Choose whether to store or retrieve data'
|
|
}
|
|
]
|
|
}
|
|
|
|
loadMethods: Record<string, (nodeData: INodeData, options?: ICommonObject) => Promise<INodeOptionsValue[]>> = {
|
|
listTables: async (nodeData: INodeData, options?: ICommonObject): Promise<INodeOptionsValue[]> => {
|
|
try {
|
|
const credentials = await getAWSCredentials(nodeData, options ?? {})
|
|
const region = (nodeData.inputs?.region as string) || DEFAULT_AWS_REGION
|
|
const dynamoClient = createDynamoDBClient(credentials, region)
|
|
|
|
const listCommand = new ListTablesCommand({})
|
|
const listResponse = await dynamoClient.send(listCommand)
|
|
|
|
if (!listResponse.TableNames || listResponse.TableNames.length === 0) {
|
|
return [
|
|
{
|
|
label: 'No tables found',
|
|
name: ERROR_PLACEHOLDER,
|
|
description: 'No DynamoDB tables found in this region'
|
|
}
|
|
]
|
|
}
|
|
|
|
const validTables: INodeOptionsValue[] = []
|
|
const invalidTables: string[] = []
|
|
|
|
// Check tables in parallel for better performance
|
|
const tableChecks = await Promise.allSettled(
|
|
listResponse.TableNames.map(async (tableName) => {
|
|
const describeCommand = new DescribeTableCommand({
|
|
TableName: tableName
|
|
})
|
|
const describeResponse = await dynamoClient.send(describeCommand)
|
|
|
|
const keySchema = describeResponse.Table?.KeySchema
|
|
if (keySchema) {
|
|
const hasPk = keySchema.some((key) => key.AttributeName === 'pk' && key.KeyType === 'HASH')
|
|
const hasSk = keySchema.some((key) => key.AttributeName === 'sk' && key.KeyType === 'RANGE')
|
|
|
|
if (hasPk && hasSk) {
|
|
return {
|
|
valid: true,
|
|
table: {
|
|
label: tableName,
|
|
name: tableName,
|
|
description: `Table with pk (partition) and sk (sort) keys`
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return { valid: false, tableName }
|
|
})
|
|
)
|
|
|
|
tableChecks.forEach((result) => {
|
|
if (result.status === 'fulfilled') {
|
|
if (result.value.valid) {
|
|
validTables.push(result.value.table!)
|
|
} else if (result.value.tableName) {
|
|
invalidTables.push(result.value.tableName)
|
|
}
|
|
}
|
|
})
|
|
|
|
if (validTables.length === 0) {
|
|
return [
|
|
{
|
|
label: 'No compatible tables found',
|
|
name: ERROR_PLACEHOLDER,
|
|
description: `No tables with partition key "pk" and sort key "sk" found. ${
|
|
invalidTables.length > 0 ? `Found ${invalidTables.length} table(s) with different schema.` : ''
|
|
} Please create a table with these keys.`
|
|
}
|
|
]
|
|
}
|
|
|
|
// Sort tables alphabetically
|
|
validTables.sort((a, b) => a.label.localeCompare(b.label))
|
|
|
|
return validTables
|
|
} catch (error) {
|
|
if (error instanceof Error && error.message.includes('AWS Access Key')) {
|
|
return [
|
|
{
|
|
label: 'AWS Credentials Required',
|
|
name: ERROR_PLACEHOLDER,
|
|
description: 'Enter AWS Access Key ID and Secret Access Key'
|
|
}
|
|
]
|
|
}
|
|
console.error('Error loading DynamoDB tables:', error)
|
|
return [
|
|
{
|
|
label: 'Error Loading Tables',
|
|
name: ERROR_PLACEHOLDER,
|
|
description: `Failed to load tables: ${error instanceof Error ? error.message : String(error)}`
|
|
}
|
|
]
|
|
}
|
|
}
|
|
}
|
|
|
|
async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
|
|
const credentials = await getAWSCredentials(nodeData, options)
|
|
|
|
const region = (nodeData.inputs?.region as string) || DEFAULT_AWS_REGION
|
|
const tableName = nodeData.inputs?.tableName as string
|
|
const keyPrefix = (nodeData.inputs?.keyPrefix as string) || ''
|
|
const operation = (nodeData.inputs?.operation as string) || Operation.STORE
|
|
|
|
if (!tableName || tableName === ERROR_PLACEHOLDER) {
|
|
throw new Error('Valid DynamoDB Table selection is required')
|
|
}
|
|
|
|
// Validate key prefix doesn't contain separator
|
|
if (keyPrefix && keyPrefix.includes(KEY_SEPARATOR)) {
|
|
throw new Error(`Key prefix cannot contain "${KEY_SEPARATOR}" character`)
|
|
}
|
|
|
|
const dynamoClient = createDynamoDBClient(credentials, region)
|
|
|
|
if (operation === Operation.STORE) {
|
|
return new DynamoDBStoreTool(dynamoClient, tableName, keyPrefix)
|
|
} else {
|
|
return new DynamoDBRetrieveTool(dynamoClient, tableName, keyPrefix)
|
|
}
|
|
}
|
|
}
|
|
|
|
module.exports = { nodeClass: AWSDynamoDBKVStorage_Tools }
|