@bdmarvin/mcp-server-memory
Version:
MCP Server for LLM Long-Term Memory using KG and Google Drive
503 lines • 28.1 kB
JavaScript
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