David Snelling
487bfa473b
### Changes: - Introduced `ExtendedBrainyDataConfig` interface to include `rootDirectory` in the storage configuration. - Enhanced S3 storage setup: - Added support for `customS3Storage` when `S3_ENDPOINT` is provided. - Retained standard `s3Storage` configuration for AWS S3 without custom endpoints. - Updated fallback logic for default storage types (`filesystem`, `memory`) with consistent handling for `rootDirectory`. - Refined `initializeBrainy` function to improve readability and configuration extensibility. - Updated TypeScript typings for improved clarity and maintainability. ### Purpose: Improved the `initializeBrainy` function by extending storage configuration capabilities, allowing seamless integration of custom S3 endpoints alongside standard AWS S3 settings. Enhanced overall code readability and adaptability for diverse storage configurations.
97 lines
2.8 KiB
TypeScript
97 lines
2.8 KiB
TypeScript
import { BrainyData, BrainyDataConfig } from '@soulcraft/brainy'
|
|
import dotenv from 'dotenv'
|
|
|
|
// Load environment variables
|
|
dotenv.config()
|
|
|
|
/**
|
|
* Extended configuration interface that includes rootDirectory
|
|
*/
|
|
interface ExtendedBrainyDataConfig extends BrainyDataConfig {
|
|
storage: BrainyDataConfig['storage'] & {
|
|
rootDirectory?: string
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Initialize the Brainy instance with appropriate configuration
|
|
* based on environment variables
|
|
*/
|
|
export async function initializeBrainy(): Promise<BrainyData> {
|
|
// Get storage configuration from environment variables
|
|
const storageType = process.env.STORAGE_TYPE || 'filesystem'
|
|
|
|
// Create configuration object
|
|
const config: ExtendedBrainyDataConfig = {
|
|
storage: {}
|
|
}
|
|
|
|
// Configure storage based on type
|
|
switch (storageType) {
|
|
case 's3':
|
|
// Configure S3 storage
|
|
if (process.env.S3_ENDPOINT) {
|
|
// Use customS3Storage for S3-compatible services with custom endpoints
|
|
config.storage.customS3Storage = {
|
|
bucketName: process.env.S3_BUCKET_NAME,
|
|
accessKeyId: process.env.S3_ACCESS_KEY_ID,
|
|
secretAccessKey: process.env.S3_SECRET_ACCESS_KEY,
|
|
region: process.env.S3_REGION || 'us-east-1',
|
|
endpoint: process.env.S3_ENDPOINT
|
|
}
|
|
} else {
|
|
// Use standard S3 storage for AWS S3
|
|
config.storage.s3Storage = {
|
|
bucketName: process.env.S3_BUCKET_NAME,
|
|
accessKeyId: process.env.S3_ACCESS_KEY_ID,
|
|
secretAccessKey: process.env.S3_SECRET_ACCESS_KEY,
|
|
region: process.env.S3_REGION || 'us-east-1'
|
|
}
|
|
}
|
|
break
|
|
|
|
case 'filesystem':
|
|
// Configure filesystem storage
|
|
config.storage.rootDirectory = process.env.STORAGE_ROOT_DIR || './data'
|
|
break
|
|
|
|
case 'memory':
|
|
// No additional configuration needed for memory storage
|
|
break
|
|
|
|
default:
|
|
// Default to filesystem storage
|
|
config.storage.rootDirectory = process.env.STORAGE_ROOT_DIR || './data'
|
|
break
|
|
}
|
|
|
|
// Note: Universal Sentence Encoder is now the only embedding option
|
|
// TensorFlow.js is required for embedding to work
|
|
|
|
// Configure HNSW index parameters if provided
|
|
if (process.env.HNSW_M) {
|
|
config.hnsw = {
|
|
M: parseInt(process.env.HNSW_M, 10),
|
|
efConstruction: process.env.HNSW_EF_CONSTRUCTION
|
|
? parseInt(process.env.HNSW_EF_CONSTRUCTION, 10)
|
|
: 200,
|
|
efSearch: process.env.HNSW_EF_SEARCH
|
|
? parseInt(process.env.HNSW_EF_SEARCH, 10)
|
|
: 50
|
|
}
|
|
}
|
|
|
|
// Create and initialize the Brainy instance
|
|
const brainy = new BrainyData(config)
|
|
await brainy.init()
|
|
|
|
return brainy
|
|
}
|
|
|
|
/**
|
|
* Get the current storage type being used by Brainy
|
|
*/
|
|
export async function getStorageType(brainy: BrainyData): Promise<string> {
|
|
const status = await brainy.status()
|
|
return status.type || 'unknown'
|
|
}
|