aimemory-core 1.0.1

package/src/rateLimiter.ts

@@ -0,0 +1,186 @@
+export interface RateLimitConfig {
+  windowMs: number;
+  maxRequests: number;
+  keyPrefix?: string;
+}
+
+export interface RateLimitResult {
+  allowed: boolean;
+  remaining: number;
+  resetTime: number;
+  retryAfter?: number;
+}
+
+export interface RateLimitStore {
+  get(key: string): Promise<{ count: number; resetTime: number } | null>;
+  set(key: string, count: number, windowMs: number, resetTime: number): Promise<void>;
+  increment(key: string): Promise<{ count: number; resetTime: number }>;
+  delete(key: string): Promise<void>;
+  reset(): Promise<void>;
+}
+
+export class InMemoryRateLimitStore implements RateLimitStore {
+  private store: Map<string, { count: number; resetTime: number }> = new Map();
+
+  async get(key: string): Promise<{ count: number; resetTime: number } | null> {
+    const entry = this.store.get(key);
+    if (!entry) return null;
+    
+    if (Date.now() > entry.resetTime) {
+      this.store.delete(key);
+      return null;
+    }
+    
+    return entry;
+  }
+
+  async set(key: string, count: number, _windowMs: number, resetTime: number): Promise<void> {
+    this.store.set(key, { count, resetTime });
+  }
+
+  async increment(key: string): Promise<{ count: number; resetTime: number }> {
+    const now = Date.now();
+    const entry = this.store.get(key);
+    
+    if (!entry || now > entry.resetTime) {
+      const resetTime = now + 60000;
+      this.store.set(key, { count: 1, resetTime });
+      return { count: 1, resetTime };
+    }
+    
+    const newCount = entry.count + 1;
+    this.store.set(key, { count: newCount, resetTime: entry.resetTime });
+    return { count: newCount, resetTime: entry.resetTime };
+  }
+
+  async delete(key: string): Promise<void> {
+    this.store.delete(key);
+  }
+
+  async reset(): Promise<void> {
+    this.store.clear();
+  }
+}
+
+export class RateLimiter {
+  private store: RateLimitStore;
+  private config: RateLimitConfig;
+
+  constructor(config: RateLimitConfig, store?: RateLimitStore) {
+    this.config = config;
+    this.store = store || new InMemoryRateLimitStore();
+  }
+
+  async check(identifier: string): Promise<RateLimitResult> {
+    const key = `${this.config.keyPrefix || 'ratelimit'}:${identifier}`;
+    const now = Date.now();
+    
+    const entry = await this.store.get(key);
+    
+    if (!entry) {
+      const resetTime = now + this.config.windowMs;
+      await this.store.set(key, 1, this.config.windowMs, resetTime);
+      
+      return {
+        allowed: true,
+        remaining: this.config.maxRequests - 1,
+        resetTime,
+      };
+    }
+
+    if (now >= entry.resetTime) {
+      const resetTime = now + this.config.windowMs;
+      await this.store.set(key, 1, this.config.windowMs, resetTime);
+      
+      return {
+        allowed: true,
+        remaining: this.config.maxRequests - 1,
+        resetTime,
+      };
+    }
+
+    if (entry.count >= this.config.maxRequests) {
+      const retryAfter = Math.ceil((entry.resetTime - now) / 1000);
+      
+      return {
+        allowed: false,
+        remaining: 0,
+        resetTime: entry.resetTime,
+        retryAfter,
+      };
+    }
+
+    const newEntry = await this.store.increment(key);
+    
+    return {
+      allowed: true,
+      remaining: this.config.maxRequests - newEntry.count,
+      resetTime: newEntry.resetTime,
+    };
+  }
+
+  async reset(identifier: string): Promise<void> {
+    const key = `${this.config.keyPrefix || 'ratelimit'}:${identifier}`;
+    await this.store.delete(key);
+  }
+
+  updateConfig(config: Partial<RateLimitConfig>): void {
+    this.config = { ...this.config, ...config };
+  }
+}
+
+export const DEFAULT_RATE_LIMITS: Record<string, RateLimitConfig> = {
+  // NPM & Extension Free
+  free: {
+    windowMs: 60 * 1000,
+    maxRequests: 100,
+    keyPrefix: 'free',
+  },
+  // Extension Starter
+  starter: {
+    windowMs: 60 * 1000,
+    maxRequests: 200,
+    keyPrefix: 'starter',
+  },
+  // NPM Pro
+  pro: {
+    windowMs: 60 * 1000,
+    maxRequests: 500,
+    keyPrefix: 'pro',
+  },
+  // Extension Plus
+  plus: {
+    windowMs: 60 * 1000,
+    maxRequests: 300,
+    keyPrefix: 'plus',
+  },
+  // Extension Premium
+  premium: {
+    windowMs: 60 * 1000,
+    maxRequests: 1000,
+    keyPrefix: 'premium',
+  },
+  // Team
+  team: {
+    windowMs: 60 * 1000,
+    maxRequests: 2000,
+    keyPrefix: 'team',
+  },
+  // Business
+  business: {
+    windowMs: 60 * 1000,
+    maxRequests: 5000,
+    keyPrefix: 'business',
+  },
+  // Enterprise
+  enterprise: {
+    windowMs: 60 * 1000,
+    maxRequests: 10000,
+    keyPrefix: 'enterprise',
+  },
+};
+
+export function createRateLimiter(plan: string = 'free', store?: RateLimitStore): RateLimiter {
+  const config = DEFAULT_RATE_LIMITS[plan] || DEFAULT_RATE_LIMITS.free;
+  return new RateLimiter(config, store);
+}

package/src/storage/index.ts

@@ -0,0 +1,2 @@
+export { RedisRateLimitStore, createRedisRateLimiter, type RedisConfig } from './redis.js';
+export { PostgresUsageStore, createPostgresUsageStore, type PostgresConfig } from './postgres.js';

package/src/storage/postgres.ts

@@ -0,0 +1,209 @@
+import { UsageStore, UsageRecord, UsageMetric } from '../usageTracker';
+
+export interface PostgresConfig {
+  host?: string;
+  port?: number;
+  database?: string;
+  user?: string;
+  password?: string;
+  connectionString?: string;
+  max?: number;
+  idleTimeoutMillis?: number;
+  connectionTimeoutMillis?: number;
+  ssl?: boolean;
+}
+
+export class PostgresUsageStore implements UsageStore {
+  private pool: any;
+  private initialized: boolean = false;
+
+  constructor(private config: PostgresConfig) {}
+
+  private async getPool(): Promise<any> {
+    if (!this.initialized) {
+      try {
+        const { default: pg } = await import('pg');
+        const { Pool } = pg;
+        
+        this.pool = new Pool({
+          host: this.config.host || 'localhost',
+          port: this.config.port || 5432,
+          database: this.config.database || 'ai_brain',
+          user: this.config.user || 'postgres',
+          password: this.config.password,
+          connectionString: this.config.connectionString,
+          max: this.config.max || 20,
+          idleTimeoutMillis: this.config.idleTimeoutMillis || 30000,
+          connectionTimeoutMillis: this.config.connectionTimeoutMillis || 5000,
+          ssl: this.config.ssl ? { rejectUnauthorized: false } : false,
+        });
+
+        await this.initialize();
+        this.initialized = true;
+      } catch (error) {
+        console.error('Failed to initialize PostgreSQL:', error);
+        throw new Error('PostgreSQL not available');
+      }
+    }
+    return this.pool;
+  }
+
+  private async initialize(): Promise<void> {
+    const pool = this.pool;
+    
+    await pool.query(`
+      CREATE TABLE IF NOT EXISTS usage_records (
+        id SERIAL PRIMARY KEY,
+        user_id VARCHAR(255) NOT NULL,
+        date DATE NOT NULL,
+        memories_added INTEGER DEFAULT 0,
+        memories_searched INTEGER DEFAULT 0,
+        context_built INTEGER DEFAULT 0,
+        api_calls INTEGER DEFAULT 0,
+        storage_used INTEGER DEFAULT 0,
+        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+        UNIQUE(user_id, date)
+      );
+    `);
+
+    await pool.query(`
+      CREATE INDEX IF NOT EXISTS idx_usage_user_date 
+      ON usage_records(user_id, date);
+    `);
+  }
+
+  async get(userId: string, date: string): Promise<UsageRecord | null> {
+    try {
+      const pool = await this.getPool();
+      const result = await pool.query(
+        'SELECT * FROM usage_records WHERE user_id = $1 AND date = $2',
+        [userId, date]
+      );
+
+      if (result.rows.length === 0) return null;
+
+      const row = result.rows[0];
+      return {
+        userId: row.user_id,
+        date: row.date.toISOString().split('T')[0],
+        metrics: {
+          memories_added: row.memories_added,
+          memories_searched: row.memories_searched,
+          context_built: row.context_built,
+          api_calls: row.api_calls,
+          storage_used: row.storage_used,
+        },
+      };
+    } catch {
+      return null;
+    }
+  }
+
+  async set(userId: string, date: string, record: UsageRecord): Promise<void> {
+    try {
+      const pool = await this.getPool();
+      await pool.query(
+        `INSERT INTO usage_records (user_id, date, memories_added, memories_searched, context_built, api_calls, storage_used)
+         VALUES ($1, $2, $3, $4, $5, $6, $7)
+         ON CONFLICT (user_id, date) DO UPDATE SET
+           memories_added = EXCLUDED.memories_added,
+           memories_searched = EXCLUDED.memories_searched,
+           context_built = EXCLUDED.context_built,
+           api_calls = EXCLUDED.api_calls,
+           storage_used = EXCLUDED.storage_used,
+           updated_at = CURRENT_TIMESTAMP`,
+        [
+          userId,
+          date,
+          record.metrics.memories_added,
+          record.metrics.memories_searched,
+          record.metrics.context_built,
+          record.metrics.api_calls,
+          record.metrics.storage_used,
+        ]
+      );
+    } catch (error) {
+      console.error('PostgreSQL set error:', error);
+    }
+  }
+
+  async increment(userId: string, metric: UsageMetric, amount: number = 1): Promise<void> {
+    const today = new Date().toISOString().split('T')[0];
+    const metricColumn = metric;
+
+    try {
+      const pool = await this.getPool();
+      await pool.query(
+        `INSERT INTO usage_records (user_id, date, ${metricColumn})
+         VALUES ($1, $2, $3)
+         ON CONFLICT (user_id, date) DO UPDATE SET
+           ${metricColumn} = usage_records.${metricColumn} + EXCLUDED.${metricColumn},
+           updated_at = CURRENT_TIMESTAMP`,
+        [userId, today, amount]
+      );
+    } catch (error) {
+      console.error('PostgreSQL increment error:', error);
+    }
+  }
+
+  async getRange(userId: string, startDate: string, endDate: string): Promise<UsageRecord[]> {
+    try {
+      const pool = await this.getPool();
+      const result = await pool.query(
+        'SELECT * FROM usage_records WHERE user_id = $1 AND date >= $2 AND date <= $3 ORDER BY date',
+        [userId, startDate, endDate]
+      );
+
+      return result.rows.map((row: any) => ({
+        userId: row.user_id,
+        date: row.date.toISOString().split('T')[0],
+        metrics: {
+          memories_added: row.memories_added,
+          memories_searched: row.memories_searched,
+          context_built: row.context_built,
+          api_calls: row.api_calls,
+          storage_used: row.storage_used,
+        },
+      }));
+    } catch {
+      return [];
+    }
+  }
+
+  async reset(userId: string, date: string): Promise<void> {
+    try {
+      const pool = await this.getPool();
+      await pool.query(
+        'DELETE FROM usage_records WHERE user_id = $1 AND date = $2',
+        [userId, date]
+      );
+    } catch (error) {
+      console.error('PostgreSQL reset error:', error);
+    }
+  }
+
+  async deleteUser(userId: string): Promise<void> {
+    try {
+      const pool = await this.getPool();
+      await pool.query('DELETE FROM usage_records WHERE user_id = $1', [userId]);
+    } catch (error) {
+      console.error('PostgreSQL deleteUser error:', error);
+    }
+  }
+
+  async disconnect(): Promise<void> {
+    if (this.pool) {
+      await this.pool.end();
+      this.initialized = false;
+    }
+  }
+
+  isInitialized(): boolean {
+    return this.initialized;
+  }
+}
+
+export function createPostgresUsageStore(config: PostgresConfig): PostgresUsageStore {
+  return new PostgresUsageStore(config);
+}

package/src/storage/redis.ts

@@ -0,0 +1,117 @@
+import { RateLimitStore } from '../rateLimiter';
+
+export interface RedisConfig {
+  host?: string;
+  port?: number;
+  url?: string;
+  password?: string;
+  db?: number;
+  keyPrefix?: string;
+  connectionTimeout?: number;
+  maxRetriesPerRequest?: number;
+}
+
+export class RedisRateLimitStore implements RateLimitStore {
+  private redis: any;
+  private keyPrefix: string;
+  private connected: boolean = false;
+
+  constructor(private config: RedisConfig) {
+    this.keyPrefix = config.keyPrefix || 'ai-memory:ratelimit:';
+  }
+
+  private async getClient(): Promise<any> {
+    if (!this.connected) {
+      try {
+        const Redis = (await import('ioredis')).default;
+        this.redis = new Redis({
+          host: this.config.host || 'localhost',
+          port: this.config.port || 6379,
+          password: this.config.password,
+          db: this.config.db || 0,
+          lazyConnect: true,
+        });
+        await this.redis.connect();
+        this.connected = true;
+      } catch (error) {
+        console.error('Failed to connect to Redis:', error);
+        throw new Error('Redis not available');
+      }
+    }
+    return this.redis;
+  }
+
+  async get(key: string): Promise<{ count: number; resetTime: number } | null> {
+    try {
+      const client = await this.getClient();
+      const data = await client.get(this.keyPrefix + key);
+      if (!data) return null;
+      return JSON.parse(data);
+    } catch {
+      return null;
+    }
+  }
+
+  async set(key: string, count: number, windowMs: number, resetTime: number): Promise<void> {
+    try {
+      const client = await this.getClient();
+      const data = JSON.stringify({ count, resetTime });
+      const ttlSeconds = Math.ceil(windowMs / 1000);
+      await client.setex(this.keyPrefix + key, ttlSeconds, data);
+    } catch (error) {
+      console.error('Redis set error:', error);
+    }
+  }
+
+  async increment(key: string): Promise<{ count: number; resetTime: number }> {
+    const client = await this.getClient();
+    const fullKey = this.keyPrefix + key;
+    
+    const existing = await this.get(key);
+    if (!existing) {
+      const resetTime = Date.now() + 60000;
+      await this.set(key, 1, 60000, resetTime);
+      return { count: 1, resetTime };
+    }
+
+    const newCount = existing.count + 1;
+    await this.set(key, newCount, 60000, existing.resetTime);
+    return { count: newCount, resetTime: existing.resetTime };
+  }
+
+  async delete(key: string): Promise<void> {
+    try {
+      const client = await this.getClient();
+      await client.del(this.keyPrefix + key);
+    } catch (error) {
+      console.error('Redis delete error:', error);
+    }
+  }
+
+  async reset(): Promise<void> {
+    try {
+      const client = await this.getClient();
+      const keys = await client.keys(this.keyPrefix + '*');
+      if (keys.length > 0) {
+        await client.del(...keys);
+      }
+    } catch (error) {
+      console.error('Redis reset error:', error);
+    }
+  }
+
+  async disconnect(): Promise<void> {
+    if (this.redis && this.connected) {
+      await this.redis.quit();
+      this.connected = false;
+    }
+  }
+
+  isConnected(): boolean {
+    return this.connected;
+  }
+}
+
+export function createRedisRateLimiter(config: RedisConfig): RedisRateLimitStore {
+  return new RedisRateLimitStore(config);
+}

package/src/types.ts

@@ -0,0 +1,114 @@
+export interface Memory {
+  id: string;
+  content: string;
+  embedding?: number[];
+  timestamp: number;
+  metadata: MemoryMetadata;
+}
+
+export interface MemoryMetadata {
+  type: MemoryType;
+  importance: number;
+  tags: string[];
+  userId?: string;
+  sessionId?: string;
+  source?: string;
+  [key: string]: unknown;
+}
+
+export type MemoryType = 
+  | 'conversation' 
+  | 'fact' 
+  | 'preference' 
+  | 'instruction' 
+  | 'context' 
+  | 'custom';
+
+export interface SearchResult {
+  memory: Memory;
+  score: number;
+}
+
+export interface MemoryQuery {
+  text?: string;
+  embedding?: number[];
+  filter?: MemoryFilter;
+  limit?: number;
+  threshold?: number;
+}
+
+export interface MemoryFilter {
+  types?: MemoryType[];
+  tags?: string[];
+  userId?: string;
+  sessionId?: string;
+  startDate?: number;
+  endDate?: number;
+  minImportance?: number;
+  maxImportance?: number;
+  metadata?: Record<string, unknown>;
+}
+
+export interface MemoryStats {
+  totalMemories: number;
+  memoriesByType: Record<MemoryType, number>;
+  memoriesByTag: Record<string, number>;
+  oldestMemory: number;
+  newestMemory: number;
+}
+
+export interface ContextConfig {
+  maxTokens: number;
+  includeSystemPrompt: boolean;
+  systemPrompt?: string;
+  relevanceThreshold: number;
+  includeMetadata: boolean;
+  memoryTypes: MemoryType[];
+  maxMemories: number;
+}
+
+export interface ContextMessage {
+  role: 'system' | 'user' | 'assistant';
+  content: string;
+}
+
+export interface ContextResult {
+  messages: ContextMessage[];
+  memories: Memory[];
+  stats: {
+    memoriesUsed: number;
+    tokensEstimated: number;
+    truncationApplied: boolean;
+  };
+}
+
+export interface EmbeddingFunction {
+  (text: string): Promise<number[]>;
+}
+
+export interface StorageAdapter {
+  save(key: string, data: unknown): Promise<void>;
+  load(key: string): Promise<unknown | null>;
+  delete(key: string): Promise<void>;
+  exists(key: string): Promise<boolean>;
+  list(prefix: string): Promise<string[]>;
+}
+
+export interface VectorStoreAdapter {
+  add(vectors: VectorEntry[]): Promise<void>;
+  search(query: number[], options: SearchOptions): Promise<SearchResult[]>;
+  delete(ids: string[]): Promise<void>;
+  getById(id: string): Promise<VectorEntry | null>;
+}
+
+export interface VectorEntry {
+  id: string;
+  vector: number[];
+  metadata: Record<string, unknown>;
+}
+
+export interface SearchOptions {
+  limit: number;
+  threshold?: number;
+  filter?: Record<string, unknown>;
+}