UNPKG

@bdmarvin/mcp-server-memory

Version:

MCP Server for LLM Long-Term Memory using KG and Google Drive

503 lines 28.1 kB
import crypto from 'crypto'; import { Mutex } from 'async-mutex'; import { google } from 'googleapis'; import { Readable } from 'stream'; import { findOrCreateProjectFolder as getProjectDriveFolderIdFromDriveService, getOauth2ClientInstance } from './driveService.js'; const KG_FILE_NAME = "knowledge_graph.json"; const projectLocks = new Map(); // Simple console logging replacements const log = { info: (...args) => console.error("INFO:", ...args.map(arg => typeof arg === 'object' ? JSON.stringify(arg) : arg)), debug: (...args) => { }, warn: (...args) => console.error("WARN:", ...args.map(arg => typeof arg === 'object' ? JSON.stringify(arg) : arg)), error: (...args) => console.error("ERROR:", ...args.map(arg => typeof arg === 'object' ? JSON.stringify(arg) : arg)), }; function getProjectLock(projectId) { if (!projectLocks.has(projectId)) { projectLocks.set(projectId, new Mutex()); } return projectLocks.get(projectId); } async function getDriveClientForKg(accessToken) { const oauth2Client = await getOauth2ClientInstance(accessToken); return google.drive({ version: 'v3', auth: oauth2Client }); } async function readKgFromDrive(accessToken, projectId) { const drive = await getDriveClientForKg(accessToken); const projectFolderId = await getProjectDriveFolderIdFromDriveService(drive, projectId); log.debug(`Attempting to read KG file from Drive. ProjectId: ${projectId}, FolderId: ${projectFolderId}, FileName: ${KG_FILE_NAME}`); try { const fileListResponse = await drive.files.list({ q: `name='${KG_FILE_NAME}' and '${projectFolderId}' in parents and trashed=false`, fields: 'files(id, name)', spaces: 'drive', }); if (fileListResponse.data.files && fileListResponse.data.files.length > 0 && fileListResponse.data.files[0].id) { const fileId = fileListResponse.data.files[0].id; log.info(`Found KG file in Drive, downloading. ProjectId: ${projectId}, FileId: ${fileId}`); const streamResponse = await drive.files.get({ fileId: fileId, alt: 'media' }, { responseType: 'stream' }); const chunks = []; for await (const chunk of streamResponse.data) { chunks.push(Buffer.from(chunk)); } return JSON.parse(Buffer.concat(chunks).toString('utf-8')); } else { log.info(`KG file not found in Drive, returning new structure. ProjectId: ${projectId}, FolderId: ${projectFolderId}`); return { nodes: {}, relationships: [], metadata: { project_id: projectId, created_at: new Date().toISOString(), updated_at: new Date().toISOString() } }; } } catch (error) { const errorReason = error.errors && error.errors.length > 0 ? error.errors[0]?.reason : 'unknown_reason'; if (error.code === 404 || errorReason === 'notFound') { log.info(`KG file not found in Drive (404 error), returning new structure. ProjectId: ${projectId}, FolderId: ${projectFolderId}, Reason: ${errorReason}`); return { nodes: {}, relationships: [], metadata: { project_id: projectId, created_at: new Date().toISOString(), updated_at: new Date().toISOString() } }; } log.error(`Error reading KG file from Drive. ProjectId: ${projectId}, FolderId: ${projectFolderId}, Error: ${error.message}`); throw error; } } async function writeKgToDrive(accessToken, projectId, data) { const drive = await getDriveClientForKg(accessToken); const projectFolderId = await getProjectDriveFolderIdFromDriveService(drive, projectId); data.metadata = data.metadata || { project_id: projectId }; data.metadata.updated_at = new Date().toISOString(); if (!data.metadata.created_at) { data.metadata.created_at = new Date().toISOString(); } const kgFileContent = JSON.stringify(data, null, 2); const media = { mimeType: 'application/json', body: Readable.from([kgFileContent]) }; log.debug(`Attempting to write KG file to Drive. ProjectId: ${projectId}, FolderId: ${projectFolderId}, FileName: ${KG_FILE_NAME}`); try { const fileListResponse = await drive.files.list({ q: `name='${KG_FILE_NAME}' and '${projectFolderId}' in parents and trashed=false`, fields: 'files(id)', spaces: 'drive', }); if (fileListResponse.data.files && fileListResponse.data.files.length > 0 && fileListResponse.data.files[0].id) { const fileId = fileListResponse.data.files[0].id; log.info(`Updating existing KG file in Drive. ProjectId: ${projectId}, FileId: ${fileId}`); await drive.files.update({ fileId: fileId, media: media, requestBody: { name: KG_FILE_NAME, mimeType: 'application/json' } }); } else { log.info(`Creating new KG file in Drive. ProjectId: ${projectId}`); await drive.files.create({ requestBody: { name: KG_FILE_NAME, parents: [projectFolderId], mimeType: 'application/json' }, media: media, }); } log.info(`Successfully wrote KG file to Drive. ProjectId: ${projectId}, FolderId: ${projectFolderId}`); } catch (error) { log.error(`Error writing KG file to Drive. ProjectId: ${projectId}, FolderId: ${projectFolderId}, Error: ${error.message}`); throw error; } } export async function updateKgNode(accessToken, args) { const lock = getProjectLock(args.project_id); return await lock.runExclusive(async () => { log.info(`updateKgNode (Drive) called. ProjectId: ${args.project_id}, NodeId: ${args.node_id}`); const kg = await readKgFromDrive(accessToken, args.project_id); kg.nodes = kg.nodes || {}; const existingNode = kg.nodes[args.node_id] || {}; kg.nodes[args.node_id] = { ...existingNode, ...args.attributes, node_id: args.node_id, updated_at: new Date().toISOString(), }; if (!existingNode.created_at) { kg.nodes[args.node_id].created_at = new Date().toISOString(); } await writeKgToDrive(accessToken, args.project_id, kg); log.info(`updateKgNode (Drive) successful. ProjectId: ${args.project_id}, NodeId: ${args.node_id}`); return { status: 'success', node_id: args.node_id, data: kg.nodes[args.node_id] }; }); } export async function addKgRelationship(accessToken, args) { const lock = getProjectLock(args.project_id); return await lock.runExclusive(async () => { log.info(`addKgRelationship (Drive) called. Args: ${JSON.stringify(args)}`); const kg = await readKgFromDrive(accessToken, args.project_id); kg.nodes = kg.nodes || {}; if (!kg.nodes[args.source_node_id]) { throw new Error(`Source node ${args.source_node_id} not found.`); } if (!kg.nodes[args.target_node_id]) { throw new Error(`Target node ${args.target_node_id} not found.`); } kg.relationships = kg.relationships || []; const relationshipId = crypto.randomUUID(); const newRelationship = { relationship_id: relationshipId, source_node_id: args.source_node_id, target_node_id: args.target_node_id, type: args.relationship_type, attributes: args.attributes || {}, created_at: new Date().toISOString(), }; kg.relationships.push(newRelationship); await writeKgToDrive(accessToken, args.project_id, kg); log.info(`addKgRelationship (Drive) successful. ProjectId: ${args.project_id}, RelationshipId: ${relationshipId}`); return { status: 'success', relationship: newRelationship }; }); } export async function logDecision(accessToken, args) { const lock = getProjectLock(args.project_id); return await lock.runExclusive(async () => { log.info(`logDecision (Drive) called. Args: ${JSON.stringify(args)}`); const decisionNodeId = `decision_${crypto.randomUUID()}`; const decisionTimestamp = new Date().toISOString(); const kg = await readKgFromDrive(accessToken, args.project_id); kg.nodes = kg.nodes || {}; kg.relationships = kg.relationships || []; kg.nodes[decisionNodeId] = { node_id: decisionNodeId, type: 'decision', summary: args.decision_summary, rationale: args.rationale, status: args.status, entities_involved_count: args.entities_involved.length, decision_timestamp: decisionTimestamp, created_at: decisionTimestamp, updated_at: decisionTimestamp, }; const relationshipsAddedDetails = []; for (const entityId of args.entities_involved) { if (!kg.nodes[entityId]) { log.warn(`Entity node for decision not found. Skipping relationship. EntityId: ${entityId}, DecisionNodeId: ${decisionNodeId}, ProjectId: ${args.project_id}`); continue; } const relationshipId = crypto.randomUUID(); const rel = { relationship_id: relationshipId, source_node_id: decisionNodeId, type: 'concerns_entity', target_node_id: entityId, attributes: {}, created_at: decisionTimestamp, }; kg.relationships.push(rel); relationshipsAddedDetails.push(rel); } await writeKgToDrive(accessToken, args.project_id, kg); log.info(`logDecision (Drive) successful. ProjectId: ${args.project_id}, DecisionNodeId: ${decisionNodeId}`); return { status: 'success', decision_id: decisionNodeId, decision_node: kg.nodes[decisionNodeId], relationships_added: relationshipsAddedDetails }; }); } export async function getKgNodeDetails(accessToken, args) { log.info(`getKgNodeDetails (Drive) called. Args: ${JSON.stringify(args)}`); const kg = await readKgFromDrive(accessToken, args.project_id); const node = kg.nodes?.[args.node_id]; if (!node) { throw new Error(`Node with ID '${args.node_id}' not found in project '${args.project_id}'.`); } const relatedRelationships = (kg.relationships || []).filter((r) => r.source_node_id === args.node_id || r.target_node_id === args.node_id); return { ...node, related_relationships: relatedRelationships }; } export async function getProjectSummary(accessToken, args) { log.info(`getProjectSummary (Drive) called. Args: ${JSON.stringify(args)}`); const kg = await readKgFromDrive(accessToken, args.project_id); const nodeCount = Object.keys(kg.nodes || {}).length; const relationshipCount = (kg.relationships || []).length; const nodeTypesCount = {}; for (const nodeId in (kg.nodes || {})) { const node = kg.nodes[nodeId]; const nodeType = node.type || 'undefined'; nodeTypesCount[nodeType] = (nodeTypesCount[nodeType] || 0) + 1; } return { project_id: args.project_id, metadata: kg.metadata || { project_id: args.project_id, updated_at: 'N/A', created_at: 'N/A' }, node_count: nodeCount, relationship_count: relationshipCount, node_types_count: nodeTypesCount, project_metadata_nodes: Object.values(kg.nodes || {}).filter((n) => n.type === 'project_metadata'), }; } export async function searchKg(accessToken, args) { log.info(`searchKg (Drive) called. Args: ${JSON.stringify(args)}`); const projectId = args.project_id; if (!projectId) { log.warn("searchKg called without project_id. Not implemented for cross-project."); return { results: [], message: "searchKg requires a project_id. Cross-project search not implemented." }; } const kg = await readKgFromDrive(accessToken, projectId); let results = []; const queryLower = args.query_description.toLowerCase(); for (const nodeId in (kg.nodes || {})) { if (results.length >= (args.max_results || 10)) break; const node = kg.nodes[nodeId]; let relevanceScore = 0; if (args.entity_types && args.entity_types.length > 0) { if (!node.type || !args.entity_types.includes(node.type)) { continue; } relevanceScore += 10; } for (const key in node) { if (typeof node[key] === 'string' && node[key].toLowerCase().includes(queryLower)) { relevanceScore += 5; if (key === 'name' || key === 'title' || key === 'summary') relevanceScore += 5; } } if (relevanceScore > 0) { results.push({ project_id: projectId, node_id: nodeId, data: node, score: relevanceScore }); } } results.sort((a, b) => b.score - a.score); log.info(`searchKg (Drive) completed. Args: ${JSON.stringify(args)}, ResultCount: ${results.length}`); return { results: results.slice(0, args.max_results || 10) }; } export async function retrieveKg(accessToken, args) { log.info(`retrieveKg (Drive) called. ProjectId: ${args.project_id}`); const kg = await readKgFromDrive(accessToken, args.project_id); return kg; } export async function traverseKg(accessToken, args) { log.info(`traverseKg (Drive) called. Args: ${JSON.stringify(args)}`); const kg = await readKgFromDrive(accessToken, args.project_id); const { start_node_id, relationship_types = [], direction = 'outgoing', max_depth = 1, include_intermediate_nodes = false, filter_target_node_types = [], max_results = 25, } = args; if (!kg.nodes || !kg.nodes[start_node_id]) { throw new Error(`Start node with ID '${start_node_id}' not found in project '${args.project_id}'.`); } const paths = []; const visitedRelationships = new Set(); const queue = [[[kg.nodes[start_node_id]], 0]]; while (queue.length > 0) { const [currentPath, currentDepth] = queue.shift(); const currentNode = currentPath[currentPath.length - 1]; if (paths.length >= max_results) break; if (currentDepth >= max_depth) { if (include_intermediate_nodes) { const terminalNodeOfPath = currentPath[currentPath.length - 1]; if (filter_target_node_types.length === 0 || (terminalNodeOfPath.type && filter_target_node_types.includes(terminalNodeOfPath.type))) { paths.push(currentPath); } } else { if (filter_target_node_types.length === 0 || (currentNode.type && filter_target_node_types.includes(currentNode.type))) { if (!paths.some(p => p.length === 1 && p[0].node_id === currentNode.node_id)) { paths.push([currentNode]); } } } continue; } const relationships = kg.relationships || []; for (const rel of relationships) { if (paths.length >= max_results && include_intermediate_nodes) break; if (rel.relationship_id && visitedRelationships.has(rel.relationship_id) && include_intermediate_nodes) continue; let nextNodeId = null; let isValidConnection = false; if ((direction === 'outgoing' || direction === 'both') && rel.source_node_id === currentNode.node_id) { isValidConnection = true; nextNodeId = rel.target_node_id; } if (!isValidConnection && (direction === 'incoming' || direction === 'both') && rel.target_node_id === currentNode.node_id) { isValidConnection = true; nextNodeId = rel.source_node_id; } if (isValidConnection && nextNodeId && kg.nodes[nextNodeId]) { if (relationship_types.length > 0 && !relationship_types.includes(rel.type)) { continue; } const nextNode = kg.nodes[nextNodeId]; const newPath = [...currentPath, rel, nextNode]; if (include_intermediate_nodes && rel.relationship_id) { visitedRelationships.add(rel.relationship_id); } queue.push([newPath, currentDepth + 1]); if (include_intermediate_nodes && (currentDepth + 1 === max_depth)) { if (filter_target_node_types.length === 0 || (nextNode.type && filter_target_node_types.includes(nextNode.type))) { paths.push(newPath); if (paths.length >= max_results) break; } } } } } if (!include_intermediate_nodes) { const uniqueTerminalNodes = Array.from(new Map(paths.map(p => [p[0].node_id, p[0]])).values()); return { project_id: args.project_id, start_node_id: start_node_id, terminal_nodes: uniqueTerminalNodes.slice(0, args.max_results) }; } return { project_id: args.project_id, start_node_id: start_node_id, paths: paths.slice(0, args.max_results) }; } export async function queryKgByAttributes(accessToken, args) { log.info(`queryKgByAttributes (Drive) called. Args: ${JSON.stringify(args)}`); const kg = await readKgFromDrive(accessToken, args.project_id); const { node_type_filter, attribute_filters, logical_operator_for_filters = 'AND', max_results = 50, } = args; const matchingNodes = []; const nodesToSearch = Object.values(kg.nodes || {}); for (const node of nodesToSearch) { if (matchingNodes.length >= max_results) break; if (node_type_filter && node.type !== node_type_filter) { continue; } let filterBlockResult = (logical_operator_for_filters === 'AND'); for (const filter of attribute_filters) { let attributeValue; if (filter.attribute_name.includes('.')) { const parts = filter.attribute_name.split('.'); if (node[parts[0]] && typeof node[parts[0]] === 'object' && node[parts[0]] !== null) { attributeValue = node[parts[0]][parts[1]]; } else { attributeValue = undefined; } } else { attributeValue = node[filter.attribute_name]; } let currentFilterSatisfied = false; const filterValue = filter.value; const caseSensitive = filter.case_sensitive === true; switch (filter.operator) { case 'exists': currentFilterSatisfied = attributeValue !== undefined && attributeValue !== null; break; case 'not_exists': currentFilterSatisfied = attributeValue === undefined || attributeValue === null; break; case 'equals': if (typeof attributeValue === 'string' && typeof filterValue === 'string') { currentFilterSatisfied = caseSensitive ? (attributeValue === filterValue) : (attributeValue.toLowerCase() === filterValue.toLowerCase()); } else if (((typeof attributeValue === 'number' && typeof filterValue === 'number') || (typeof attributeValue === 'boolean' && typeof filterValue === 'boolean'))) { currentFilterSatisfied = attributeValue === filterValue; } else if (attributeValue === null && filterValue === null) { currentFilterSatisfied = true; } break; case 'not_equals': if (typeof attributeValue === 'string' && typeof filterValue === 'string') { currentFilterSatisfied = caseSensitive ? (attributeValue !== filterValue) : (attributeValue.toLowerCase() !== filterValue.toLowerCase()); } else if (((typeof attributeValue === 'number' && typeof filterValue === 'number') || (typeof attributeValue === 'boolean' && typeof filterValue === 'boolean'))) { currentFilterSatisfied = attributeValue !== filterValue; } else if (attributeValue === null && filterValue === null) { currentFilterSatisfied = false; } else if (attributeValue === undefined || attributeValue === null) { currentFilterSatisfied = true; } break; case 'contains': if (typeof attributeValue === 'string' && typeof filterValue === 'string') { currentFilterSatisfied = caseSensitive ? attributeValue.includes(filterValue) : attributeValue.toLowerCase().includes(filterValue.toLowerCase()); } break; case 'not_contains': if (typeof attributeValue === 'string' && typeof filterValue === 'string') { currentFilterSatisfied = caseSensitive ? !attributeValue.includes(filterValue) : !attributeValue.toLowerCase().includes(filterValue.toLowerCase()); } else if (attributeValue === undefined || attributeValue === null) { currentFilterSatisfied = true; } break; case 'startswith': if (typeof attributeValue === 'string' && typeof filterValue === 'string') { currentFilterSatisfied = caseSensitive ? attributeValue.startsWith(filterValue) : attributeValue.toLowerCase().startsWith(filterValue.toLowerCase()); } break; case 'endswith': if (typeof attributeValue === 'string' && typeof filterValue === 'string') { currentFilterSatisfied = caseSensitive ? attributeValue.endsWith(filterValue) : attributeValue.toLowerCase().endsWith(filterValue.toLowerCase()); } break; case 'gt': currentFilterSatisfied = typeof attributeValue === typeof filterValue && attributeValue > filterValue; break; case 'lt': currentFilterSatisfied = typeof attributeValue === typeof filterValue && attributeValue < filterValue; break; case 'gte': currentFilterSatisfied = typeof attributeValue === typeof filterValue && attributeValue >= filterValue; break; case 'lte': currentFilterSatisfied = typeof attributeValue === typeof filterValue && attributeValue <= filterValue; break; case 'in_array': if (Array.isArray(filterValue) && (typeof attributeValue === 'string' || typeof attributeValue === 'number' || typeof attributeValue === 'boolean')) { currentFilterSatisfied = filterValue.includes(attributeValue); } else if (Array.isArray(attributeValue) && (typeof filterValue === 'string' || typeof filterValue === 'number' || typeof filterValue === 'boolean')) { currentFilterSatisfied = attributeValue.includes(filterValue); } else if (Array.isArray(attributeValue) && Array.isArray(filterValue)) { currentFilterSatisfied = attributeValue.some(item => filterValue.includes(item)); } break; case 'not_in_array': if (Array.isArray(filterValue) && (typeof attributeValue === 'string' || typeof attributeValue === 'number' || typeof attributeValue === 'boolean')) { currentFilterSatisfied = !filterValue.includes(attributeValue); } else if (Array.isArray(attributeValue) && (typeof filterValue === 'string' || typeof filterValue === 'number' || typeof filterValue === 'boolean')) { currentFilterSatisfied = !attributeValue.includes(filterValue); } else if (Array.isArray(attributeValue) && Array.isArray(filterValue)) { currentFilterSatisfied = !attributeValue.some(item => filterValue.includes(item)); } else if (attributeValue === undefined || attributeValue === null) { currentFilterSatisfied = true; } break; default: log.warn(`Unsupported operator: ${filter.operator}`); } if (logical_operator_for_filters === 'AND') { if (!currentFilterSatisfied) { filterBlockResult = false; break; } } else { // OR if (currentFilterSatisfied) { filterBlockResult = true; break; } } } if (filterBlockResult) { matchingNodes.push(node); } } log.info(`queryKgByAttributes (Drive) completed. Args: ${JSON.stringify(args)}, ResultCount: ${matchingNodes.length}`); return { results: matchingNodes.slice(0, max_results) }; } export async function deleteKgNode(accessToken, args) { const lock = getProjectLock(args.project_id); return await lock.runExclusive(async () => { log.info(`deleteKgNode (Drive) called. ProjectId: ${args.project_id}, NodeId: ${args.node_id}`); const kg = await readKgFromDrive(accessToken, args.project_id); if (!kg.nodes[args.node_id]) { log.warn(`Node with ID '${args.node_id}' not found in project '${args.project_id}'. Cannot delete.`); return { status: 'not_found', node_id: args.node_id, message: `Node with ID '${args.node_id}' not found.` }; } delete kg.nodes[args.node_id]; const initialRelationshipCount = kg.relationships.length; kg.relationships = kg.relationships.filter(rel => rel.source_node_id !== args.node_id && rel.target_node_id !== args.node_id); const relationshipsDeletedCount = initialRelationshipCount - kg.relationships.length; await writeKgToDrive(accessToken, args.project_id, kg); log.info(`deleteKgNode (Drive) successful. ProjectId: ${args.project_id}, NodeId: ${args.node_id}, RelationshipsDeleted: ${relationshipsDeletedCount}`); return { status: 'success', node_id: args.node_id, relationships_deleted_count: relationshipsDeletedCount }; }); } export async function deleteKgRelationship(accessToken, args) { const lock = getProjectLock(args.project_id); return await lock.runExclusive(async () => { log.info(`deleteKgRelationship (Drive) called. ProjectId: ${args.project_id}, RelationshipId: ${args.relationship_id}`); const kg = await readKgFromDrive(accessToken, args.project_id); const initialRelationshipCount = kg.relationships.length; kg.relationships = kg.relationships.filter(rel => rel.relationship_id !== args.relationship_id); const deleted = initialRelationshipCount > kg.relationships.length; if (!deleted) { log.warn(`Relationship with ID '${args.relationship_id}' not found in project '${args.project_id}'. Cannot delete.`); return { status: 'not_found', relationship_id: args.relationship_id, message: `Relationship with ID '${args.relationship_id}' not found.` }; } await writeKgToDrive(accessToken, args.project_id, kg); log.info(`deleteKgRelationship (Drive) successful. ProjectId: ${args.project_id}, RelationshipId: ${args.relationship_id}`); return { status: 'success', relationship_id: args.relationship_id }; }); } //# sourceMappingURL=kgService.js.map