index.js

/**
 * @fileoverview Hybrid file-based cache with optional semantic 
 * (embedding-based) lookup.
 *
 * Entries are stored as JSON files in a sharded directory tree:
 *
 *   <dir>/<segment>/<sem>/<a>/<ab>/<abc>/<md5hash>.json
 *
 * Where:
 *   - <dir>      — base cache directory (e.g. './cache')
 *   - <segment>  — logical namespace (e.g. 'treatments')
 *   - <sem>      — '+' for semantic entries, '-' for exact-match entries
 *   - <a/ab/abc> — three-level sharding prefix derived from the MD5 of the 
 *                  query
 *   - <md5hash>  — full MD5 hex digest of the lowercased query string
 *
 * Sharding keeps any single directory from accumulating too many files,
 * which degrades filesystem performance on large caches.
 *
 * @module cache
 */

import fs from 'node:fs/promises';
import node_path from 'node:path';
import crypto from 'node:crypto';
import { pipeline, cos_sim } from '@huggingface/transformers';

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

/**
 * @typedef {Object} CacheConfig
 * @property {string}  dir                 - Base directory for all cache files.
 * @property {string}  segment             - Logical namespace / partition name.
 * @property {number}  ttl                 - Default time-to-live in ms. 
 *                                           Use -1 for no expiry.
 * @property {number}  similarityThreshold - Cosine-similarity threshold for 
 *                                           semantic hits (0–1).
 */

/**
 * @typedef {Object} CacheEntry
 * @property {string}   query       - Original query string.
 * @property {*}        response    - Cached response value.
 * @property {number}   stored      - Unix timestamp (ms) when the entry was 
 *                                    written.
 * @property {number}   ttl         - TTL in ms at write time. -1 means 
 *                                    immortal.
 * @property {number[]} [embedding] - Pre-computed embedding vector (semantic 
 *                                    entries only).
 */

/**
 * @typedef {Object} GetOpts
 * @property {string}   query              - The query to look up.
 * @property {string}   [segment]          - Override the instance default 
 *                                           segment.
 * @property {boolean}  [isSemantic=false] - Also search for semantically 
 *                                           similar queries.
 * @property {string[]} [omit]             - Keys to strip from the returned 
 *                                           entry before returning it to the 
 *                                           caller. Useful for removing large 
 *                                           internal fields (e.g. 'embedding')
 *                                           or environment-specific fields 
 *                                           (e.g. 'debug') that should never 
 *                                           reach the client. The on-disk 
 *                                           entry is never modified.
 */

/**
 * @typedef {Object} SetOpts
 * @property {string}  query              - The query to cache.
 * @property {*}       response           - The value to store.
 * @property {string}  [segment]          - Override the instance default 
 *                                          segment.
 * @property {boolean} [isSemantic=false] - Store with a pre-computed embedding.
 * @property {number}  [ttl]              - Override the instance default TTL.
 */

/**
 * @typedef {Object} KeyOpts
 * @property {string}  query              - The query to target.
 * @property {string}  [segment]          - Override the instance default 
 *                                          segment.
 * @property {boolean} [isSemantic=false] - Target the semantic (+) subtree.
 */

/**
 * @typedef {Object} ListOpts
 * @property {string} [segment]            - Override the instance default 
 *                                           segment.
 * @property {boolean} [isSemantic=false]  - List entries in the semantic (+) 
 *                                           subtree.
 * @property {string[]} [omit]             - Keys to strip from every returned 
 *                                           entry.
 *                                           Follows the same semantics as {@link GetOpts#omit}.
 */

// ---------------------------------------------------------------------------
// Embedding singleton
// ---------------------------------------------------------------------------

/**
 * Lazy singleton that loads the HuggingFace feature-extraction pipeline once
 * and reuses it for every subsequent embedding request.
 *
 * Using a static class keeps the singleton scoped to this module while still
 * allowing the model / task to be overridden in one place.
 *
 * @private
 */
class EmbeddingModel {
    static task  = 'feature-extraction';
    static model = 'Supabase/gte-small';

    /** @type {Promise<import('@huggingface/transformers').FeatureExtractionPipeline> | null} */
    static #instance = null;

    /**
     * Returns the shared pipeline instance, creating it on first call.
     *
     * The result is a Promise that is stored immediately so that concurrent
     * callers await the same initialisation rather than starting multiple loads.
     *
     * @param {Function} [progressCallback] - Optional download-progress hook.
     * @returns {Promise<import('@huggingface/transformers').FeatureExtractionPipeline>}
     */
    static getInstance(progressCallback = null) {
        if (this.#instance === null) {
            this.#instance = pipeline(this.task, this.model, {
                progress_callback: progressCallback,
                dtype: 'fp32',
            });
        }
        return this.#instance;
    }
}

/**
 * Generates a normalised mean-pooled embedding vector for the given text.
 *
 * @param {string} text - Input text to embed.
 * @returns {Promise<number[]>} Flat array of floats representing the embedding.
 * @private
 */
async function generateEmbedding(text) {
    const model    = await EmbeddingModel.getInstance();
    const response = await model(text, { pooling: 'mean', normalize: true });
    return Array.from(response.data);
}

// ---------------------------------------------------------------------------
// Path helpers
// ---------------------------------------------------------------------------

/**
 * Derives the shard directory and filename for a given query.
 *
 * The MD5 of the lowercased query is used as the filename, and its first three
 * characters drive three levels of directory sharding:
 *
 *   hash = 'acbd18db...'  →  dirs: a / ac / acb  →  file: acbd18db....json
 *
 * This keeps each directory to at most ~16 entries on average (for a uniform
 * hash distribution), which is filesystem-friendly.
 *
 * @param {object}  opts
 * @param {string}  opts.dir        - Base cache directory.
 * @param {string}  opts.segment    - Logical namespace.
 * @param {boolean} [opts.isSemantic=false] - Whether this is a semantic entry.
 * @param {string}  opts.query      - The query string to derive a path for.
 * @returns {{ filepath: string, filename: string }}
 * @private
 */
function derivePaths({ dir, segment, isSemantic=false, query }) {
    const hash = crypto
        .createHash('md5')
        .update(query.toLowerCase())
        .digest('hex');

    // Three shard levels: first char, first two chars, first three chars.
    const shard = [hash[0], hash.slice(0, 2), hash.slice(0, 3)];

    const sem = isSemantic ? '+' : '-';

    return {
        filepath: node_path.join(dir, segment, sem, ...shard),
        filename: `${hash}.json`,
    };
}

// ---------------------------------------------------------------------------
// TTL helper
// ---------------------------------------------------------------------------

/**
 * Returns true if the cache entry's TTL has elapsed.
 *
 * A TTL of -1 (or any negative number) means the entry never expires.
 *
 * @param {CacheEntry} entry - The entry to test.
 * @returns {boolean}
 * @private
 */
function isExpired(entry) {
    if (entry.ttl < 0) return false;
    return (entry.stored + entry.ttl) <= Date.now();
}

// ---------------------------------------------------------------------------
// Projection helper
// ---------------------------------------------------------------------------

/**
 * Returns a shallow copy of `entry` with the specified keys removed.
 *
 * This is applied at the boundary between the cache internals and the caller,
 * so internal bookkeeping fields (e.g. `embedding`, `debug`) are never
 * accidentally forwarded to a client.  The on-disk file is never touched.
 *
 * Passing an empty or absent `omit` list is a no-op — the original entry
 * object is returned as-is (no copy is made) to avoid unnecessary allocation
 * on the common path.
 *
 * @param {CacheEntry}  entry - The entry read from disk.
 * @param {string[]}    [omit=[]] - Field names to exclude from the result.
 * @returns {CacheEntry} The projected entry (may be the original reference).
 * @private
 *
 * @example
 * // Strip the embedding vector and a debug field before sending to the client:
 * const entry = await cache.get({ query, isSemantic: true, omit: ['embedding', 'debug'] });
 */
function projectEntry(entry, omit = []) {
    if (omit.length === 0) return entry;

    const projected = { ...entry };
    for (const key of omit) {
        delete projected[key];
    }
    return projected;
}



/**
 * Resolves to true if the given path exists on disk, false otherwise.
 *
 * Uses stat() rather than access() so it works for both files and directories.
 *
 * @param {string} p - Absolute or relative filesystem path.
 * @returns {Promise<boolean>}
 * @private
 */
async function pathExists(p) {
    return fs.stat(p).then(() => true, () => false);
}

/**
 * Recursively creates a directory, ignoring errors only when the directory
 * already exists (which `recursive: true` handles natively in Node ≥ 10).
 *
 * @param {string} dir - Directory path to create.
 * @returns {Promise<void>}
 * @throws {Error} If the directory cannot be created for any reason.
 * @private
 */
async function mkdirp(dir) {
    await fs.mkdir(dir, { recursive: true });
}

/**
 * Walks a directory tree depth-first and invokes `cb` for every `.json` file
 * encountered.  The callback is awaited before moving to the next entry so
 * that callers can perform sequential I/O without managing concurrency.
 *
 * @param {string}                      dir - Root directory to walk.
 * @param {(file: string) => Promise<void>} cb  - Called with the absolute path of each JSON file.
 * @returns {Promise<void>}
 * @private
 */
async function walkJsonFiles(dir, cb) {
    const entries = await fs.readdir(dir, { withFileTypes: true });

    for (const entry of entries) {
        const full = node_path.join(dir, entry.name);

        if (entry.isDirectory()) {
            await walkJsonFiles(full, cb);
        } else if (node_path.extname(entry.name) === '.json') {
            await cb(full);
        }
    }
}

// ---------------------------------------------------------------------------
// Cache class
// ---------------------------------------------------------------------------

/**
 * Hybrid file-based cache supporting both exact-match and semantic lookups.
 *
 * ### Exact-match lookup
 * The query is MD5-hashed and mapped to a JSON file.  This is O(1) and works
 * for any value type.
 *
 * ### Semantic lookup (`isSemantic: true`)
 * When no exact match exists, the cache walks every semantic entry, compares
 * embeddings using cosine similarity, and returns the best match above the
 * configured threshold.  Embeddings are stored alongside entries at write time
 * so that only the *source* query requires an embedding call at read time.
 *
 * @example
 * const cache = new Cache({ dir: './cache', segment: 'llm', ttl: 86_400_000 });
 * await cache.init();
 *
 * await cache.set({ query: 'hello', response: 'world' });
 * const entry = await cache.get({ query: 'hello' });
 * console.log(entry.response); // 'world'
 */
class Cache {

    /**
     * @param {Partial<CacheConfig>} opts - Configuration overrides.
     */
    constructor(opts = {}) {
        /** @type {CacheConfig} */
        this.config = {
            dir:                './cache',
            segment:            'default',
            ttl:                1 * 24 * 60 * 60 * 1000, // 1 day in ms
            similarityThreshold: 0.9,
            ...opts,
        };
    }

    // -----------------------------------------------------------------------
    // Lifecycle
    // -----------------------------------------------------------------------

    /**
     * Ensures the cache base directory exists.
     *
     * Must be called once before using any other method, e.g.:
     * ```js
     * const cache = new Cache({ dir: './cache' });
     * await cache.init();
     * ```
     *
     * @returns {Promise<void>}
     */
    async init() {
        await mkdirp(this.config.dir);
    }

    // -----------------------------------------------------------------------
    // Public API
    // -----------------------------------------------------------------------

    /**
     * Retrieves a cached entry by query.
     *
     * Resolution order:
     * 1. Exact match in the non-semantic ('-') subtree.
     * 2. Exact match in the semantic ('+') subtree (same hash, different dir).
     * 3. Semantic nearest-neighbour search (only when `isSemantic` is true).
     *
     * Expired entries are deleted on access and not returned.
     *
     * @param {GetOpts} opts
     * @returns {Promise<CacheEntry|null>} The cached entry, or null on miss.
     */
    async get(opts) {
        if (!this.#assertQuery('get', opts)) return null;

        const { query, isSemantic = false, omit = [] } = opts;
        const segment = opts.segment ?? this.config.segment;

        // Step 1 & 2 — exact match in either subtree.
        const exact = await this.#readEntry({ segment, query, isSemantic: false })
                   ?? await this.#readEntry({ segment, query, isSemantic: true  });
        if (exact) return projectEntry(exact, omit);

        // Step 3 — semantic nearest-neighbour search.
        if (isSemantic) {
            const match = await this.#semanticSearch({ segment, query });
            return match ? projectEntry(match, omit) : null;
        }

        return null;
    }

    /**
     * Stores a query/response pair in the cache.
     *
     * When `isSemantic` is true the entry's embedding is computed and stored
     * alongside the response so that future semantic lookups do not need to
     * re-embed cached entries.
     *
     * @param {SetOpts} opts
     * @returns {Promise<CacheEntry|null>} The written entry, or null on error.
     */
    async set(opts) {
        if (!this.#assertQuery('set', opts))    return null;
        if (!this.#assertResponse('set', opts)) return null;

        const { query, response, isSemantic = false } = opts;
        const segment = opts.segment ?? this.config.segment;
        const ttl     = opts.ttl     ?? this.config.ttl;

        /** @type {CacheEntry} */
        const entry = { query, response, stored: Date.now(), ttl };

        // Pre-compute and store the embedding at write time so semantic reads
        // only ever need to embed the *source* query, not every cached entry.
        if (isSemantic) {
            entry.embedding = await generateEmbedding(query);
        }

        const { filepath, filename } = derivePaths({
            dir: this.config.dir,
            segment,
            isSemantic,
            query,
        });

        const cacheFile = node_path.join(filepath, filename);

        await mkdirp(filepath);
        await fs.writeFile(cacheFile, JSON.stringify(entry), 'utf8');

        return entry;
    }

    /**
     * Removes a single cached entry.
     *
     * @param {KeyOpts} opts
     * @returns {Promise<boolean>} True if the file was deleted, false if it did not exist.
     */
    async rm(opts) {
        if (!this.#assertQuery('rm', opts)) return false;

        const { query, isSemantic = false } = opts;
        const segment = opts.segment ?? this.config.segment;

        const { filepath, filename } = derivePaths({
            dir: this.config.dir,
            segment,
            isSemantic,
            query,
        });

        const cacheFile = node_path.join(filepath, filename);

        if (!(await pathExists(cacheFile))) return false;

        await fs.unlink(cacheFile);
        return true;
    }

    /**
     * Alias for {@link Cache#rm}.
     * @param {KeyOpts} opts
     * @returns {Promise<boolean>}
     */
    del(opts) { return this.rm(opts); }

    /**
     * Alias for {@link Cache#rm}.
     * @param {KeyOpts} opts
     * @returns {Promise<boolean>}
     */
    delete(opts) { return this.rm(opts); }

    /**
     * Returns true when a non-expired entry exists for the given query.
     *
     * Unlike a raw filesystem check, `has()` respects TTL, so an expired entry
     * returns false (and is cleaned up).
     *
     * @param {KeyOpts} opts
     * @returns {Promise<boolean>}
     */
    async has(opts) {
        if (!this.#assertQuery('has', opts)) return false;

        const entry = await this.get({ ...opts, isSemantic: false });
        return entry !== null;
    }

    /**
     * Lists all valid (non-expired) entries in the specified subtree.
     *
     * Expired entries encountered during the walk are deleted as a side effect.
     *
     * @param {ListOpts} [opts]
     * @returns {Promise<CacheEntry[]>} Array of live cache entries, with any
     *   `omit` fields stripped from each element.
     */
    async queries(opts) {
        const segment    = opts?.segment    ?? this.config.segment;
        const isSemantic = opts?.isSemantic ?? false;
        const omit       = opts?.omit       ?? [];
        const { live }   = await this.#walkEntries({ segment, isSemantic });
        return omit.length === 0 ? live : live.map(e => projectEntry(e, omit));
    }

    /**
     * Deletes all expired entries in the specified subtree and returns a count.
     *
     * Internally shares its walk logic with {@link Cache#queries} to avoid
     * duplicating the traversal code.
     *
     * @param {ListOpts} [opts]
     * @returns {Promise<number>} Number of expired entries removed.
     */
    async prune(opts) {
        const segment    = opts?.segment    ?? this.config.segment;
        const isSemantic = opts?.isSemantic ?? false;
        const { pruned } = await this.#walkEntries({ segment, isSemantic });
        return pruned;
    }

    // -----------------------------------------------------------------------
    // Private helpers
    // -----------------------------------------------------------------------

    /**
     * Reads a single cache entry from disk, handling TTL expiry.
     *
     * Returns `null` on any of: file absent, parse failure, expired TTL.
     *
     * @param {{ segment: string, query: string, isSemantic: boolean }} opts
     * @returns {Promise<CacheEntry|null>}
     * @private
     */
    async #readEntry({ segment, query, isSemantic }) {
        const { filepath, filename } = derivePaths({
            dir: this.config.dir,
            segment,
            isSemantic,
            query,
        });

        const cacheFile = node_path.join(filepath, filename);

        if (!(await pathExists(cacheFile))) return null;

        let entry;
        try {
            const raw = await fs.readFile(cacheFile, 'utf8');
            entry = JSON.parse(raw);
        } catch {
            // Corrupt file — treat as a miss.
            return null;
        }

        if (isExpired(entry)) {
            // Lazy deletion: remove the stale file and report a miss.
            await fs.unlink(cacheFile).catch(() => {});
            return null;
        }

        return entry;
    }

    /**
     * Searches the semantic subtree for the closest matching entry.
     *
     * Algorithm:
     *  1. Embed the source query once.
     *  2. Walk every `.json` file in the semantic ('+') directory.
     *  3. For each entry, retrieve its stored embedding (no re-embedding needed).
     *  4. Compute cosine similarity; track the best match above the threshold.
     *
     * This is O(n) in the number of semantic entries.  For very large caches
     * consider an approximate nearest-neighbour index instead.
     *
     * @param {{ segment: string, query: string }} opts
     * @returns {Promise<CacheEntry|null>}
     * @private
     */
    async #semanticSearch({ segment, query }) {
        const semanticDir = node_path.join(this.config.dir, segment, '+');
        if (!(await pathExists(semanticDir))) return null;

        const srcEmbedding = await generateEmbedding(query);

        let bestMatch         = null;
        let highestSimilarity = 0;

        await walkJsonFiles(semanticDir, async (file) => {
            let entry;
            try {
                const raw = await fs.readFile(file, 'utf8');
                entry = JSON.parse(raw);
            } catch {
                return; // Skip corrupt files.
            }

            // Lazy-delete expired entries found during the walk.
            if (isExpired(entry)) {
                await fs.unlink(file).catch(() => {});
                return;
            }

            // Entries written by set() always carry a stored embedding.
            // Skip any legacy entries that pre-date this field.
            if (!Array.isArray(entry.embedding)) return;

            const similarity = cos_sim(srcEmbedding, entry.embedding);

            if (similarity >= this.config.similarityThreshold
                && similarity > highestSimilarity) {
                highestSimilarity = similarity;
                bestMatch         = entry;
            }
        });

        return bestMatch;
    }

    /**
     * Walks an entire subtree, separating live entries from prunable ones.
     *
     * Shared by {@link Cache#queries} and {@link Cache#prune} to avoid
     * duplicating traversal logic.
     *
     * @param {{ segment: string, isSemantic: boolean }} opts
     * @returns {Promise<{ live: CacheEntry[], pruned: number }>}
     * @private
     */
    async #walkEntries({ segment, isSemantic }) {
        const sem      = isSemantic ? '+' : '-';
        const cacheDir = node_path.join(this.config.dir, segment, sem);

        const result = { live: [], pruned: 0 };

        if (!(await pathExists(cacheDir))) return result;

        await walkJsonFiles(cacheDir, async (file) => {
            let entry;
            try {
                const raw = await fs.readFile(file, 'utf8');
                entry = JSON.parse(raw);
            } catch {
                return; // Skip corrupt files.
            }

            if (isExpired(entry)) {
                await fs.unlink(file).catch(() => {});
                result.pruned++;
            } else {
                result.live.push(entry);
            }
        });

        return result;
    }

    // -----------------------------------------------------------------------
    // Guard helpers
    // -----------------------------------------------------------------------

    /**
     * Logs an error and returns false when `opts.query` is absent.
     *
     * @param {string} method - Calling method name, used in the error message.
     * @param {object} opts   - Options object passed to the public method.
     * @returns {boolean} True when the query is present and non-empty.
     * @private
     */
    #assertQuery(method, opts) {
        if (opts?.query) return true;
        console.error(`[Cache] ${method}(): 'query' is required.`);
        return false;
    }

    /**
     * Logs an error and returns false when `opts.response` is absent.
     *
     * Note: `undefined` is treated as absent; `null` and `false` are valid
     * response values and will pass this check.
     *
     * @param {string} method - Calling method name, used in the error message.
     * @param {object} opts   - Options object passed to the public method.
     * @returns {boolean} True when the response field is present.
     * @private
     */
    #assertResponse(method, opts) {
        if (opts?.response !== undefined) return true;
        console.error(`[Cache] ${method}(): 'response' is required.`);
        return false;
    }
}

export { Cache };