@ever_cheng/memory-task-mcp
Version:
Memory and task management MCP Server
424 lines • 15.2 kB
JavaScript
"use strict";
/**
* Vector Store for MemTask
*
* ChromaDB integration for semantic search and vector similarity operations.
* Handles storage and retrieval of memory chunk embeddings.
*/
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.VectorStore = exports.DEFAULT_VECTOR_STORE_CONFIG = void 0;
const chromadb_1 = require("chromadb");
const logger_1 = __importDefault(require("./logger"));
/**
* Default vector store configuration
*/
exports.DEFAULT_VECTOR_STORE_CONFIG = {
collectionName: 'memtask_memories',
maxResults: 20,
similarityThreshold: 0.5,
distanceMetric: 'cosine'
};
/**
* Vector Store Class
*/
class VectorStore {
constructor(config = exports.DEFAULT_VECTOR_STORE_CONFIG) {
this.client = null;
this.collection = null;
this.initialized = false;
this.config = {
collectionName: config.collectionName,
maxResults: config.maxResults,
similarityThreshold: config.similarityThreshold,
distanceMetric: config.distanceMetric
};
}
/**
* 初始化 ChromaDB 客戶端和集合
*/
async initialize() {
if (this.initialized)
return;
try {
// 創建 ChromaDB 客戶端 - 使用環境變數或默認配置
const chromaUrl = process.env.CHROMADB_URL || 'http://localhost:8000';
logger_1.default.info(`正在連接到 ChromaDB: ${chromaUrl}`);
this.client = new chromadb_1.ChromaClient({
path: chromaUrl
});
// 獲取或創建集合
this.collection = await this.client.getOrCreateCollection({
name: this.config.collectionName,
metadata: {
'hnsw:space': this.config.distanceMetric,
'description': 'MemTask memory chunks for semantic search'
},
embeddingFunction: undefined // 使用自己的 embedding
});
this.initialized = true;
logger_1.default.info(`✅ ChromaDB 初始化成功，集合: ${this.config.collectionName}`, {
collectionName: this.config.collectionName,
distanceMetric: this.config.distanceMetric
});
}
catch (error) {
const chromaUrl = process.env.CHROMADB_URL || 'http://localhost:8000';
logger_1.default.error(`❌ ChromaDB 初始化失敗`, error, { chromaUrl });
throw new Error(`Failed to initialize ChromaDB: ${error}`);
}
}
/**
* 添加單個 memory chunk
*/
async addMemoryChunk(chunk) {
await this.initialize();
if (!this.collection || !chunk.embedding) {
throw new Error('Collection not initialized or chunk missing embedding');
}
try {
const metadata = {
memoryId: chunk.memoryId,
chunkIndex: chunk.chunkIndex,
totalChunks: chunk.totalChunks,
created_at: chunk.metadata.created_at,
originalLength: chunk.metadata.originalLength,
chunkLength: chunk.metadata.chunkLength,
tags: chunk.metadata.tags.join(',') // Convert tags array to string for ChromaDB
};
await this.collection.add({
ids: [chunk.id],
embeddings: [chunk.embedding],
metadatas: [metadata],
documents: [chunk.content]
});
logger_1.default.debug(`✅ Added chunk: ${chunk.id}`, { chunkId: chunk.id, memoryId: chunk.memoryId });
}
catch (error) {
logger_1.default.error(`❌ Failed to add chunk ${chunk.id}`, error, { chunkId: chunk.id });
throw error;
}
}
/**
* 批量添加 memory chunks
*/
async addMemoryChunks(chunks) {
await this.initialize();
if (!this.collection) {
throw new Error('Collection not initialized');
}
let processed = 0;
let failed = 0;
const errors = [];
// 過濾出有 embedding 的 chunks
const validChunks = chunks.filter(chunk => chunk.embedding);
if (validChunks.length === 0) {
return {
success: true,
processed: 0,
failed: 0,
errors: []
};
}
try {
const ids = validChunks.map(chunk => chunk.id);
const embeddings = validChunks.map(chunk => chunk.embedding);
const metadatas = validChunks.map(chunk => ({
memoryId: chunk.memoryId,
chunkIndex: chunk.chunkIndex,
totalChunks: chunk.totalChunks,
tags: chunk.metadata.tags.join(','), // Convert tags array to string for ChromaDB
created_at: chunk.metadata.created_at,
originalLength: chunk.metadata.originalLength,
chunkLength: chunk.metadata.chunkLength
}));
const documents = validChunks.map(chunk => chunk.content);
await this.collection.add({
ids,
embeddings,
metadatas,
documents
});
processed = validChunks.length;
logger_1.default.info(`✅ Batch added ${processed} chunks`, { processed, totalChunks: validChunks.length });
}
catch (error) {
failed = validChunks.length;
errors.push(String(error));
logger_1.default.error(`❌ Batch add failed`, error, { totalChunks: validChunks.length });
}
return {
success: failed === 0,
processed,
failed,
errors
};
}
/**
* 向量相似度搜尋
*/
async searchSimilar(queryEmbedding, limit = this.config.maxResults, filters) {
await this.initialize();
if (!this.collection) {
throw new Error('Collection not initialized');
}
try {
// 構建查詢條件
const whereClause = {};
if (filters?.tags && filters.tags.length > 0) {
whereClause.tags = { '$in': filters.tags };
}
if (filters?.memoryIds && filters.memoryIds.length > 0) {
whereClause.memoryId = { '$in': filters.memoryIds };
}
const queryResults = await this.collection.query({
queryEmbeddings: [queryEmbedding],
nResults: limit,
where: Object.keys(whereClause).length > 0 ? whereClause : undefined,
include: [chromadb_1.IncludeEnum.Distances, chromadb_1.IncludeEnum.Metadatas, chromadb_1.IncludeEnum.Documents]
});
// 轉換結果格式
const results = [];
if (queryResults.ids[0]) {
for (let i = 0; i < queryResults.ids[0].length; i++) {
const distance = queryResults.distances[0][i];
const similarity = this.distanceToSimilarity(distance);
// 過濾低相似度結果
if (similarity >= this.config.similarityThreshold) {
const metadata = queryResults.metadatas[0][i];
results.push({
chunkId: queryResults.ids[0][i],
memoryId: metadata?.memoryId,
similarity,
distance,
metadata: metadata,
content: queryResults.documents[0][i]
});
}
}
}
return results;
}
catch (error) {
logger_1.default.error('Vector search failed', error, { limit, hasFilters: !!filters });
throw error;
}
}
/**
* 根據 memory ID 搜尋相關 chunks
*/
async getChunksByMemoryId(memoryId) {
await this.initialize();
if (!this.collection) {
throw new Error('Collection not initialized');
}
try {
const results = await this.collection.get({
where: { memoryId },
include: [chromadb_1.IncludeEnum.Metadatas, chromadb_1.IncludeEnum.Documents]
});
return results.ids?.map((id, index) => ({
chunkId: id,
memoryId,
similarity: 1.0, // 完全相關
distance: 0,
metadata: results.metadatas[index],
content: results.documents[index]
})) || [];
}
catch (error) {
logger_1.default.error(`Failed to get chunks for memory ${memoryId}`, error, { memoryId });
throw error;
}
}
/**
* 刪除 memory 的所有 chunks
*/
async deleteMemoryChunks(memoryId) {
await this.initialize();
if (!this.collection) {
throw new Error('Collection not initialized');
}
try {
// 先獲取所有相關的 chunk IDs
const existingChunks = await this.collection.get({
where: { memoryId },
include: []
});
if (existingChunks.ids && existingChunks.ids.length > 0) {
await this.collection.delete({
ids: existingChunks.ids
});
logger_1.default.info(`✅ Deleted ${existingChunks.ids.length} chunks for memory ${memoryId}`, { memoryId, deletedCount: existingChunks.ids.length });
return true;
}
return false;
}
catch (error) {
logger_1.default.error(`Failed to delete chunks for memory ${memoryId}`, error, { memoryId });
return false;
}
}
/**
* 更新 memory chunks（先刪除再添加）
*/
async updateMemoryChunks(memoryId, newChunks) {
try {
// 先刪除現有 chunks
await this.deleteMemoryChunks(memoryId);
// 添加新 chunks
const result = await this.addMemoryChunks(newChunks);
return result.success;
}
catch (error) {
logger_1.default.error(`Failed to update chunks for memory ${memoryId}`, error, { memoryId });
return false;
}
}
/**
* 獲取 vector store 統計信息
*/
async getStats() {
await this.initialize();
if (!this.collection) {
throw new Error('Collection not initialized');
}
try {
const result = await this.collection.get({
include: [chromadb_1.IncludeEnum.Metadatas]
});
const totalChunks = result.ids?.length || 0;
const uniqueMemories = new Set(result.metadatas?.map(m => m?.memoryId)).size;
return {
totalChunks,
uniqueMemories,
collectionSize: totalChunks,
avgSimilarityScore: 0, // 需要額外計算
metadata: {
collectionName: this.config.collectionName,
distanceMetric: this.config.distanceMetric
}
};
}
catch (error) {
logger_1.default.error('Failed to get stats', error, { collectionName: this.config.collectionName });
throw error;
}
}
/**
* 清空整個集合
*/
async clearCollection() {
await this.initialize();
if (!this.client) {
throw new Error('Client not initialized');
}
try {
// 刪除集合
await this.client.deleteCollection({
name: this.config.collectionName
});
// 重新創建集合
this.collection = await this.client.getOrCreateCollection({
name: this.config.collectionName,
metadata: {
'hnsw:space': this.config.distanceMetric,
'description': 'MemTask memory chunks for semantic search'
},
embeddingFunction: undefined
});
logger_1.default.info(`✅ Collection ${this.config.collectionName} cleared`, { collectionName: this.config.collectionName });
return true;
}
catch (error) {
logger_1.default.error('Failed to clear collection', error, { collectionName: this.config.collectionName });
return false;
}
}
/**
* 將距離轉換為相似度 (0-1)
*/
distanceToSimilarity(distance) {
if (this.config.distanceMetric === 'cosine') {
// 餘弦距離轉相似度
return 1 - distance;
}
else {
// 歐幾里得距離或曼哈頓距離的簡單轉換
return 1 / (1 + distance);
}
}
/**
* 健康檢查
*/
async healthCheck() {
try {
await this.initialize();
// 測試集合操作
if (this.collection) {
await this.collection.count();
}
return {
status: 'healthy',
initialized: this.initialized,
collectionExists: this.collection !== null
};
}
catch (error) {
return {
status: 'unhealthy',
initialized: this.initialized,
collectionExists: this.collection !== null,
error: String(error)
};
}
}
/**
* 更新配置
*/
updateConfig(newConfig) {
this.config = { ...this.config, ...newConfig };
}
/**
* 獲取當前配置
*/
getConfig() {
return { ...this.config };
}
/**
* 清理資源 (實現 Disposable 介面)
*/
async dispose() {
try {
// 清理 Collection 資源
if (this.collection) {
// ChromaDB collection doesn't need explicit cleanup
this.collection = null;
}
// 清理 Client 資源
if (this.client) {
// ChromaDB client doesn't need explicit cleanup in current version
// but we null it out for memory management
this.client = null;
}
// 重置狀態
this.initialized = false;
logger_1.default.info('✅ VectorStore disposed successfully');
}
catch (error) {
logger_1.default.error('❌ Error disposing VectorStore', error);
throw error;
}
}
/**
* 清理資源 (向後兼容)
* @deprecated Use dispose() instead
*/
async cleanup() {
await this.dispose();
}
}
exports.VectorStore = VectorStore;
//# sourceMappingURL=vector_store.js.map