PoC: Refactor ChromaManager to a Neo.mjs Class #7537

Author: tobiu

ai/mcp/server/memory-core/mcp-stdio.mjs

@@ -1,5 +1,9 @@
-import { Server } from '@modelcontextprotocol/sdk/server/index.js';
-import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
+import {Server}               from '@modelcontextprotocol/sdk/server/index.js';
+import {StdioServerTransport} from '@modelcontextprotocol/sdk/server/stdio.js';
+import Neo                    from '../../../../src/Neo.mjs';
+import * as core              from '../../../../src/core/_export.mjs';
+import InstanceManager        from '../../../../src/manager/Instance.mjs';
+
 import {
     CallToolRequestSchema,
     ListToolsRequestSchema,

ai/mcp/server/memory-core/services/ChromaManager.mjs

@@ -0,0 +1,141 @@
+import {ChromaClient} from 'chromadb';
+import aiConfig       from '../../config.mjs';
+import Base           from '../../../../../src/core/Base.mjs';
+
+/**
+ * Simple manager around the Chroma client that lazily caches frequently used collections.
+ * @class AI.mcp.server.memory.ChromaManager
+ * @extends Neo.core.Base
+ * @singleton
+ */
+class ChromaManager extends Base {
+    static config = {
+        /**
+         * @member {String} className='AI.mcp.server.memory.ChromaManager'
+         * @protected
+         */
+        className: 'AI.mcp.server.memory.ChromaManager',
+        /**
+         * @member {ChromaClient|null} client_=null
+         * @protected
+         * @reactive
+         */
+        client_: null,
+        /**
+         * @member {Boolean} connected_=false
+         * @reactive
+         */
+        connected_: false,
+        /**
+         * @member {import('chromadb').Collection|null} memoryCollection_=null
+         * @protected
+         * @reactive
+         */
+        memoryCollection_: null,
+        /**
+         * @member {Boolean} singleton=true
+         * @protected
+         */
+        singleton: true,
+        /**
+         * @member {import('chromadb').Collection|null} summaryCollection_=null
+         * @protected
+         * @reactive
+         */
+        summaryCollection_: null
+    }
+
+    /**
+     * Triggered after the connected config got changed
+     * @param {Boolean} value
+     * @param {Boolean} oldValue
+     * @protected
+     */
+    afterSetConnected(value, oldValue) {
+        // console.error(`[ChromaManager] ${value ? 'Connected' : 'Disconnected'}`); // Use logger instead
+        // this.fire?.('connectionChange', { connected: value }); // Not part of PoC yet
+    }
+
+    /**
+     *
+     */
+    onConstructed() {
+        super.onConstructed();
+
+        // The client is created here, but the connection is established in initAsync
+        const {host, port} = aiConfig.memory;
+        this.client = new ChromaClient({host, port, ssl: false});
+    }
+
+    /**
+     * @returns {Promise<void>}
+     */
+    async initAsync() {
+        await super.initAsync();
+        await this.connect();
+    }
+
+    /**
+     * Establishes connection to ChromaDB.
+     */
+    async connect() {
+        try {
+            await this.client.heartbeat();
+            this.connected = true;
+        } catch (e) {
+            this.connected = false;
+            throw e;
+        }
+    }
+
+    /**
+     * Ensures the process can reach the Chroma server and both collections are available.
+     * @returns {Promise<{heartbeat: number, memoryCollection: string, summaryCollection: string}>}
+     */
+    async checkConnectivity() {
+        const heartbeat = await this.client.heartbeat();
+
+        const memory    = await this.getMemoryCollection();
+        const summaries = await this.getSummaryCollection();
+
+        return {
+            heartbeat,
+            memoryCollection : memory.name,
+            summaryCollection: summaries.name
+        };
+    }
+
+    /**
+     * @returns {Promise<import('chromadb').Collection>}
+     */
+    async getMemoryCollection() {
+        if (!this.memoryCollection) {
+            const {collectionName} = aiConfig.memory;
+
+            this.memoryCollection = await this.client.getCollection({
+                name             : collectionName,
+                embeddingFunction: aiConfig.dummyEmbeddingFunction
+            });
+        }
+
+        return this.memoryCollection;
+    }
+
+    /**
+     * @returns {Promise<import('chromadb').Collection>}
+     */
+    async getSummaryCollection() {
+        if (!this.summaryCollection) {
+            const {collectionName} = aiConfig.sessions;
+
+            this.summaryCollection = await this.client.getOrCreateCollection({
+                name             : collectionName,
+                embeddingFunction: aiConfig.dummyEmbeddingFunction
+            });
+        }
+
+        return this.summaryCollection;
+    }
+}
+
+export default Neo.setupClass(ChromaManager);