feat: add Firestore vector search integration and noun type-based search enhancements
Introduced Firestore vector search integration using `@firebase/firestore-vector-search`. Added support for restricting searches by specific noun types to optimize performance and relevance. Updated `README.md` to document the new functionality and usage examples.
This commit is contained in:
parent
b10a32e8de
commit
725610fbe4
13 changed files with 2439 additions and 178 deletions
480
src/augmentations/firestoreStorageAugmentation.ts
Normal file
480
src/augmentations/firestoreStorageAugmentation.ts
Normal file
|
|
@ -0,0 +1,480 @@
|
|||
import {
|
||||
AugmentationType,
|
||||
IMemoryAugmentation,
|
||||
AugmentationResponse
|
||||
} from '../types/augmentations.js'
|
||||
import { Vector } from '../coreTypes.js'
|
||||
import { cosineDistance } from '../utils/distance.js'
|
||||
import { initializeApp, getApp } from 'firebase/app'
|
||||
import {
|
||||
getFirestore,
|
||||
collection,
|
||||
doc,
|
||||
setDoc,
|
||||
getDoc,
|
||||
updateDoc,
|
||||
deleteDoc,
|
||||
query,
|
||||
getDocs,
|
||||
limit
|
||||
} from 'firebase/firestore'
|
||||
import { findNearest } from '@firebase/firestore-vector-search'
|
||||
|
||||
/**
|
||||
* Configuration for Firestore storage augmentation
|
||||
*/
|
||||
export interface FirestoreStorageConfig {
|
||||
/**
|
||||
* Firestore project ID
|
||||
*/
|
||||
projectId: string
|
||||
|
||||
/**
|
||||
* Firestore collection name for storing data
|
||||
*/
|
||||
collection: string
|
||||
|
||||
/**
|
||||
* Optional Firestore credentials
|
||||
*/
|
||||
credentials?: any
|
||||
|
||||
/**
|
||||
* Optional Firestore database URL
|
||||
*/
|
||||
databaseURL?: string
|
||||
|
||||
/**
|
||||
* Optional Firestore app name
|
||||
*/
|
||||
appName?: string
|
||||
}
|
||||
|
||||
/**
|
||||
* Storage augmentation that uses Firestore for storage
|
||||
*/
|
||||
export class FirestoreStorageAugmentation implements IMemoryAugmentation {
|
||||
readonly name: string
|
||||
readonly description: string = 'Storage augmentation that stores data in Firestore'
|
||||
enabled: boolean = true
|
||||
private config: FirestoreStorageConfig
|
||||
private isInitialized = false
|
||||
private firestore: any = null
|
||||
private collection: any = null
|
||||
|
||||
constructor(name: string, config: FirestoreStorageConfig) {
|
||||
this.name = name
|
||||
this.config = config
|
||||
}
|
||||
|
||||
getType(): AugmentationType {
|
||||
return AugmentationType.MEMORY
|
||||
}
|
||||
|
||||
async initialize(): Promise<void> {
|
||||
if (this.isInitialized) {
|
||||
return
|
||||
}
|
||||
|
||||
try {
|
||||
// Initialize Firebase if not already initialized
|
||||
let app
|
||||
try {
|
||||
app = getApp(this.config.appName || '[DEFAULT]')
|
||||
} catch (e) {
|
||||
// App doesn't exist, initialize it
|
||||
// Create a config object with the required properties
|
||||
const firebaseConfig: any = {
|
||||
projectId: this.config.projectId
|
||||
}
|
||||
|
||||
// Add optional properties if they exist
|
||||
if (this.config.databaseURL) {
|
||||
firebaseConfig.databaseURL = this.config.databaseURL
|
||||
}
|
||||
|
||||
// Add credentials if they exist
|
||||
if (this.config.credentials) {
|
||||
// Use credentials as part of the config object
|
||||
// instead of as a 'credential' property
|
||||
Object.assign(firebaseConfig, this.config.credentials)
|
||||
}
|
||||
|
||||
app = initializeApp(firebaseConfig, this.config.appName)
|
||||
}
|
||||
|
||||
// Get Firestore instance
|
||||
this.firestore = getFirestore(app)
|
||||
this.collection = collection(this.firestore, this.config.collection)
|
||||
this.isInitialized = true
|
||||
} catch (error) {
|
||||
console.error(`Failed to initialize FirestoreStorageAugmentation:`, error)
|
||||
throw new Error(`Failed to initialize FirestoreStorageAugmentation: ${error}`)
|
||||
}
|
||||
}
|
||||
|
||||
async shutDown(): Promise<void> {
|
||||
this.isInitialized = false
|
||||
this.firestore = null
|
||||
this.collection = null
|
||||
}
|
||||
|
||||
async getStatus(): Promise<'active' | 'inactive' | 'error'> {
|
||||
if (!this.isInitialized) {
|
||||
return 'inactive'
|
||||
}
|
||||
|
||||
try {
|
||||
// Try a simple operation to check if Firestore is working
|
||||
await getDocs(query(collection(this.firestore, '__test__'), limit(1)))
|
||||
return 'active'
|
||||
} catch (error) {
|
||||
console.error('Firestore connection error:', error)
|
||||
return 'error'
|
||||
}
|
||||
}
|
||||
|
||||
async storeData(
|
||||
key: string,
|
||||
data: unknown,
|
||||
options?: Record<string, unknown>
|
||||
): Promise<AugmentationResponse<boolean>> {
|
||||
await this.ensureInitialized()
|
||||
|
||||
try {
|
||||
await setDoc(doc(this.collection, key), this.prepareForFirestore(data))
|
||||
return { success: true, data: true }
|
||||
} catch (error) {
|
||||
console.error(`Failed to store data for key ${key}:`, error)
|
||||
return {
|
||||
success: false,
|
||||
data: false,
|
||||
error: `Failed to store data: ${error}`
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async retrieveData(
|
||||
key: string,
|
||||
options?: Record<string, unknown>
|
||||
): Promise<AugmentationResponse<unknown>> {
|
||||
await this.ensureInitialized()
|
||||
|
||||
try {
|
||||
const docSnapshot = await getDoc(doc(this.collection, key))
|
||||
|
||||
if (!docSnapshot.exists()) {
|
||||
return {
|
||||
success: true,
|
||||
data: null
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
success: true,
|
||||
data: docSnapshot.data()
|
||||
}
|
||||
} catch (error) {
|
||||
console.error(`Failed to retrieve data for key ${key}:`, error)
|
||||
return {
|
||||
success: false,
|
||||
data: null,
|
||||
error: `Failed to retrieve data: ${error}`
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async updateData(
|
||||
key: string,
|
||||
data: unknown,
|
||||
options?: Record<string, unknown>
|
||||
): Promise<AugmentationResponse<boolean>> {
|
||||
await this.ensureInitialized()
|
||||
|
||||
try {
|
||||
await updateDoc(doc(this.collection, key), this.prepareForFirestore(data))
|
||||
return { success: true, data: true }
|
||||
} catch (error) {
|
||||
console.error(`Failed to update data for key ${key}:`, error)
|
||||
return {
|
||||
success: false,
|
||||
data: false,
|
||||
error: `Failed to update data: ${error}`
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async deleteData(
|
||||
key: string,
|
||||
options?: Record<string, unknown>
|
||||
): Promise<AugmentationResponse<boolean>> {
|
||||
await this.ensureInitialized()
|
||||
|
||||
try {
|
||||
await deleteDoc(doc(this.collection, key))
|
||||
return { success: true, data: true }
|
||||
} catch (error) {
|
||||
console.error(`Failed to delete data for key ${key}:`, error)
|
||||
return {
|
||||
success: false,
|
||||
data: false,
|
||||
error: `Failed to delete data: ${error}`
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async listDataKeys(
|
||||
pattern?: string,
|
||||
options?: Record<string, unknown>
|
||||
): Promise<AugmentationResponse<string[]>> {
|
||||
await this.ensureInitialized()
|
||||
|
||||
try {
|
||||
// Create a query from the collection
|
||||
const q = query(this.collection)
|
||||
|
||||
// If pattern is provided, use it to filter keys
|
||||
// Note: Firestore doesn't support wildcard queries directly,
|
||||
// so we'll need to do some client-side filtering
|
||||
const snapshot = await getDocs(q)
|
||||
|
||||
let keys = snapshot.docs.map((docSnapshot: any) => docSnapshot.id)
|
||||
|
||||
// Apply pattern filtering if provided
|
||||
if (pattern) {
|
||||
// Convert wildcard pattern to regex
|
||||
const regexPattern = new RegExp(
|
||||
'^' + pattern.replace(/\*/g, '.*') + '$'
|
||||
)
|
||||
keys = keys.filter((key: string) => regexPattern.test(key))
|
||||
}
|
||||
|
||||
return {
|
||||
success: true,
|
||||
data: keys
|
||||
}
|
||||
} catch (error) {
|
||||
console.error(`Failed to list data keys:`, error)
|
||||
return {
|
||||
success: false,
|
||||
data: [],
|
||||
error: `Failed to list data keys: ${error}`
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Searches for data in Firestore using vector similarity.
|
||||
* Uses Firestore's built-in findNearest function for efficient vector search.
|
||||
* @param queryData The query vector or data to search for
|
||||
* @param k Number of results to return (default: 10)
|
||||
* @param options Optional search options
|
||||
*/
|
||||
async search(
|
||||
queryData: unknown,
|
||||
k: number = 10,
|
||||
options?: Record<string, unknown>
|
||||
): Promise<AugmentationResponse<Array<{
|
||||
id: string;
|
||||
score: number;
|
||||
data: unknown;
|
||||
}>>> {
|
||||
await this.ensureInitialized()
|
||||
|
||||
try {
|
||||
// Check if queryData is a vector
|
||||
let queryVector: Vector
|
||||
|
||||
if (Array.isArray(queryData) && queryData.every(item => typeof item === 'number')) {
|
||||
queryVector = queryData as Vector
|
||||
} else {
|
||||
// If queryData is not a vector, we can't perform vector search
|
||||
return {
|
||||
success: false,
|
||||
data: [],
|
||||
error: 'Query must be a vector (array of numbers) for vector search'
|
||||
}
|
||||
}
|
||||
|
||||
// Get vector field name from options or use default 'vector'
|
||||
const vectorField = options?.vectorField as string || 'vector'
|
||||
|
||||
// Get distance measure from options or use default 'COSINE'
|
||||
// Firestore supports 'COSINE', 'EUCLIDEAN', and 'DOT_PRODUCT'
|
||||
const distanceMeasure = options?.distanceMeasure as string || 'COSINE'
|
||||
|
||||
try {
|
||||
// Note: In Firebase v9+, vector search requires the Firebase Extensions for Firestore Vector Search
|
||||
// This code attempts to use it if available, but will fall back to client-side search
|
||||
|
||||
// Use the vector search extension imported at the top of the file
|
||||
try {
|
||||
const vectorSearchOptions = {
|
||||
collection: this.collection,
|
||||
vectorField: vectorField,
|
||||
queryVector: queryVector,
|
||||
limit: k,
|
||||
distanceMeasure: distanceMeasure
|
||||
}
|
||||
|
||||
const searchResults = await findNearest(vectorSearchOptions)
|
||||
|
||||
// Process results
|
||||
const results: Array<{
|
||||
id: string;
|
||||
score: number;
|
||||
data: unknown;
|
||||
}> = searchResults.map((result: any) => {
|
||||
// Calculate the similarity score based on the distance measure
|
||||
let score: number
|
||||
|
||||
if (distanceMeasure === 'DOT_PRODUCT') {
|
||||
// For dot product, higher is already better
|
||||
score = result.distance || 0
|
||||
} else {
|
||||
// For COSINE and EUCLIDEAN, convert to similarity score
|
||||
// where 1 is perfect match and 0 is completely dissimilar
|
||||
score = 1 / (1 + (result.distance || 0))
|
||||
}
|
||||
|
||||
return {
|
||||
id: result.id,
|
||||
score,
|
||||
data: result.data
|
||||
}
|
||||
})
|
||||
|
||||
return {
|
||||
success: true,
|
||||
data: results
|
||||
}
|
||||
} catch (vectorSearchError) {
|
||||
// Vector search extension not available, fall back to client-side search
|
||||
console.warn('Firestore vector search extension not available, falling back to client-side search:', vectorSearchError)
|
||||
throw vectorSearchError
|
||||
}
|
||||
} catch (vectorSearchError) {
|
||||
console.warn('Firestore vector search failed, falling back to client-side search:', vectorSearchError)
|
||||
|
||||
// Fallback to client-side search if findNearest is not available
|
||||
// This can happen if the Firestore instance doesn't support vector search
|
||||
// or if the collection isn't configured for vector search
|
||||
|
||||
// Get all documents from Firestore using the modern API
|
||||
const q = query(this.collection)
|
||||
const snapshot = await getDocs(q)
|
||||
|
||||
// Calculate distances and prepare results
|
||||
const results: Array<{
|
||||
id: string;
|
||||
score: number;
|
||||
data: unknown;
|
||||
}> = []
|
||||
|
||||
for (const docSnapshot of snapshot.docs) {
|
||||
const data = docSnapshot.data() as Record<string, any>
|
||||
|
||||
// Skip documents that don't have a vector field
|
||||
if (!data[vectorField] || !Array.isArray(data[vectorField])) {
|
||||
continue
|
||||
}
|
||||
|
||||
// Calculate distance between query vector and document vector
|
||||
const distance = cosineDistance(queryVector, data[vectorField] as number[])
|
||||
|
||||
// Convert distance to similarity score (1 - distance for cosine)
|
||||
// This way higher scores are better (more similar)
|
||||
const score = 1 - distance
|
||||
|
||||
results.push({
|
||||
id: docSnapshot.id,
|
||||
score,
|
||||
data
|
||||
})
|
||||
}
|
||||
|
||||
// Sort results by score (descending) and take top k
|
||||
results.sort((a, b) => b.score - a.score)
|
||||
const topResults = results.slice(0, k)
|
||||
|
||||
return {
|
||||
success: true,
|
||||
data: topResults
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
console.error(`Failed to search in Firestore:`, error)
|
||||
return {
|
||||
success: false,
|
||||
data: [],
|
||||
error: `Failed to search in Firestore: ${error}`
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async ensureInitialized(): Promise<void> {
|
||||
if (!this.isInitialized) {
|
||||
await this.initialize()
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Prepare data for Firestore by handling special types
|
||||
*/
|
||||
private prepareForFirestore(obj: any): any {
|
||||
if (obj === null || obj === undefined) {
|
||||
return null
|
||||
}
|
||||
|
||||
if (
|
||||
typeof obj === 'string' ||
|
||||
typeof obj === 'number' ||
|
||||
typeof obj === 'boolean'
|
||||
) {
|
||||
return obj
|
||||
}
|
||||
|
||||
if (obj instanceof Date) {
|
||||
return obj
|
||||
}
|
||||
|
||||
if (Array.isArray(obj)) {
|
||||
return obj.map(item => this.prepareForFirestore(item))
|
||||
}
|
||||
|
||||
if (obj instanceof Map) {
|
||||
const result: Record<string, any> = {}
|
||||
for (const [key, value] of obj.entries()) {
|
||||
result[String(key)] = this.prepareForFirestore(value)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
if (obj instanceof Set) {
|
||||
return Array.from(obj).map(item => this.prepareForFirestore(item))
|
||||
}
|
||||
|
||||
if (typeof obj === 'object') {
|
||||
const result: Record<string, any> = {}
|
||||
for (const key in obj) {
|
||||
if (Object.prototype.hasOwnProperty.call(obj, key)) {
|
||||
result[key] = this.prepareForFirestore(obj[key])
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// Default fallback
|
||||
return null
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a Firestore storage augmentation
|
||||
*/
|
||||
export function createFirestoreStorageAugmentation(
|
||||
name: string,
|
||||
config: FirestoreStorageConfig
|
||||
): IMemoryAugmentation {
|
||||
return new FirestoreStorageAugmentation(name, config)
|
||||
}
|
||||
341
src/augmentations/memoryAugmentations.ts
Normal file
341
src/augmentations/memoryAugmentations.ts
Normal file
|
|
@ -0,0 +1,341 @@
|
|||
import {
|
||||
AugmentationType,
|
||||
IMemoryAugmentation,
|
||||
AugmentationResponse
|
||||
} from '../types/augmentations.js'
|
||||
import { StorageAdapter, Vector } from '../coreTypes.js'
|
||||
import { MemoryStorage } from '../storage/opfsStorage.js'
|
||||
import { FileSystemStorage } from '../storage/fileSystemStorage.js'
|
||||
import { OPFSStorage } from '../storage/opfsStorage.js'
|
||||
import { cosineDistance } from '../utils/distance.js'
|
||||
import {
|
||||
FirestoreStorageAugmentation,
|
||||
FirestoreStorageConfig,
|
||||
createFirestoreStorageAugmentation
|
||||
} from './firestoreStorageAugmentation.js'
|
||||
|
||||
/**
|
||||
* Base class for memory augmentations that wrap a StorageAdapter
|
||||
*/
|
||||
abstract class BaseMemoryAugmentation implements IMemoryAugmentation {
|
||||
readonly name: string
|
||||
readonly description: string = 'Base memory augmentation'
|
||||
enabled: boolean = true
|
||||
protected storage: StorageAdapter
|
||||
protected isInitialized = false
|
||||
|
||||
constructor(name: string, storage: StorageAdapter) {
|
||||
this.name = name
|
||||
this.storage = storage
|
||||
}
|
||||
|
||||
async initialize(): Promise<void> {
|
||||
if (this.isInitialized) {
|
||||
return
|
||||
}
|
||||
|
||||
try {
|
||||
await this.storage.init()
|
||||
this.isInitialized = true
|
||||
} catch (error) {
|
||||
console.error(`Failed to initialize ${this.name}:`, error)
|
||||
throw new Error(`Failed to initialize ${this.name}: ${error}`)
|
||||
}
|
||||
}
|
||||
|
||||
async shutDown(): Promise<void> {
|
||||
this.isInitialized = false
|
||||
}
|
||||
|
||||
async getStatus(): Promise<'active' | 'inactive' | 'error'> {
|
||||
return this.isInitialized ? 'active' : 'inactive'
|
||||
}
|
||||
|
||||
async storeData(
|
||||
key: string,
|
||||
data: unknown,
|
||||
options?: Record<string, unknown>
|
||||
): Promise<AugmentationResponse<boolean>> {
|
||||
await this.ensureInitialized()
|
||||
|
||||
try {
|
||||
await this.storage.saveMetadata(key, data)
|
||||
return { success: true, data: true }
|
||||
} catch (error) {
|
||||
console.error(`Failed to store data for key ${key}:`, error)
|
||||
return {
|
||||
success: false,
|
||||
data: false,
|
||||
error: `Failed to store data: ${error}`
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async retrieveData(
|
||||
key: string,
|
||||
options?: Record<string, unknown>
|
||||
): Promise<AugmentationResponse<unknown>> {
|
||||
await this.ensureInitialized()
|
||||
|
||||
try {
|
||||
const data = await this.storage.getMetadata(key)
|
||||
return {
|
||||
success: true,
|
||||
data
|
||||
}
|
||||
} catch (error) {
|
||||
console.error(`Failed to retrieve data for key ${key}:`, error)
|
||||
return {
|
||||
success: false,
|
||||
data: null,
|
||||
error: `Failed to retrieve data: ${error}`
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async updateData(
|
||||
key: string,
|
||||
data: unknown,
|
||||
options?: Record<string, unknown>
|
||||
): Promise<AugmentationResponse<boolean>> {
|
||||
await this.ensureInitialized()
|
||||
|
||||
try {
|
||||
await this.storage.saveMetadata(key, data)
|
||||
return { success: true, data: true }
|
||||
} catch (error) {
|
||||
console.error(`Failed to update data for key ${key}:`, error)
|
||||
return {
|
||||
success: false,
|
||||
data: false,
|
||||
error: `Failed to update data: ${error}`
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async deleteData(
|
||||
key: string,
|
||||
options?: Record<string, unknown>
|
||||
): Promise<AugmentationResponse<boolean>> {
|
||||
await this.ensureInitialized()
|
||||
|
||||
try {
|
||||
// There's no direct deleteMetadata method, so we save null
|
||||
await this.storage.saveMetadata(key, null)
|
||||
return { success: true, data: true }
|
||||
} catch (error) {
|
||||
console.error(`Failed to delete data for key ${key}:`, error)
|
||||
return {
|
||||
success: false,
|
||||
data: false,
|
||||
error: `Failed to delete data: ${error}`
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async listDataKeys(
|
||||
pattern?: string,
|
||||
options?: Record<string, unknown>
|
||||
): Promise<AugmentationResponse<string[]>> {
|
||||
// This is a limitation of the current StorageAdapter interface
|
||||
// It doesn't provide a way to list all metadata keys
|
||||
// We could implement this in the future by extending the StorageAdapter interface
|
||||
return {
|
||||
success: false,
|
||||
data: [],
|
||||
error: 'listDataKeys is not supported by this storage adapter'
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Searches for data in the storage using vector similarity.
|
||||
* Implements the findNearest functionality by calculating distances client-side.
|
||||
* @param query The query vector or data to search for
|
||||
* @param k Number of results to return (default: 10)
|
||||
* @param options Optional search options
|
||||
*/
|
||||
async search(
|
||||
query: unknown,
|
||||
k: number = 10,
|
||||
options?: Record<string, unknown>
|
||||
): Promise<AugmentationResponse<Array<{
|
||||
id: string;
|
||||
score: number;
|
||||
data: unknown;
|
||||
}>>> {
|
||||
await this.ensureInitialized()
|
||||
|
||||
try {
|
||||
// Check if query is a vector
|
||||
let queryVector: Vector
|
||||
|
||||
if (Array.isArray(query) && query.every(item => typeof item === 'number')) {
|
||||
queryVector = query as Vector
|
||||
} else {
|
||||
// If query is not a vector, we can't perform vector search
|
||||
return {
|
||||
success: false,
|
||||
data: [],
|
||||
error: 'Query must be a vector (array of numbers) for vector search'
|
||||
}
|
||||
}
|
||||
|
||||
// Get all nodes from storage
|
||||
const nodes = await this.storage.getAllNodes()
|
||||
|
||||
// Calculate distances and prepare results
|
||||
const results: Array<{
|
||||
id: string;
|
||||
score: number;
|
||||
data: unknown;
|
||||
}> = []
|
||||
|
||||
for (const node of nodes) {
|
||||
// Skip nodes that don't have a vector
|
||||
if (!node.vector || !Array.isArray(node.vector)) {
|
||||
continue
|
||||
}
|
||||
|
||||
// Get metadata for the node
|
||||
const metadata = await this.storage.getMetadata(node.id)
|
||||
|
||||
// Calculate distance between query vector and node vector
|
||||
const distance = cosineDistance(queryVector, node.vector)
|
||||
|
||||
// Convert distance to similarity score (1 - distance for cosine)
|
||||
// This way higher scores are better (more similar)
|
||||
const score = 1 - distance
|
||||
|
||||
results.push({
|
||||
id: node.id,
|
||||
score,
|
||||
data: metadata
|
||||
})
|
||||
}
|
||||
|
||||
// Sort results by score (descending) and take top k
|
||||
results.sort((a, b) => b.score - a.score)
|
||||
const topResults = results.slice(0, k)
|
||||
|
||||
return {
|
||||
success: true,
|
||||
data: topResults
|
||||
}
|
||||
} catch (error) {
|
||||
console.error(`Failed to search in storage:`, error)
|
||||
return {
|
||||
success: false,
|
||||
data: [],
|
||||
error: `Failed to search in storage: ${error}`
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
protected async ensureInitialized(): Promise<void> {
|
||||
if (!this.isInitialized) {
|
||||
await this.initialize()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Memory augmentation that uses in-memory storage
|
||||
*/
|
||||
export class MemoryStorageAugmentation extends BaseMemoryAugmentation {
|
||||
readonly description = 'Memory augmentation that stores data in memory'
|
||||
enabled = true
|
||||
|
||||
constructor(name: string) {
|
||||
super(name, new MemoryStorage())
|
||||
}
|
||||
|
||||
getType(): AugmentationType {
|
||||
return AugmentationType.MEMORY
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Memory augmentation that uses file system storage
|
||||
*/
|
||||
export class FileSystemStorageAugmentation extends BaseMemoryAugmentation {
|
||||
readonly description = 'Memory augmentation that stores data in the file system'
|
||||
enabled = true
|
||||
|
||||
constructor(name: string, rootDirectory?: string) {
|
||||
super(name, new FileSystemStorage(rootDirectory))
|
||||
}
|
||||
|
||||
getType(): AugmentationType {
|
||||
return AugmentationType.MEMORY
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Memory augmentation that uses OPFS (Origin Private File System) storage
|
||||
*/
|
||||
export class OPFSStorageAugmentation extends BaseMemoryAugmentation {
|
||||
readonly description = 'Memory augmentation that stores data in the Origin Private File System'
|
||||
enabled = true
|
||||
|
||||
constructor(name: string) {
|
||||
super(name, new OPFSStorage())
|
||||
}
|
||||
|
||||
getType(): AugmentationType {
|
||||
return AugmentationType.MEMORY
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Factory function to create the appropriate memory augmentation based on the environment
|
||||
*/
|
||||
export async function createMemoryAugmentation(
|
||||
name: string,
|
||||
options: {
|
||||
storageType?: 'memory' | 'filesystem' | 'opfs' | 'firestore'
|
||||
rootDirectory?: string
|
||||
requestPersistentStorage?: boolean
|
||||
firestoreConfig?: FirestoreStorageConfig
|
||||
} = {}
|
||||
): Promise<IMemoryAugmentation> {
|
||||
// If a specific storage type is requested, use that
|
||||
if (options.storageType) {
|
||||
switch (options.storageType) {
|
||||
case 'memory':
|
||||
return new MemoryStorageAugmentation(name)
|
||||
case 'filesystem':
|
||||
return new FileSystemStorageAugmentation(name, options.rootDirectory)
|
||||
case 'opfs':
|
||||
return new OPFSStorageAugmentation(name)
|
||||
case 'firestore':
|
||||
if (!options.firestoreConfig) {
|
||||
throw new Error('firestoreConfig is required when using Firestore storage')
|
||||
}
|
||||
return createFirestoreStorageAugmentation(name, options.firestoreConfig)
|
||||
}
|
||||
}
|
||||
|
||||
// Otherwise, select based on environment
|
||||
const isNode = typeof process !== 'undefined' &&
|
||||
process.versions != null &&
|
||||
process.versions.node != null
|
||||
|
||||
if (isNode) {
|
||||
// In Node.js, use FileSystemStorage
|
||||
return new FileSystemStorageAugmentation(name, options.rootDirectory)
|
||||
} else {
|
||||
// In browser, try OPFS first
|
||||
const opfsStorage = new OPFSStorage()
|
||||
|
||||
if (opfsStorage.isOPFSAvailable()) {
|
||||
// Request persistent storage if specified
|
||||
if (options.requestPersistentStorage) {
|
||||
await opfsStorage.requestPersistentStorage()
|
||||
}
|
||||
return new OPFSStorageAugmentation(name)
|
||||
} else {
|
||||
// Fall back to memory storage
|
||||
return new MemoryStorageAugmentation(name)
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue