{"version":3,"sources":["../src/index.ts"],"sourcesContent":["import {\n    ChromaValueError,\n    EmbeddingFunction,\n    EmbeddingFunctionSpace,\n    registerEmbeddingFunction,\n} from \"chromadb\";\nimport {\n    validateConfigSchema,\n} from \"@chroma-core/ai-embeddings-common\";\nimport { pipeline } from \"@huggingface/transformers\";\n\nconst NAME = \"sentence_transformer\";\n\nexport interface SentenceTransformersConfig {\n    model_name: string;\n    device: string;\n    normalize_embeddings: boolean;\n    kwargs?: Record<string, any>;\n}\n\nexport interface SentenceTransformersArgs {\n    modelName?: string;\n    device?: string;\n    normalizeEmbeddings?: boolean;\n    kwargs?: Record<string, any>;\n}\n\nexport class SentenceTransformersEmbeddingFunction\n    implements EmbeddingFunction {\n    public readonly name = NAME;\n    private readonly modelName: string;\n    private readonly device: string;\n    private readonly normalizeEmbeddings: boolean;\n    private readonly kwargs: Record<string, any>;\n    private pipelinePromise: Promise<any> | null = null;\n    private pipeline: any = null;\n\n    constructor(args: SentenceTransformersArgs = {}) {\n        const {\n            modelName = \"all-MiniLM-L6-v2\",\n            device = \"cpu\",\n            normalizeEmbeddings = false,\n            kwargs = {},\n        } = args;\n\n        // Validate kwargs are JSON-serializable (no functions or symbols)\n        for (const [key, value] of Object.entries(kwargs)) {\n            if (typeof value === \"function\" || typeof value === \"symbol\") {\n                throw new ChromaValueError(\n                    `Keyword argument '${key}' has a value of type '${typeof value}', which is not supported. Only JSON-serializable values are allowed.`\n                );\n            }\n        }\n\n        this.modelName = modelName;\n        this.device = device;\n        this.normalizeEmbeddings = normalizeEmbeddings;\n        this.kwargs = kwargs;\n    }\n\n    private async getPipeline(): Promise<any> {\n        if (this.pipeline) {\n            return this.pipeline;\n        }\n\n        if (!this.pipelinePromise) {\n            // Resolve model name: if it doesn't contain a '/', prefix with 'Xenova/'\n            // to form a full model identifier for transformers.js\n            // This allows short names like \"all-MiniLM-L6-v2\" to work while maintaining\n            // compatibility with Python client which uses short names\n            let resolvedModelName = this.modelName;\n            if (!resolvedModelName.includes(\"/\")) {\n                resolvedModelName = `Xenova/${resolvedModelName}`;\n            }\n\n            this.pipelinePromise = pipeline(\n                \"feature-extraction\",\n                resolvedModelName,\n                {\n                    device: this.device as any,\n                    ...this.kwargs,\n                } as any\n            ).catch((error) => {\n                // Reset pipelinePromise on error to allow retry on next call\n                this.pipelinePromise = null;\n                throw error;\n            });\n        }\n\n        this.pipeline = await this.pipelinePromise;\n        return this.pipeline;\n    }\n\n    public async generate(texts: string[]): Promise<number[][]> {\n        if (!texts || texts.length === 0) {\n            return [];\n        }\n\n        const pipe = await this.getPipeline();\n\n        // Process all texts in batch\n        const output = await pipe(texts, {\n            pooling: \"mean\",\n            normalize: this.normalizeEmbeddings,\n        });\n\n        // Convert tensor output to JavaScript array\n        return output.tolist();\n    }\n\n    public async dispose(): Promise<void> {\n        // To avoid race conditions, we capture the promise and then nullify the\n        // instance properties to prevent new operations from starting.\n        const promiseToDispose = this.pipelinePromise;\n        this.pipeline = null;\n        this.pipelinePromise = null;\n\n        if (!promiseToDispose) return;\n\n        try {\n            // If pipeline is already initialized, this will resolve immediately.\n            // Otherwise, it will wait for initialization to complete.\n            const pipeline = await promiseToDispose;\n            if (pipeline && typeof pipeline.dispose === \"function\") {\n                await pipeline.dispose();\n            }\n        } catch {\n            // If the pipeline promise fails, there's nothing to dispose.\n            // This error will be handled by callers of generate(), so we can ignore it here.\n        }\n    }\n\n    public defaultSpace(): EmbeddingFunctionSpace {\n        return \"cosine\";\n    }\n\n    public supportedSpaces(): EmbeddingFunctionSpace[] {\n        return [\"cosine\", \"l2\", \"ip\"];\n    }\n\n    public static buildFromConfig(\n        config: SentenceTransformersConfig,\n    ): SentenceTransformersEmbeddingFunction {\n        const { model_name, device, normalize_embeddings, kwargs } = config;\n\n        if (model_name === undefined || device === undefined || normalize_embeddings === undefined) {\n            throw new ChromaValueError(\"model_name, device, and normalize_embeddings are required\");\n        }\n\n        return new SentenceTransformersEmbeddingFunction({\n            modelName: model_name,\n            device,\n            normalizeEmbeddings: normalize_embeddings,\n            kwargs: kwargs || {},\n        });\n    }\n\n    public getConfig(): SentenceTransformersConfig {\n        return {\n            model_name: this.modelName,\n            device: this.device,\n            normalize_embeddings: this.normalizeEmbeddings,\n            kwargs: this.kwargs,\n        };\n    }\n\n    public validateConfigUpdate(newConfig: Record<string, any>): void {\n        // Model name is also used as the identifier for model path if stored locally.\n        // Users should be able to change the path if needed, so we should not validate that.\n        // e.g. moving file path from /v1/my-model.bin to /v2/my-model.bin\n        return;\n    }\n\n    public static validateConfig(config: SentenceTransformersConfig): void {\n        validateConfigSchema(config, \"sentence-transformer\");\n    }\n}\n\n// Register with both the Python name (sentence_transformer) and the mapped JS name (sentence-transformer)\nregisterEmbeddingFunction(NAME, SentenceTransformersEmbeddingFunction);\nregisterEmbeddingFunction(\"sentence-transformer\", SentenceTransformersEmbeddingFunction);\n"],"mappings":";AAAA;AAAA,EACI;AAAA,EAGA;AAAA,OACG;AACP;AAAA,EACI;AAAA,OACG;AACP,SAAS,gBAAgB;AAEzB,IAAM,OAAO;AAgBN,IAAM,wCAAN,MAAM,uCACoB;AAAA,EAS7B,YAAY,OAAiC,CAAC,GAAG;AARjD,SAAgB,OAAO;AAKvB,SAAQ,kBAAuC;AAC/C,SAAQ,WAAgB;AAGpB,UAAM;AAAA,MACF,YAAY;AAAA,MACZ,SAAS;AAAA,MACT,sBAAsB;AAAA,MACtB,SAAS,CAAC;AAAA,IACd,IAAI;AAGJ,eAAW,CAAC,KAAK,KAAK,KAAK,OAAO,QAAQ,MAAM,GAAG;AAC/C,UAAI,OAAO,UAAU,cAAc,OAAO,UAAU,UAAU;AAC1D,cAAM,IAAI;AAAA,UACN,qBAAqB,GAAG,0BAA0B,OAAO,KAAK;AAAA,QAClE;AAAA,MACJ;AAAA,IACJ;AAEA,SAAK,YAAY;AACjB,SAAK,SAAS;AACd,SAAK,sBAAsB;AAC3B,SAAK,SAAS;AAAA,EAClB;AAAA,EAEA,MAAc,cAA4B;AACtC,QAAI,KAAK,UAAU;AACf,aAAO,KAAK;AAAA,IAChB;AAEA,QAAI,CAAC,KAAK,iBAAiB;AAKvB,UAAI,oBAAoB,KAAK;AAC7B,UAAI,CAAC,kBAAkB,SAAS,GAAG,GAAG;AAClC,4BAAoB,UAAU,iBAAiB;AAAA,MACnD;AAEA,WAAK,kBAAkB;AAAA,QACnB;AAAA,QACA;AAAA,QACA;AAAA,UACI,QAAQ,KAAK;AAAA,UACb,GAAG,KAAK;AAAA,QACZ;AAAA,MACJ,EAAE,MAAM,CAAC,UAAU;AAEf,aAAK,kBAAkB;AACvB,cAAM;AAAA,MACV,CAAC;AAAA,IACL;AAEA,SAAK,WAAW,MAAM,KAAK;AAC3B,WAAO,KAAK;AAAA,EAChB;AAAA,EAEA,MAAa,SAAS,OAAsC;AACxD,QAAI,CAAC,SAAS,MAAM,WAAW,GAAG;AAC9B,aAAO,CAAC;AAAA,IACZ;AAEA,UAAM,OAAO,MAAM,KAAK,YAAY;AAGpC,UAAM,SAAS,MAAM,KAAK,OAAO;AAAA,MAC7B,SAAS;AAAA,MACT,WAAW,KAAK;AAAA,IACpB,CAAC;AAGD,WAAO,OAAO,OAAO;AAAA,EACzB;AAAA,EAEA,MAAa,UAAyB;AAGlC,UAAM,mBAAmB,KAAK;AAC9B,SAAK,WAAW;AAChB,SAAK,kBAAkB;AAEvB,QAAI,CAAC,iBAAkB;AAEvB,QAAI;AAGA,YAAMA,YAAW,MAAM;AACvB,UAAIA,aAAY,OAAOA,UAAS,YAAY,YAAY;AACpD,cAAMA,UAAS,QAAQ;AAAA,MAC3B;AAAA,IACJ,QAAQ;AAAA,IAGR;AAAA,EACJ;AAAA,EAEO,eAAuC;AAC1C,WAAO;AAAA,EACX;AAAA,EAEO,kBAA4C;AAC/C,WAAO,CAAC,UAAU,MAAM,IAAI;AAAA,EAChC;AAAA,EAEA,OAAc,gBACV,QACqC;AACrC,UAAM,EAAE,YAAY,QAAQ,sBAAsB,OAAO,IAAI;AAE7D,QAAI,eAAe,UAAa,WAAW,UAAa,yBAAyB,QAAW;AACxF,YAAM,IAAI,iBAAiB,2DAA2D;AAAA,IAC1F;AAEA,WAAO,IAAI,uCAAsC;AAAA,MAC7C,WAAW;AAAA,MACX;AAAA,MACA,qBAAqB;AAAA,MACrB,QAAQ,UAAU,CAAC;AAAA,IACvB,CAAC;AAAA,EACL;AAAA,EAEO,YAAwC;AAC3C,WAAO;AAAA,MACH,YAAY,KAAK;AAAA,MACjB,QAAQ,KAAK;AAAA,MACb,sBAAsB,KAAK;AAAA,MAC3B,QAAQ,KAAK;AAAA,IACjB;AAAA,EACJ;AAAA,EAEO,qBAAqB,WAAsC;AAI9D;AAAA,EACJ;AAAA,EAEA,OAAc,eAAe,QAA0C;AACnE,yBAAqB,QAAQ,sBAAsB;AAAA,EACvD;AACJ;AAGA,0BAA0B,MAAM,qCAAqC;AACrE,0BAA0B,wBAAwB,qCAAqC;","names":["pipeline"]}