This is page 2 of 2. Use http://codebase.md/j3k0/mcp-elastic-memory?page={x} to view the full context.
# Directory Structure
```
├── .cursor
│ └── rules
│ └── using-git.mdc
├── .gitignore
├── BUILD-NOTES.md
├── docker-compose.yml
├── Dockerfile
├── jest.config.cjs
├── jest.config.js
├── launch.example
├── legacy
│ ├── cli.ts
│ ├── index.ts
│ ├── query-language.test.ts
│ ├── query-language.ts
│ ├── test-memory.json
│ └── types.ts
├── package.json
├── README.md
├── src
│ ├── admin-cli.ts
│ ├── ai-service.ts
│ ├── es-types.ts
│ ├── filesystem
│ │ └── index.ts
│ ├── index.ts
│ ├── json-to-es.ts
│ ├── kg-client.ts
│ ├── kg-inspection.ts
│ └── logger.ts
├── tests
│ ├── boolean-search.test.ts
│ ├── cross-zone-relations.test.ts
│ ├── empty-name-validation.test.ts
│ ├── entity-type-filtering.test.ts
│ ├── fuzzy-search.test.ts
│ ├── non-existent-entity-relationships.test.ts
│ ├── simple.test.js
│ ├── test-config.ts
│ ├── test-cross-zone.js
│ ├── test-empty-name.js
│ ├── test-non-existent-entity.js
│ ├── test-relationship-cleanup.js
│ ├── test-relevance-score.js
│ └── test-zone-management.js
├── tsconfig.json
└── vitest.config.ts
```
# Files
--------------------------------------------------------------------------------
/src/admin-cli.ts:
--------------------------------------------------------------------------------
```typescript
#!/usr/bin/env node
import { promises as fs } from 'fs';
import path from 'path';
import { KnowledgeGraphClient } from './kg-client.js';
import { ESEntity, ESRelation, KG_RELATIONS_INDEX } from './es-types.js';
import { importFromJsonFile, exportToJsonFile } from './json-to-es.js';
import readline from 'readline';
import GroqAI from './ai-service.js';
// Environment configuration for Elasticsearch
const ES_NODE = process.env.ES_NODE || 'http://localhost:9200';
const ES_USERNAME = process.env.ES_USERNAME;
const ES_PASSWORD = process.env.ES_PASSWORD;
const DEFAULT_ZONE = process.env.KG_DEFAULT_ZONE || 'default';
// Configure ES client with authentication if provided
const esOptions: {
node: string;
auth?: { username: string; password: string };
defaultZone?: string;
} = {
node: ES_NODE,
defaultZone: DEFAULT_ZONE
};
if (ES_USERNAME && ES_PASSWORD) {
esOptions.auth = { username: ES_USERNAME, password: ES_PASSWORD };
}
// Create KG client
const kgClient = new KnowledgeGraphClient(esOptions);
/**
* Display help information
*/
function showHelp() {
console.log('Knowledge Graph Admin CLI');
console.log('========================');
console.log('');
console.log('Commands:');
console.log(' init Initialize the Elasticsearch index');
console.log(' import <file> [zone] Import data from a JSON file (optionally to a specific zone)');
console.log(' export <file> [zone] Export data to a JSON file (optionally from a specific zone)');
console.log(' backup <file> Backup all zones and relations to a file');
console.log(' restore <file> [--yes] Restore all zones and relations from a backup file');
console.log(' stats [zone] Display statistics about the knowledge graph');
console.log(' search <query> [zone] Search the knowledge graph');
console.log(' reset [zone] [--yes] Reset the knowledge graph (delete all data)');
console.log(' entity <n> [zone] Display information about a specific entity');
console.log(' zones list List all memory zones');
console.log(' zones add <name> [desc] Add a new memory zone');
console.log(' zones delete <name> [--yes] Delete a memory zone and all its data');
console.log(' zones stats <name> Show statistics for a specific zone');
console.log(' zones update_descriptions <name> [limit] [prompt]');
console.log(' Generate AI descriptions based on zone content');
console.log(' (limit: optional entity limit, prompt: optional description of zone purpose)');
console.log(' relations <entity> [zone] Show relations for a specific entity');
console.log(' help Show this help information');
console.log('');
console.log('Options:');
console.log(' --yes, -y Automatically confirm all prompts (for scripts)');
console.log('');
console.log('Environment variables:');
console.log(' ES_NODE Elasticsearch node URL (default: http://localhost:9200)');
console.log(' ES_USERNAME Elasticsearch username (if authentication is required)');
console.log(' ES_PASSWORD Elasticsearch password (if authentication is required)');
console.log(' KG_DEFAULT_ZONE Default zone to use (default: "default")');
}
/**
* Initialize the Elasticsearch index
*/
async function initializeIndex() {
try {
await kgClient.initialize();
console.log('Elasticsearch indices initialized successfully');
} catch (error) {
console.error('Error initializing index:', (error as Error).message);
process.exit(1);
}
}
/**
* Display statistics about the knowledge graph or a specific zone
*/
async function showStats(zone?: string) {
try {
// Initialize client
await kgClient.initialize(zone);
if (zone) {
// Get zone-specific stats
const stats = await kgClient.getMemoryZoneStats(zone);
console.log(`Knowledge Graph Statistics for Zone: ${zone}`);
console.log('=============================================');
console.log(`Total entities: ${stats.entityCount}`);
console.log(`Total relations: ${stats.relationCount}`);
console.log('');
console.log('Entity types:');
Object.entries(stats.entityTypes).forEach(([type, count]) => {
console.log(` ${type}: ${count}`);
});
console.log('');
console.log('Relation types:');
Object.entries(stats.relationTypes).forEach(([type, count]) => {
console.log(` ${type}: ${count}`);
});
} else {
// Get all zone metadata
const zones = await kgClient.listMemoryZones();
console.log('Knowledge Graph Multi-zone Statistics');
console.log('====================================');
console.log(`Total zones: ${zones.length}`);
console.log('');
console.log('Zones:');
for (const zone of zones) {
console.log(`Zone: ${zone.name}`);
console.log(` Description: ${zone.description || 'N/A'}`);
console.log(` Created: ${zone.createdAt}`);
console.log(` Last modified: ${zone.lastModified}`);
// Get zone stats
const stats = await kgClient.getMemoryZoneStats(zone.name);
console.log(` Entities: ${stats.entityCount}`);
console.log(` Relations: ${stats.relationCount}`);
console.log('');
}
// Get relation stats
const data = await kgClient.exportData();
const relations = data.filter(item => item.type === 'relation');
console.log(`Total relations in all zones: ${relations.length}`);
// Count relation types
const relationTypes = new Map<string, number>();
relations.forEach(relation => {
const type = (relation as any).relationType;
relationTypes.set(type, (relationTypes.get(type) || 0) + 1);
});
console.log('Relation types:');
relationTypes.forEach((count, type) => {
console.log(` ${type}: ${count}`);
});
}
} catch (error) {
console.error('Error getting statistics:', (error as Error).message);
process.exit(1);
}
}
/**
* Search the knowledge graph
* @param query The search query
* @param zone Optional zone to search in
*/
async function searchGraph(query: string, zone?: string) {
try {
// Initialize client
await kgClient.initialize(zone);
// Search for entities
const results = await kgClient.search({
query,
limit: 10,
sortBy: 'relevance',
zone
});
// Display results
console.log(`Search Results for "${query}"${zone ? ` in zone "${zone}"` : ''}`);
console.log('====================================');
console.log(`Found ${results.hits.total.value} matches`);
console.log('');
// Extract all entities from search results
const entities = results.hits.hits
.filter(hit => hit._source.type === 'entity')
.map(hit => hit._source as ESEntity);
// Get entity names for relation lookup
const entityNames = entities.map(entity => entity.name);
// Create a set of entity names for faster lookup
const entityNameSet = new Set(entityNames);
// Collect all relations
const allRelations: ESRelation[] = [];
const relatedEntities = new Map<string, ESEntity>();
// For each found entity, get all its relations
for (const entityName of entityNames) {
const { relations } = await kgClient.getRelatedEntities(entityName, 1, zone);
for (const relation of relations) {
// Add relation if not already added
if (!allRelations.some(r =>
r.from === relation.from &&
r.fromZone === relation.fromZone &&
r.to === relation.to &&
r.toZone === relation.toZone &&
r.relationType === relation.relationType
)) {
allRelations.push(relation);
// Track related entities that weren't in the search results
// If 'from' entity is not in our set and not already tracked
if (!entityNameSet.has(relation.from) && !relatedEntities.has(relation.from)) {
const entity = await kgClient.getEntityWithoutUpdatingLastRead(relation.from, relation.fromZone);
if (entity) relatedEntities.set(relation.from, entity);
}
// If 'to' entity is not in our set and not already tracked
if (!entityNameSet.has(relation.to) && !relatedEntities.has(relation.to)) {
const entity = await kgClient.getEntityWithoutUpdatingLastRead(relation.to, relation.toZone);
if (entity) relatedEntities.set(relation.to, entity);
}
}
}
}
// Display each entity from search results
entities.forEach((entity, index) => {
const hit = results.hits.hits.find(h =>
h._source.type === 'entity' && (h._source as ESEntity).name === entity.name
);
const score = hit && hit._score !== null && hit._score !== undefined ? hit._score.toFixed(2) : 'N/A';
console.log(`${index + 1}. ${entity.name} (${entity.entityType}) [Score: ${score}]`);
console.log(` Zone: ${entity.zone || 'default'}`);
console.log(` Observations: ${entity.observations.length}`);
// Show highlights if available
if (hit && hit.highlight) {
console.log(' Matches:');
Object.entries(hit.highlight).forEach(([field, highlights]) => {
highlights.forEach(highlight => {
console.log(` - ${field}: ${highlight}`);
});
});
}
console.log('');
});
// Display relations if any
if (allRelations.length > 0) {
console.log('Relations for these entities:');
console.log('====================================');
allRelations.forEach(relation => {
// Lookup entity types for more context
const fromType = entityNameSet.has(relation.from)
? entities.find(e => e.name === relation.from)?.entityType
: relatedEntities.get(relation.from)?.entityType || '?';
const toType = entityNameSet.has(relation.to)
? entities.find(e => e.name === relation.to)?.entityType
: relatedEntities.get(relation.to)?.entityType || '?';
console.log(`${relation.from} [${relation.fromZone}] (${fromType}) → ${relation.relationType} → ${relation.to} [${relation.toZone}] (${toType})`);
});
console.log('');
}
} catch (error) {
console.error('Error searching knowledge graph:', (error as Error).message);
process.exit(1);
}
}
/**
* Reset the knowledge graph (delete all data)
* @param zone Optional zone to reset, if not provided resets all zones
*/
async function resetIndex(zone?: string, args: string[] = []) {
try {
const confirmMessage = zone
? `Are you sure you want to delete all data in zone "${zone}"? This cannot be undone. (y/N) `
: 'Are you sure you want to delete ALL DATA IN ALL ZONES? This cannot be undone. (y/N) ';
const confirmed = await confirmAction(confirmMessage, args);
if (confirmed) {
if (zone) {
// Delete specific zone
if (zone === 'default') {
// For default zone, just delete all entities but keep the index
await kgClient.initialize(zone);
const allEntities = await kgClient.exportData(zone);
for (const item of allEntities) {
if (item.type === 'entity') {
await kgClient.deleteEntity(item.name, zone);
}
}
// Delete relations involving this zone
const client = kgClient['client']; // Access the private client property
await client.deleteByQuery({
index: KG_RELATIONS_INDEX,
body: {
query: {
bool: {
should: [
{ term: { fromZone: zone } },
{ term: { toZone: zone } }
],
minimum_should_match: 1
}
}
},
refresh: true
});
console.log(`Zone "${zone}" has been reset (entities and relations deleted)`);
} else {
// For non-default zones, delete the zone completely
const success = await kgClient.deleteMemoryZone(zone);
if (success) {
console.log(`Zone "${zone}" has been completely deleted`);
} else {
console.error(`Failed to delete zone "${zone}"`);
process.exit(1);
}
}
} else {
// Delete all zones
const zones = await kgClient.listMemoryZones();
// Delete all indices
const client = kgClient['client']; // Access the private client property
for (const zone of zones) {
if (zone.name === 'default') {
// Clear default zone but don't delete it
const indexName = `knowledge-graph@default`;
try {
await client.indices.delete({ index: indexName });
console.log(`Deleted index: ${indexName}`);
} catch (error) {
console.error(`Error deleting index ${indexName}:`, error);
}
} else {
// Delete non-default zones
await kgClient.deleteMemoryZone(zone.name);
console.log(`Deleted zone: ${zone.name}`);
}
}
// Delete relations index
try {
await client.indices.delete({ index: KG_RELATIONS_INDEX });
console.log('Deleted relations index');
} catch (error) {
console.error('Error deleting relations index:', error);
}
// Re-initialize everything
await kgClient.initialize();
console.log('Knowledge graph has been completely reset');
}
} else {
console.log('Operation cancelled');
}
} catch (error) {
console.error('Error resetting index:', (error as Error).message);
process.exit(1);
}
}
/**
* Display information about a specific entity
* @param name Entity name
* @param zone Optional zone name
*/
async function showEntity(name: string, zone?: string) {
try {
// Initialize client
await kgClient.initialize(zone);
// Get entity
const entity = await kgClient.getEntityWithoutUpdatingLastRead(name, zone);
if (!entity) {
console.error(`Entity "${name}" not found${zone ? ` in zone "${zone}"` : ''}`);
process.exit(1);
}
// Get related entities
const related = await kgClient.getRelatedEntities(name, 1, zone);
// Display entity information
console.log(`Entity: ${entity.name}`);
console.log(`Type: ${entity.entityType}`);
console.log(`Zone: ${entity.zone || 'default'}`);
console.log(`Last read: ${entity.lastRead}`);
console.log(`Last write: ${entity.lastWrite}`);
console.log(`Read count: ${entity.readCount}`);
console.log(`Relevance score: ${typeof entity.relevanceScore === 'number' ? entity.relevanceScore.toFixed(2) : '1.00'} (higher = more important)`);
console.log('');
console.log('Observations:');
entity.observations.forEach((obs: string, i: number) => {
console.log(` ${i+1}. ${obs}`);
});
console.log('');
console.log('Relations:');
for (const relation of related.relations) {
if (relation.from === name && relation.fromZone === (entity.zone || 'default')) {
console.log(` → ${relation.relationType} → ${relation.to} [${relation.toZone}]`);
} else {
console.log(` ← ${relation.relationType} ← ${relation.from} [${relation.fromZone}]`);
}
}
} catch (error) {
console.error('Error getting entity:', (error as Error).message);
process.exit(1);
}
}
/**
* List all memory zones
*/
async function listZones() {
try {
await kgClient.initialize();
const zones = await kgClient.listMemoryZones();
console.log('Memory Zones:');
console.log('=============');
for (const zone of zones) {
console.log(`${zone.name}`);
console.log(` Description: ${zone.description || 'N/A'}`);
console.log(` Created: ${zone.createdAt}`);
console.log(` Last modified: ${zone.lastModified}`);
console.log('');
}
console.log(`Total: ${zones.length} zones`);
} catch (error) {
console.error('Error listing zones:', (error as Error).message);
process.exit(1);
}
}
/**
* Add a new memory zone
*/
async function addZone(name: string, description?: string) {
try {
await kgClient.initialize();
await kgClient.addMemoryZone(name, description);
console.log(`Zone "${name}" created successfully`);
} catch (error) {
console.error('Error adding zone:', (error as Error).message);
process.exit(1);
}
}
/**
* Delete a memory zone
*/
async function deleteZone(name: string, args: string[] = []) {
try {
const confirmMessage = `Are you sure you want to delete zone "${name}" and all its data? This cannot be undone. (y/N) `;
const confirmed = await confirmAction(confirmMessage, args);
if (confirmed) {
await kgClient.initialize();
const success = await kgClient.deleteMemoryZone(name);
if (success) {
console.log(`Zone "${name}" deleted successfully`);
} else {
console.error(`Failed to delete zone "${name}"`);
process.exit(1);
}
} else {
console.log('Operation cancelled');
}
} catch (error) {
console.error('Error deleting zone:', (error as Error).message);
process.exit(1);
}
}
/**
* Show relations for a specific entity
*/
async function showRelations(name: string, zone?: string) {
try {
await kgClient.initialize(zone);
// Check if entity exists
const entity = await kgClient.getEntityWithoutUpdatingLastRead(name, zone);
if (!entity) {
console.error(`Entity "${name}" not found${zone ? ` in zone "${zone}"` : ''}`);
process.exit(1);
}
const actualZone = zone || 'default';
// Get all relations for this entity
const { relations } = await kgClient.getRelationsForEntities([name], actualZone);
console.log(`Relations for entity "${name}" in zone "${actualZone}":"`);
console.log('====================================');
if (relations.length === 0) {
console.log('No relations found.');
return;
}
// Group by relation type
const relationsByType = new Map<string, ESRelation[]>();
for (const relation of relations) {
if (!relationsByType.has(relation.relationType)) {
relationsByType.set(relation.relationType, []);
}
relationsByType.get(relation.relationType)!.push(relation);
}
// Display grouped relations
for (const [type, rels] of relationsByType.entries()) {
console.log(`\n${type} (${rels.length}):`);
console.log('----------------');
for (const rel of rels) {
if (rel.from === name && rel.fromZone === actualZone) {
// This entity is the source
console.log(`→ ${rel.to} [${rel.toZone}]`);
} else {
// This entity is the target
console.log(`← ${rel.from} [${rel.fromZone}]`);
}
}
}
} catch (error) {
console.error('Error showing relations:', (error as Error).message);
process.exit(1);
}
}
/**
* Backup all zones to a file
*/
async function backupAll(filePath: string) {
try {
await kgClient.initialize();
// Export all data from all zones
console.log('Exporting all zones and relations...');
const data = await kgClient.exportAllData();
console.log(`Found ${data.entities.length} entities, ${data.relations.length} relations, and ${data.zones.length} zones`);
// Write to file
const jsonData = JSON.stringify(data, null, 2);
await fs.writeFile(filePath, jsonData);
console.log(`Backup saved to ${filePath}`);
console.log(`Entities: ${data.entities.length}`);
console.log(`Relations: ${data.relations.length}`);
console.log(`Zones: ${data.zones.length}`);
} catch (error) {
console.error('Error creating backup:', (error as Error).message);
process.exit(1);
}
}
/**
* Restore all zones from a backup file
*/
async function restoreAll(filePath: string, args: string[] = []) {
try {
// Read the backup file
const jsonData = await fs.readFile(filePath, 'utf8');
const data = JSON.parse(jsonData);
if (!data.entities || !data.relations || !data.zones) {
console.error('Invalid backup file format');
process.exit(1);
}
console.log(`Found ${data.entities.length} entities, ${data.relations.length} relations, and ${data.zones.length} zones in backup`);
// Confirm with user
const confirmMessage = 'This will merge the backup with existing data. Continue? (y/N) ';
const confirmed = await confirmAction(confirmMessage, args);
if (!confirmed) {
console.log('Operation cancelled');
return;
}
// Import the data
await kgClient.initialize();
const result = await kgClient.importAllData(data);
console.log('Restore completed:');
console.log(`Zones added: ${result.zonesAdded}`);
console.log(`Entities added: ${result.entitiesAdded}`);
console.log(`Relations added: ${result.relationsAdded}`);
} catch (error) {
console.error('Error restoring backup:', (error as Error).message);
process.exit(1);
}
}
/**
* Helper function to check if --yes flag is present
*/
function hasYesFlag(args: string[]): boolean {
return args.includes('--yes') || args.includes('-y');
}
/**
* Remove --yes or -y flags from arguments if present
*/
function cleanArgs(args: string[]): string[] {
return args.filter(arg => arg !== '--yes' && arg !== '-y');
}
/**
* Confirm an action with the user
* @param message The confirmation message to display
* @param args Command line arguments to check for --yes flag
*/
async function confirmAction(message: string, args: string[]): Promise<boolean> {
// Skip confirmation if --yes flag is present
if (hasYesFlag(args)) {
return true;
}
// Otherwise, ask for confirmation
const rl = readline.createInterface({
input: process.stdin,
output: process.stdout
});
const answer = await new Promise<string>((resolve) => {
rl.question(message, (ans: string) => {
resolve(ans);
rl.close();
});
});
return answer.toLowerCase() === 'y';
}
/**
* Update zone descriptions based on content
* @param zone Zone name to update
* @param limit Optional maximum number of entities to analyze
* @param userPrompt Optional user-provided description of the zone's purpose
*/
async function updateZoneDescriptions(zone: string, limit: number = 20, userPrompt?: string) {
try {
console.log(`Updating descriptions for zone "${zone}" based on content...`);
// Initialize client
await kgClient.initialize(zone);
// Get current zone metadata
const zoneMetadata = await kgClient.getZoneMetadata(zone);
if (!zoneMetadata) {
console.error(`Zone "${zone}" not found`);
process.exit(1);
}
console.log(`Finding the most representative entities to answer "What is ${zone}?"...`);
// Try multiple search strategies to get the most representative entities
const relevantEntities = [];
// If user provided a prompt, search for it specifically
if (userPrompt) {
console.log(`Using user-provided description: "${userPrompt}"`);
const { entities: promptEntities } = await kgClient.userSearch({
query: userPrompt,
limit: limit,
sortBy: 'importance',
includeObservations: true,
informationNeeded: zone !== 'default' ? `What is ${zone}?` : undefined,
reason: `Trying to figure out what ${zone} is about, in order to update the zone description.`,
zone: zone
});
relevantEntities.push(...promptEntities);
}
// Strategy 1: First get most important entities
if (relevantEntities.length < limit) {
const { entities: importantEntities } = await kgClient.userSearch({
query: "*", // Get all entities
limit: Math.floor(limit / 2),
sortBy: 'importance',
includeObservations: true,
informationNeeded: zone !== 'default' ? `What is ${zone}?` : undefined,
reason: `Trying to figure out what ${zone} is about, in order to update the zone description.`,
zone: zone
});
// Add only new entities
for (const entity of importantEntities) {
if (!relevantEntities.some(e =>
e.entityType && // Make sure we're comparing entities
entity.entityType &&
e.name === entity.name
)) {
relevantEntities.push(entity);
}
}
}
// Strategy 2: Use zone name as search query to find semantically related entities
if (relevantEntities.length < limit) {
const { entities: nameEntities } = await kgClient.userSearch({
query: zone, // Use zone name as search query
limit: Math.ceil(limit / 4),
sortBy: 'relevance',
includeObservations: true,
informationNeeded: zone !== 'default' ? `What is ${zone}?` : undefined,
reason: `Trying to figure out what ${zone} is about, in order to update the zone description.`,
zone: zone
});
// Add only new entities not already in the list
for (const entity of nameEntities) {
if (!relevantEntities.some(e =>
e.entityType && // Make sure we're comparing entities
entity.entityType &&
e.name === entity.name
)) {
relevantEntities.push(entity);
}
}
}
// Strategy 3: Get most frequently accessed entities
if (relevantEntities.length < limit) {
const { entities: recentEntities } = await kgClient.userSearch({
query: "*", // Get all entities
limit: Math.ceil(limit / 4),
sortBy: 'recent',
includeObservations: true,
informationNeeded: zone !== 'default' ? `What is ${zone}?` : undefined,
reason: `Trying to figure out what ${zone} is about, in order to update the zone description.`,
zone: zone
});
// Add only new entities not already in the list
for (const entity of recentEntities) {
if (!relevantEntities.some(e =>
e.entityType && // Make sure we're comparing entities
entity.entityType &&
e.name === entity.name
)) {
relevantEntities.push(entity);
}
}
}
if (relevantEntities.length === 0) {
console.log(`No entities found in zone "${zone}" to analyze.`);
return;
}
// Trim to limit
const finalEntities = relevantEntities.slice(0, limit);
console.log(`Found ${finalEntities.length} representative entities to analyze for zone description.`);
// Generate descriptions using AI
console.log("\nGenerating descriptions...");
try {
const descriptions = await GroqAI.generateZoneDescriptions(
zone,
zoneMetadata.description || '',
finalEntities,
userPrompt
);
// Update the zone with new descriptions
await kgClient.updateZoneDescriptions(
zone,
descriptions.description,
descriptions.shortDescription
);
console.log(`\nUpdated descriptions for zone "${zone}":`);
console.log(`\nShort Description: ${descriptions.shortDescription}`);
console.log(`\nFull Description: ${descriptions.description}`);
} catch (error) {
console.error(`\nError generating descriptions: ${error.message}`);
console.log('\nFalling back to existing description. Please try again or provide a more specific prompt.');
}
} catch (error) {
console.error('Error updating zone descriptions:', (error as Error).message);
process.exit(1);
}
}
/**
* Main function to parse and execute commands
*/
async function main() {
const args = process.argv.slice(2);
const cleanedArgs = cleanArgs(args);
const command = cleanedArgs[0];
if (!command || command === 'help') {
showHelp();
return;
}
switch (command) {
case 'init':
await initializeIndex();
break;
case 'import':
if (!cleanedArgs[1]) {
console.error('Error: File path is required for import');
process.exit(1);
}
await importFromJsonFile(cleanedArgs[1], {
...esOptions,
defaultZone: cleanedArgs[2] || DEFAULT_ZONE
});
break;
case 'export':
if (!cleanedArgs[1]) {
console.error('Error: File path is required for export');
process.exit(1);
}
await exportToJsonFile(cleanedArgs[1], {
...esOptions,
defaultZone: cleanedArgs[2] || DEFAULT_ZONE
});
break;
case 'backup':
if (!cleanedArgs[1]) {
console.error('Error: File path is required for backup');
process.exit(1);
}
await backupAll(cleanedArgs[1]);
break;
case 'restore':
if (!cleanedArgs[1]) {
console.error('Error: File path is required for restore');
process.exit(1);
}
await restoreAll(cleanedArgs[1], args);
break;
case 'stats':
await showStats(cleanedArgs[1]);
break;
case 'search':
if (!cleanedArgs[1]) {
console.error('Error: Search query is required');
process.exit(1);
}
await searchGraph(cleanedArgs[1], cleanedArgs[2]);
break;
case 'reset':
await resetIndex(cleanedArgs[1], args);
break;
case 'entity':
if (!cleanedArgs[1]) {
console.error('Error: Entity name is required');
process.exit(1);
}
await showEntity(cleanedArgs[1], cleanedArgs[2]);
break;
case 'zones':
const zonesCommand = cleanedArgs[1];
if (!zonesCommand) {
await listZones();
break;
}
switch (zonesCommand) {
case 'list':
await listZones();
break;
case 'add':
if (!cleanedArgs[2]) {
console.error('Error: Zone name is required');
process.exit(1);
}
await addZone(cleanedArgs[2], cleanedArgs[3]);
break;
case 'delete':
if (!cleanedArgs[2]) {
console.error('Error: Zone name is required');
process.exit(1);
}
await deleteZone(cleanedArgs[2], args);
break;
case 'stats':
if (!cleanedArgs[2]) {
console.error('Error: Zone name is required');
process.exit(1);
}
await showStats(cleanedArgs[2]);
break;
case 'update_descriptions':
if (!cleanedArgs[2]) {
console.error('Error: Zone name is required');
process.exit(1);
}
let limit = 20;
let userPrompt = undefined;
// Check if the third argument is a number (limit) or a string (userPrompt)
if (cleanedArgs[3]) {
const parsedLimit = parseInt(cleanedArgs[3], 10);
if (!isNaN(parsedLimit) && parsedLimit.toString() === cleanedArgs[3]) {
// Only interpret as a limit if it's a pure number with no text
limit = parsedLimit;
// If there's a fourth argument, it's the user prompt
if (cleanedArgs[4]) {
userPrompt = cleanedArgs.slice(4).join(' ');
}
} else {
// If third argument isn't a pure number, it's the start of the user prompt
userPrompt = cleanedArgs.slice(3).join(' ');
}
}
await updateZoneDescriptions(cleanedArgs[2], limit, userPrompt);
break;
default:
console.error(`Unknown zones command: ${zonesCommand}`);
showHelp();
process.exit(1);
}
break;
case 'relations':
if (!cleanedArgs[1]) {
console.error('Error: Entity name is required');
process.exit(1);
}
await showRelations(cleanedArgs[1], cleanedArgs[2]);
break;
default:
console.error(`Unknown command: ${command}`);
showHelp();
process.exit(1);
}
}
// Run the CLI
main().catch(error => {
console.error('Error:', (error as Error).message);
process.exit(1);
});
```
--------------------------------------------------------------------------------
/src/index.ts:
--------------------------------------------------------------------------------
```typescript
#!/usr/bin/env node
// @ts-ignore
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
// @ts-ignore
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import {
CallToolRequestSchema,
ListToolsRequestSchema,
ListResourcesRequestSchema,
ListPromptsRequestSchema
} from "@modelcontextprotocol/sdk/types.js";
import { KnowledgeGraphClient } from './kg-client.js';
import { ESEntity, ESRelation, ESSearchParams } from './es-types.js';
import GroqAI from './ai-service.js';
import { inspectFile } from './filesystem/index.js';
// Environment configuration for Elasticsearch
const ES_NODE = process.env.ES_NODE || 'http://localhost:9200';
const ES_USERNAME = process.env.ES_USERNAME;
const ES_PASSWORD = process.env.ES_PASSWORD;
const DEBUG = process.env.DEBUG === 'true';
// Configure ES client with authentication if provided
const esOptions: { node: string; auth?: { username: string; password: string } } = {
node: ES_NODE
};
if (ES_USERNAME && ES_PASSWORD) {
esOptions.auth = { username: ES_USERNAME, password: ES_PASSWORD };
}
// Create KG client
const kgClient = new KnowledgeGraphClient(esOptions);
// Helper function to format dates in YYYY-MM-DD format
function formatDate(date: Date = new Date()): string {
return date.toISOString().split('T')[0]; // Returns YYYY-MM-DD
}
// Start the MCP server
async function startServer() {
try {
// Initialize the knowledge graph
await kgClient.initialize();
// Use stderr for logging, not stdout
console.error('Elasticsearch Knowledge Graph initialized');
} catch (error) {
console.error('Warning: Failed to connect to Elasticsearch:', error.message);
console.error('The memory server will still start, but operations requiring Elasticsearch will fail');
}
// Create the MCP server
const server = new Server({
name: "memory",
version: "1.0.0",
}, {
capabilities: {
tools: {},
// Add empty resource and prompt capabilities to support list requests
resources: {},
prompts: {}
},
});
console.error('Starting MCP server...');
// Handle resources/list requests (return empty list)
server.setRequestHandler(ListResourcesRequestSchema, async () => {
if (DEBUG) {
console.error('ListResourcesRequestSchema');
}
return {
resources: []
};
});
// Handle prompts/list requests (return empty list)
server.setRequestHandler(ListPromptsRequestSchema, async () => {
if (DEBUG) {
console.error('ListPromptsRequestSchema');
}
return {
prompts: []
};
});
// Register the tools handler to list all available tools
server.setRequestHandler(ListToolsRequestSchema, async () => {
return {
tools: [
{
name: "inspect_files",
description: "Agent driven file inspection that uses AI to retrieve relevant content from multiple files.",
inputSchema: {
type: "object",
properties: {
file_paths: {
type: "array",
items: { type: "string" },
description: "Paths to the files (or directories) to inspect"
},
information_needed: {
type: "string",
description: "Full description of what information is needed from the files, including the context of the information needed. Do not be vague, be specific. The AI agent does not have access to your context, only this \"information needed\" and \"reason\" fields. That's all it will use to decide that a line is relevant to the information needed. So provide a detailed specific description, listing all the details about what you are looking for."
},
reason: {
type: "string",
description: "Explain why this information is needed to help the AI agent give better results. The more context you provide, the better the results will be."
},
include_lines: {
type: "boolean",
description: "Whether to include the actual line content in the response, which uses more of your limited token quota, but gives more informatiom (default: false)"
},
keywords: {
type: "array",
items: { type: "string" },
description: "Array of specific keywords related to the information needed. AI will target files that contain one of these keywords. REQUIRED and cannot be null or empty - the more keywords you provide, the better the results. Include variations, synonyms, and related terms."
}
},
required: ["file_paths", "information_needed", "include_lines", "keywords"],
additionalProperties: false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
},
{
name: "inspect_knowledge_graph",
description: "Agent driven knowledge graph inspection that uses AI to retrieve relevant entities and relations based on a query.",
inputSchema: {
type: "object",
properties: {
information_needed: {
type: "string",
description: "Full description of what information is needed from the knowledge graph, including the context of the information needed. Do not be vague, be specific. The AI agent does not have access to your context, only this \"information needed\" and \"reason\" fields. That's all it will use to decide that an entity is relevant to the information needed."
},
reason: {
type: "string",
description: "Explain why this information is needed to help the AI agent give better results. The more context you provide, the better the results will be."
},
include_entities: {
type: "boolean",
description: "Whether to include the full entity details in the response, which uses more of your limited token quota, but gives more information (default: false)"
},
include_relations: {
type: "boolean",
description: "Whether to include the entity relations in the response (default: false)"
},
keywords: {
type: "array",
items: { type: "string" },
description: "Array of specific keywords related to the information needed. AI will target entities that match one of these keywords. REQUIRED and cannot be null or empty - the more keywords you provide, the better the results. Include variations, synonyms, and related terms."
},
memory_zone: {
type: "string",
description: "Memory zone to search in. If not provided, uses the default zone."
},
entity_types: {
type: "array",
items: { type: "string" },
description: "Optional filter to specific entity types"
}
},
required: ["information_needed", "keywords"],
additionalProperties: false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
},
{
name: "create_entities",
description: "Create entities in knowledge graph (memory)",
inputSchema: {
type: "object",
properties: {
entities: {
type: "array",
items: {
type: "object",
properties: {
name: {type: "string", description: "Entity name"},
entityType: {type: "string", description: "Entity type"},
observations: {
type: "array",
items: {type: "string"},
description: "Observations about this entity"
}
},
required: ["name", "entityType"]
},
description: "List of entities to create"
},
memory_zone: {
type: "string",
description: "Memory zone to create entities in."
}
},
required: ["entities", "memory_zone"],
additionalProperties: false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
},
{
name: "update_entities",
description: "Update entities in knowledge graph (memory)",
inputSchema: {
type: "object",
properties: {
entities: {
type: "array",
description: "List of entities to update",
items: {
type: "object",
properties: {
name: {type: "string"},
entityType: {type: "string"},
observations: {
type: "array",
items: {type: "string"}
},
isImportant: {type: "boolean"}
},
required: ["name"]
}
},
memory_zone: {
type: "string",
description: "Memory zone specifier. Entities will be updated in this zone."
}
},
required: ["entities", "memory_zone"],
additionalProperties: false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
},
{
name: "delete_entities",
description: "Delete entities from knowledge graph (memory)",
inputSchema: {
type: "object",
properties: {
names: {
type: "array",
items: {type: "string"},
description: "Names of entities to delete"
},
memory_zone: {
type: "string",
description: "Memory zone specifier. Entities will be deleted from this zone."
},
cascade_relations: {
type: "boolean",
description: "Whether to delete relations involving these entities (default: true)",
default: true
}
},
required: ["names", "memory_zone"],
additionalProperties: false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
},
{
name: "create_relations",
description: "Create relationships between entities in knowledge graph (memory)",
inputSchema: {
type: "object",
properties: {
relations: {
type: "array",
description: "List of relations to create",
items: {
type: "object",
properties: {
from: {type: "string", description: "Source entity name"},
fromZone: {type: "string", description: "Optional zone for source entity, defaults to memory_zone or default zone. Must be one of the existing zones."},
to: {type: "string", description: "Target entity name"},
toZone: {type: "string", description: "Optional zone for target entity, defaults to memory_zone or default zone. Must be one of the existing zones."},
type: {type: "string", description: "Relationship type"}
},
required: ["from", "to", "type"]
}
},
memory_zone: {
type: "string",
description: "Optional default memory zone specifier. Used if a relation doesn't specify fromZone or toZone."
},
auto_create_missing_entities: {
type: "boolean",
description: "Whether to automatically create missing entities in the relations (default: true)",
default: true
}
},
required: ["relations"],
additionalProperties: false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
},
{
name: "delete_relations",
description: "Delete relationships from knowledge graph (memory)",
inputSchema: {
type: "object",
properties: {
relations: {
type: "array",
description: "List of relations to delete",
items: {
type: "object",
properties: {
from: {type: "string", description: "Source entity name"},
to: {type: "string", description: "Target entity name"},
type: {type: "string", description: "Relationship type"}
},
required: ["from", "to", "type"]
}
},
memory_zone: {
type: "string",
description: "Optional memory zone specifier. If provided, relations will be deleted from this zone."
}
},
required: ["relations"],
additionalProperties: false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
},
{
name: "search_nodes",
description: "Search entities using ElasticSearch query syntax. Supports boolean operators (AND, OR, NOT), fuzzy matching (~), phrases (\"term\"), proximity (\"terms\"~N), wildcards (*, ?), and boosting (^N). Examples: 'meeting AND notes', 'Jon~', '\"project plan\"~2'. All searches respect zone isolation.",
inputSchema: {
type: "object",
properties: {
query: {
type: "string",
description: "ElasticSearch query string."
},
informationNeeded: {
type: "string",
description: "Important. Describe what information you are looking for, to give a precise context to the search engine AI agent. What questions are you trying to answer? Helps get more useful results."
},
reason: {
type: "string",
description: "Explain why this information is needed to help the AI agent give better results. The more context you provide, the better the results will be."
},
entityTypes: {
type: "array",
items: {type: "string"},
description: "Filter to specific entity types (OR condition if multiple)."
},
limit: {
type: "integer",
description: "Max results (default: 20, or 5 with observations)."
},
sortBy: {
type: "string",
enum: ["relevance", "recency", "importance"],
description: "Sort by match quality, access time, or importance."
},
includeObservations: {
type: "boolean",
description: "Include full entity observations (default: false).",
default: false
},
memory_zone: {
type: "string",
description: "Limit search to specific zone. Omit for default zone."
},
},
required: ["query", "memory_zone", "informationNeeded", "reason"],
additionalProperties: false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
},
{
name: "open_nodes",
description: "Get details about specific entities in knowledge graph (memory) and their relations",
inputSchema: {
type: "object",
properties: {
names: {
type: "array",
items: {type: "string"},
description: "Names of entities to retrieve"
},
memory_zone: {
type: "string",
description: "Optional memory zone to retrieve entities from. If not specified, uses the default zone."
}
},
required: ["names", "memory_zone"],
additionalProperties: false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
},
{
name: "add_observations",
description: "Add observations to an existing entity in knowledge graph (memory)",
inputSchema: {
type: "object",
properties: {
name: {
type: "string",
description: "Name of entity to add observations to"
},
observations: {
type: "array",
items: {type: "string"},
description: "Observations to add to the entity"
},
memory_zone: {
type: "string",
description: "Optional memory zone where the entity is stored. If not specified, uses the default zone."
}
},
required: ["memory_zone", "name", "observations"],
additionalProperties: false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
},
{
name: "mark_important",
description: "Mark entity as important in knowledge graph (memory) by boosting its relevance score",
inputSchema: {
type: "object",
properties: {
name: {
type: "string",
description: "Entity name"
},
important: {
type: "boolean",
description: "Set as important (true - multiply relevance by 10) or not (false - divide relevance by 10)"
},
memory_zone: {
type: "string",
description: "Optional memory zone specifier. If provided, entity will be marked in this zone."
},
auto_create: {
type: "boolean",
description: "Whether to automatically create the entity if it doesn't exist (default: false)",
default: false
}
},
required: ["memory_zone", "name", "important"],
additionalProperties: false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
},
{
name: "get_recent",
description: "Get recently accessed entities from knowledge graph (memory) and their relations",
inputSchema: {
type: "object",
properties: {
limit: {
type: "integer",
description: "Max results (default: 20 if includeObservations is false, 5 if true)"
},
includeObservations: {
type: "boolean",
description: "Whether to include full entity observations in results (default: false)",
default: false
},
memory_zone: {
type: "string",
description: "Optional memory zone to get recent entities from. If not specified, uses the default zone."
}
},
required: ["memory_zone"],
additionalProperties: false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
},
{
name: "list_zones",
description: "List all available memory zones with metadata. When a reason is provided, zones will be filtered and prioritized based on relevance to your needs.",
inputSchema: {
type: "object",
properties: {
reason: {
type: "string",
description: "Reason for listing zones. What zones are you looking for? Why are you looking for them? The AI will use this to prioritize and filter relevant zones."
}
},
additionalProperties: false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
},
{
name: "create_zone",
description: "Create a new memory zone with optional description.",
inputSchema: {
type: "object",
properties: {
name: {
type: "string",
description: "Zone name (cannot be 'default')"
},
shortDescription: {
type: "string",
description: "Short description of the zone."
},
description: {
type: "string",
description: "Full zone description. Make it very descriptive and detailed."
}
},
required: ["name"]
}
},
{
name: "delete_zone",
description: "Delete a memory zone and all its entities/relations.",
inputSchema: {
type: "object",
properties: {
name: {
type: "string",
description: "Zone name to delete (cannot be 'default')"
},
confirm: {
type: "boolean",
description: "Confirmation flag, must be true",
default: false
}
},
required: ["name", "confirm"]
}
},
{
name: "copy_entities",
description: "Copy entities between zones with optional relation handling.",
inputSchema: {
type: "object",
properties: {
names: {
type: "array",
items: { type: "string" },
description: "Entity names to copy"
},
source_zone: {
type: "string",
description: "Source zone"
},
target_zone: {
type: "string",
description: "Target zone"
},
copy_relations: {
type: "boolean",
description: "Copy related relationships (default: true)",
default: true
},
overwrite: {
type: "boolean",
description: "Overwrite if entity exists (default: false)",
default: false
}
},
required: ["names", "source_zone", "target_zone"]
}
},
{
name: "move_entities",
description: "Move entities between zones (copy + delete from source).",
inputSchema: {
type: "object",
properties: {
names: {
type: "array",
items: { type: "string" },
description: "Entity names to move"
},
source_zone: {
type: "string",
description: "Source zone"
},
target_zone: {
type: "string",
description: "Target zone"
},
move_relations: {
type: "boolean",
description: "Move related relationships (default: true)",
default: true
},
overwrite: {
type: "boolean",
description: "Overwrite if entity exists (default: false)",
default: false
}
},
required: ["names", "source_zone", "target_zone"]
}
},
{
name: "merge_zones",
description: "Merge multiple zones with conflict resolution options.",
inputSchema: {
type: "object",
properties: {
source_zones: {
type: "array",
items: { type: "string" },
description: "Source zones to merge from"
},
target_zone: {
type: "string",
description: "Target zone to merge into"
},
delete_source_zones: {
type: "boolean",
description: "Delete source zones after merging",
default: false
},
overwrite_conflicts: {
type: "string",
enum: ["skip", "overwrite", "rename"],
description: "How to handle name conflicts",
default: "skip"
}
},
required: ["source_zones", "target_zone"]
}
},
{
name: "zone_stats",
description: "Get statistics for entities and relationships in a zone.",
inputSchema: {
type: "object",
properties: {
zone: {
type: "string",
description: "Zone name (omit for default zone)"
}
},
required: ["zone"]
}
},
{
name: "get_time_utc",
description: "Get the current UTC time in YYYY-MM-DD hh:mm:ss format",
inputSchema: {
type: "object",
properties: {},
additionalProperties: false
}
}
]
};
});
// Register the call tool handler to handle tool executions
server.setRequestHandler(CallToolRequestSchema, async (request) => {
if (DEBUG) {
console.error('Tool request received:', request.params.name);
console.error('Tool request params:', JSON.stringify(request.params));
}
const toolName = request.params.name;
const params = request.params.arguments as any;
if (DEBUG) {
console.error('Parsed parameters:', JSON.stringify(params));
}
// Helper function to format response for Claude
const formatResponse = (data: any) => {
const stringifiedData = JSON.stringify(data, null, 2);
return {
content: [
{
type: "text",
text: stringifiedData,
},
],
};
};
if (toolName === "inspect_files") {
const { file_paths, information_needed, reason, include_lines, keywords } = params;
// Validate keywords
if (!keywords || !Array.isArray(keywords) || keywords.length === 0) {
return formatResponse({
success: false,
error: "Keywords parameter is required and cannot be null or empty. Please provide an array of keywords to help find relevant information. The more keywords you provide (including variations, synonyms, and related terms), the better the results."
});
}
const results = [];
for (const filePath of file_paths) {
try {
const fileResults = await inspectFile(filePath, information_needed, reason, keywords);
results.push({
filePath,
linesContent: `lines as returned by cat -n ${filePath}`,
lines: include_lines ? fileResults.lines.map(line => `${line.lineNumber}\t${line.content}`) : [],
tentativeAnswer: fileResults.tentativeAnswer
});
} catch (error) {
results.push({
filePath,
error: error.message
});
}
}
return formatResponse({
success: true,
results
});
}
else if (toolName === "inspect_knowledge_graph") {
const { information_needed, reason, include_entities, include_relations, keywords, memory_zone, entity_types } = params;
// Validate keywords
if (!keywords || !Array.isArray(keywords) || keywords.length === 0) {
return formatResponse({
success: false,
error: "Keywords parameter is required and cannot be null or empty. Please provide an array of keywords to help find relevant entities. The more keywords you provide (including variations, synonyms, and related terms), the better the results."
});
}
// Import the inspectKnowledgeGraph function
const { inspectKnowledgeGraph } = await import('./kg-inspection.js');
try {
// Call the inspectKnowledgeGraph function
const results = await inspectKnowledgeGraph(
kgClient,
information_needed,
reason,
keywords,
memory_zone,
entity_types
);
// Format the response based on include_entities and include_relations flags
return formatResponse({
success: true,
tentativeAnswer: results.tentativeAnswer,
entities: include_entities ? results.entities : results.entities.map(e => ({ name: e.name, entityType: e.entityType })),
relations: include_relations ? results.relations : []
});
} catch (error) {
return formatResponse({
success: false,
error: `Error inspecting knowledge graph: ${error.message}`
});
}
}
else if (toolName === "create_entities") {
const entities = params.entities;
const zone = params.memory_zone;
// First, check if any entities already exist or have empty names
const conflictingEntities = [];
const invalidEntities = [];
for (const entity of entities) {
// Check for empty names
if (!entity.name || entity.name.trim() === '') {
invalidEntities.push({
name: "[empty]",
reason: "Entity name cannot be empty"
});
continue;
}
const existingEntity = await kgClient.getEntity(entity.name, zone);
if (existingEntity) {
conflictingEntities.push(entity.name);
}
}
// If there are conflicts or invalid entities, reject the operation
if (conflictingEntities.length > 0 || invalidEntities.length > 0) {
const zoneMsg = zone ? ` in zone "${zone}"` : "";
// Fetch existing entity details if there are conflicts
const existingEntitiesData = [];
if (conflictingEntities.length > 0) {
for (const entityName of conflictingEntities) {
const existingEntity = await kgClient.getEntity(entityName, zone);
if (existingEntity) {
existingEntitiesData.push(existingEntity);
}
}
}
return formatResponse({
success: false,
error: `Entity creation failed${zoneMsg}, no entities were created.`,
conflicts: conflictingEntities.length > 0 ? conflictingEntities : undefined,
existingEntities: existingEntitiesData.length > 0 ? existingEntitiesData : undefined,
invalidEntities: invalidEntities.length > 0 ? invalidEntities : undefined,
message: conflictingEntities.length > 0 ?
"Feel free to extend existing entities with more information if needed, or create entities with different names. Use update_entities to modify existing entities." :
"Please provide valid entity names for all entities."
});
}
// If no conflicts, proceed with entity creation
const createdEntities = [];
for (const entity of entities) {
const savedEntity = await kgClient.saveEntity({
name: entity.name,
entityType: entity.entityType,
observations: entity.observations,
relevanceScore: entity.relevanceScore ?? 1.0
}, zone);
createdEntities.push(savedEntity);
}
return formatResponse({
success: true,
entities: createdEntities.map(e => ({
name: e.name,
entityType: e.entityType,
observations: e.observations
}))
});
}
else if (toolName === "update_entities") {
const entities = params.entities;
const zone = params.memory_zone;
const updatedEntities = [];
for (const entity of entities) {
// Get the existing entity first, then update with new values
const existingEntity = await kgClient.getEntity(entity.name, zone);
if (!existingEntity) {
const zoneMsg = zone ? ` in zone "${zone}"` : "";
throw new Error(`Entity "${entity.name}" not found${zoneMsg}`);
}
// Update with new values, preserving existing values for missing fields
const updatedEntity = await kgClient.saveEntity({
name: entity.name,
entityType: entity.entityType || existingEntity.entityType,
observations: entity.observations || existingEntity.observations,
relevanceScore: entity.relevanceScore || ((existingEntity.relevanceScore ?? 1.0) * 2.0)
}, zone);
updatedEntities.push(updatedEntity);
}
return formatResponse({
entities: updatedEntities.map(e => ({
name: e.name,
entityType: e.entityType,
observations: e.observations
}))
});
}
else if (toolName === "delete_entities") {
const names = params.names;
const zone = params.memory_zone;
const cascadeRelations = params.cascade_relations !== false; // Default to true
const results = [];
const invalidNames = [];
// Validate names before attempting deletion
for (const name of names) {
if (!name || name.trim() === '') {
invalidNames.push("[empty]");
continue;
}
}
// If there are invalid names, reject those operations
if (invalidNames.length > 0) {
return formatResponse({
success: false,
error: "Entity deletion failed for some entities",
invalidNames,
message: "Entity names cannot be empty"
});
}
// Delete each valid entity individually
for (const name of names) {
try {
const success = await kgClient.deleteEntity(name, zone, {
cascadeRelations
});
results.push({ name, deleted: success });
} catch (error) {
results.push({ name, deleted: false, error: error.message });
}
}
return formatResponse({
success: true,
results
});
}
else if (toolName === "create_relations") {
const relations = params.relations;
const defaultZone = params.memory_zone;
const autoCreateMissingEntities = params.auto_create_missing_entities !== false; // Default to true for backward compatibility
const savedRelations = [];
const failedRelations = [];
for (const relation of relations) {
const fromZone = relation.fromZone || defaultZone;
const toZone = relation.toZone || defaultZone;
try {
const savedRelation = await kgClient.saveRelation({
from: relation.from,
to: relation.to,
relationType: relation.type
}, fromZone, toZone, { autoCreateMissingEntities });
savedRelations.push(savedRelation);
} catch (error) {
failedRelations.push({
relation,
error: error.message
});
}
}
// If there were any failures, include them in the response
if (failedRelations.length > 0) {
return formatResponse({
success: savedRelations.length > 0,
relations: savedRelations.map(r => ({
from: r.from,
to: r.to,
type: r.relationType,
fromZone: r.fromZone,
toZone: r.toZone
})),
failedRelations
});
}
return formatResponse({
success: true,
relations: savedRelations.map(r => ({
from: r.from,
to: r.to,
type: r.relationType,
fromZone: r.fromZone,
toZone: r.toZone
}))
});
}
else if (toolName === "delete_relations") {
const relations = params.relations;
const zone = params.memory_zone;
const results = [];
// Delete each relation individually
for (const relation of relations) {
const success = await kgClient.deleteRelation(
relation.from,
relation.to,
relation.type,
zone,
zone
);
results.push({
from: relation.from,
to: relation.to,
type: relation.type,
deleted: success
});
}
return formatResponse({
success: true,
results
});
}
else if (toolName === "search_nodes") {
const includeObservations = params.includeObservations ?? false;
const zone = params.memory_zone;
// Use the high-level userSearch method that handles AI filtering internally
const { entities: filteredEntities, relations: formattedRelations } = await kgClient.userSearch({
query: params.query,
entityTypes: params.entityTypes,
limit: params.limit,
sortBy: params.sortBy,
includeObservations,
zone,
informationNeeded: params.informationNeeded,
reason: params.reason
});
return formatResponse({ entities: filteredEntities, relations: formattedRelations });
}
else if (toolName === "open_nodes") {
const names = params.names || [];
const zone = params.memory_zone;
// Get the entities
const entities: ESEntity[] = [];
for (const name of names) {
const entity = await kgClient.getEntity(name, zone);
if (entity) {
entities.push(entity);
}
}
// Format entities
const formattedEntities = entities.map(e => ({
name: e.name,
entityType: e.entityType,
observations: e.observations
}));
// Get relations between these entities
const entityNames = formattedEntities.map(e => e.name);
const { relations } = await kgClient.getRelationsForEntities(entityNames, zone);
// Map relations to the expected format
const formattedRelations = relations.map(r => ({
from: r.from,
to: r.to,
type: r.relationType,
fromZone: r.fromZone,
toZone: r.toZone
}));
return formatResponse({ entities: formattedEntities, relations: formattedRelations });
}
else if (toolName === "add_observations") {
const name = params.name;
const observations = params.observations;
const zone = params.memory_zone;
// Get existing entity
const entity = await kgClient.getEntity(name, zone);
if (!entity) {
const zoneMsg = zone ? ` in zone "${zone}"` : "";
return formatResponse({
success: false,
error: `Entity "${name}" not found${zoneMsg}`,
message: "Please create the entity before adding observations."
});
}
// Add observations to the entity
const updatedEntity = await kgClient.addObservations(name, observations, zone);
return formatResponse({
success: true,
entity: updatedEntity
});
}
else if (toolName === "mark_important") {
const name = params.name;
const important = params.important;
const zone = params.memory_zone;
const autoCreate = params.auto_create === true;
try {
// Mark the entity as important, with auto-creation if specified
const updatedEntity = await kgClient.markImportant(name, important, zone, {
autoCreateMissingEntities: autoCreate
});
return formatResponse({
success: true,
entity: updatedEntity,
auto_created: autoCreate && !(await kgClient.getEntity(name, zone))
});
} catch (error) {
const zoneMsg = zone ? ` in zone "${zone}"` : "";
return formatResponse({
success: false,
error: `Entity "${name}" not found${zoneMsg}`,
message: "Please create the entity before marking it as important, or set auto_create to true."
});
}
}
else if (toolName === "get_recent") {
const limit = params.limit || 20;
const includeObservations = params.includeObservations ?? false;
const zone = params.memory_zone;
const recentEntities = await kgClient.getRecentEntities(limit, includeObservations, zone);
return formatResponse({
entities: recentEntities.map(e => ({
name: e.name,
entityType: e.entityType,
observations: e.observations
})),
total: recentEntities.length
});
}
else if (toolName === "list_zones") {
const reason = params.reason;
const zones = await kgClient.listMemoryZones(reason);
// If reason is provided and GroqAI is available, use AI to score zone usefulness
if (reason && GroqAI.isEnabled && zones.length > 0) {
try {
// Get usefulness scores for each zone
const usefulnessScores = await GroqAI.classifyZoneUsefulness(zones, reason);
// Process zones based on their usefulness scores
const processedZones = zones.map(zone => {
const usefulness = usefulnessScores[zone.name] !== undefined ?
usefulnessScores[zone.name] : 2; // Default to very useful (2) if not classified
// Format zone info based on usefulness score
if (usefulness === 0) { // Not useful
return {
name: zone.name,
usefulness: 'not useful'
};
} else if (usefulness === 1) { // A little useful
return {
name: zone.name,
description: zone.description,
usefulness: 'a little useful'
};
} else { // Very useful (2) or default
return {
name: zone.name,
description: zone.description,
created_at: zone.createdAt,
last_modified: zone.lastModified,
config: zone.config,
usefulness: 'very useful'
};
}
});
// Sort zones by usefulness (most useful first)
processedZones.sort((a, b) => {
const scoreA = usefulnessScores[a.name] !== undefined ? usefulnessScores[a.name] : 2;
const scoreB = usefulnessScores[b.name] !== undefined ? usefulnessScores[b.name] : 2;
return scoreB - scoreA; // Descending order
});
return formatResponse({
zones: processedZones
});
} catch (error) {
console.error('Error classifying zones:', error);
// Fall back to default behavior
}
}
// Default behavior (no reason provided or AI failed)
return formatResponse({
zones: zones.map(zone => ({
name: zone.name,
description: zone.description,
created_at: zone.createdAt,
last_modified: zone.lastModified,
usefulness: 'very useful' // Default to very useful when no AI filtering is done
}))
});
}
else if (toolName === "create_zone") {
const name = params.name;
const description = params.description;
try {
await kgClient.addMemoryZone(name, description);
return formatResponse({
success: true,
zone: name,
message: `Zone "${name}" created successfully`
});
} catch (error) {
return formatResponse({
success: false,
error: `Failed to create zone: ${(error as Error).message}`
});
}
}
else if (toolName === "delete_zone") {
const name = params.name;
const confirm = params.confirm === true;
if (!confirm) {
return formatResponse({
success: false,
error: "Confirmation required. Set confirm=true to proceed with deletion."
});
}
try {
const result = await kgClient.deleteMemoryZone(name);
if (result) {
return formatResponse({
success: true,
message: `Zone "${name}" deleted successfully`
});
} else {
return formatResponse({
success: false,
error: `Failed to delete zone "${name}"`
});
}
} catch (error) {
return formatResponse({
success: false,
error: `Error deleting zone: ${(error as Error).message}`
});
}
}
else if (toolName === "copy_entities") {
const names = params.names;
const sourceZone = params.source_zone;
const targetZone = params.target_zone;
const copyRelations = params.copy_relations !== false;
const overwrite = params.overwrite === true;
try {
const result = await kgClient.copyEntitiesBetweenZones(
names,
sourceZone,
targetZone,
{
copyRelations,
overwrite
}
);
return formatResponse({
success: result.entitiesCopied.length > 0,
entities_copied: result.entitiesCopied,
entities_skipped: result.entitiesSkipped,
relations_copied: result.relationsCopied
});
} catch (error) {
return formatResponse({
success: false,
error: `Error copying entities: ${(error as Error).message}`
});
}
}
else if (toolName === "move_entities") {
const names = params.names;
const sourceZone = params.source_zone;
const targetZone = params.target_zone;
const moveRelations = params.move_relations !== false;
const overwrite = params.overwrite === true;
try {
const result = await kgClient.moveEntitiesBetweenZones(
names,
sourceZone,
targetZone,
{
moveRelations,
overwrite
}
);
return formatResponse({
success: result.entitiesMoved.length > 0,
entities_moved: result.entitiesMoved,
entities_skipped: result.entitiesSkipped,
relations_moved: result.relationsMoved
});
} catch (error) {
return formatResponse({
success: false,
error: `Error moving entities: ${(error as Error).message}`
});
}
}
else if (toolName === "merge_zones") {
const sourceZones = params.source_zones;
const targetZone = params.target_zone;
const deleteSourceZones = params.delete_source_zones === true;
const overwriteConflicts = params.overwrite_conflicts || 'skip';
try {
const result = await kgClient.mergeZones(
sourceZones,
targetZone,
{
deleteSourceZones,
overwriteConflicts: overwriteConflicts as 'skip' | 'overwrite' | 'rename'
}
);
return formatResponse({
success: result.mergedZones.length > 0,
merged_zones: result.mergedZones,
failed_zones: result.failedZones,
entities_copied: result.entitiesCopied,
entities_skipped: result.entitiesSkipped,
relations_copied: result.relationsCopied
});
} catch (error) {
return formatResponse({
success: false,
error: `Error merging zones: ${(error as Error).message}`
});
}
}
else if (toolName === "zone_stats") {
const zone = params.zone;
try {
const stats = await kgClient.getMemoryZoneStats(zone);
return formatResponse({
zone: stats.zone,
entity_count: stats.entityCount,
relation_count: stats.relationCount,
entity_types: stats.entityTypes,
relation_types: stats.relationTypes
});
} catch (error) {
return formatResponse({
success: false,
error: `Error getting zone stats: ${(error as Error).message}`
});
}
}
else if (toolName === "get_time_utc") {
// Get current time in UTC
const now = new Date();
// Format time as YYYY-MM-DD hh:mm:ss
const pad = (num: number) => num.toString().padStart(2, '0');
const year = now.getUTCFullYear();
const month = pad(now.getUTCMonth() + 1); // months are 0-indexed
const day = pad(now.getUTCDate());
const hours = pad(now.getUTCHours());
const minutes = pad(now.getUTCMinutes());
const seconds = pad(now.getUTCSeconds());
const formattedTime = `${year}-${month}-${day} ${hours}:${minutes}:${seconds}`;
return formatResponse({
utc_time: formattedTime
});
}
});
return server;
}
// Start the server with proper transport and error handling
async function initServer() {
const server = await startServer();
// Connect the server to the transport
const transport = new StdioServerTransport();
await server.connect(transport);
console.error('MCP server running on stdio');
}
// Initialize with error handling
initServer().catch(error => {
console.error('Error starting server:', error);
process.exit(1);
});
```
--------------------------------------------------------------------------------
/src/kg-client.ts:
--------------------------------------------------------------------------------
```typescript
// WARNING: no console.log in this file, it will break MCP server. Use console.error instead
import { Client } from '@elastic/elasticsearch';
import {
KG_INDEX_CONFIG,
KG_INDEX_PREFIX,
KG_RELATIONS_INDEX,
KG_METADATA_INDEX,
getIndexName,
ESEntity,
ESRelation,
ESHighlightResponse,
ESSearchParams,
ESSearchResponse
} from './es-types.js';
// Zone metadata document structure
interface ZoneMetadata {
name: string;
description?: string;
shortDescription?: string;
createdAt: string;
lastModified: string;
config?: Record<string, any>;
}
// Import the AI service
import GroqAI from './ai-service.js';
/**
* Knowledge Graph Client
*
* Core library for interacting with the Elasticsearch-backed knowledge graph
*/
export class KnowledgeGraphClient {
private client: Client;
private initialized: boolean = false;
private defaultZone: string;
// Cache of initialized indices to avoid repeated checks
private initializedIndices: Set<string> = new Set();
// Cache of existing zones to avoid repeated database queries when checking zone existence
// This improves performance for operations that check the same zones multiple times
private existingZonesCache: Record<string, boolean> = {};
/**
* Create a new KnowledgeGraphClient
* @param options Connection options
*/
constructor(private options: {
node: string;
auth?: { username: string; password: string };
defaultZone?: string;
}) {
this.client = new Client({
node: options.node,
auth: options.auth,
});
this.defaultZone = options.defaultZone || process.env.KG_DEFAULT_ZONE || 'default';
}
private getIndexForZone(zone?: string): string {
return getIndexName(zone || this.defaultZone);
}
/**
* Initialize the knowledge graph (create index if needed)
*/
async initialize(zone?: string): Promise<void> {
if (!this.initialized) {
this.client = new Client({
node: this.options.node,
auth: this.options.auth,
});
this.initialized = true;
// Initialize the metadata index if it doesn't exist yet
const metadataIndexExists = await this.client.indices.exists({
index: KG_METADATA_INDEX
});
if (!metadataIndexExists) {
await this.client.indices.create({
index: KG_METADATA_INDEX,
mappings: {
properties: {
name: { type: 'keyword' },
description: { type: 'text' },
createdAt: { type: 'date' },
lastModified: { type: 'date' },
config: { type: 'object', enabled: false }
}
}
});
console.error(`Created metadata index: ${KG_METADATA_INDEX}`);
// Add default zone metadata
await this.saveZoneMetadata('default', 'Default knowledge zone');
}
// Initialize the relations index if it doesn't exist yet
const relationsIndexExists = await this.client.indices.exists({
index: KG_RELATIONS_INDEX
});
if (!relationsIndexExists) {
await this.client.indices.create({
index: KG_RELATIONS_INDEX,
...KG_INDEX_CONFIG
});
console.error(`Created relations index: ${KG_RELATIONS_INDEX}`);
}
}
// Continue with zone-specific index initialization
const indexName = this.getIndexForZone(zone);
// If we've already initialized this index, skip
if (this.initializedIndices.has(indexName)) {
return;
}
const indexExists = await this.client.indices.exists({ index: indexName });
if (!indexExists) {
await this.client.indices.create({
index: indexName,
...KG_INDEX_CONFIG
});
console.error(`Created index: ${indexName}`);
}
this.initializedIndices.add(indexName);
}
/**
* Create or update an entity
* @param entity Entity to create or update
* @param zone Optional memory zone name, uses defaultZone if not specified
* @param options Optional configuration options
* @param options.validateZones Whether to validate that zones exist before creating entities (default: true)
*/
async saveEntity(
entity: Omit<ESEntity, 'type' | 'readCount' | 'lastRead' | 'lastWrite' | 'zone'>,
zone?: string,
options?: {
validateZones?: boolean;
}
): Promise<ESEntity> {
// Validate entity name is not empty
if (!entity.name || entity.name.trim() === '') {
throw new Error('Entity name cannot be empty');
}
const actualZone = zone || this.defaultZone;
await this.initialize(actualZone);
// Default to true for zone validation
const validateZones = options?.validateZones ?? true;
// Validate that zone exists if required
if (validateZones && actualZone !== this.defaultZone) {
const zoneExists = await this.zoneExists(actualZone);
if (!zoneExists) {
throw new Error(`Cannot create entity: Zone '${actualZone}' does not exist. Create the zone first.`);
}
}
const now = new Date().toISOString();
const existingEntity = await this.getEntity(entity.name, actualZone);
const newEntity: ESEntity = {
type: 'entity',
name: entity.name,
entityType: entity.entityType,
observations: entity.observations || [],
// If entity exists, preserve its readCount and lastRead, but update lastWrite
readCount: existingEntity?.readCount ?? 0,
lastRead: existingEntity?.lastRead ?? now,
lastWrite: now,
relevanceScore: entity.relevanceScore ?? (existingEntity?.relevanceScore ?? 1.0),
zone: actualZone
};
const indexName = this.getIndexForZone(actualZone);
await this.client.index({
index: indexName,
id: `entity:${entity.name}`,
document: newEntity,
refresh: true // Make sure it's immediately available for search
});
return newEntity;
}
/**
* Get an entity by name without updating lastRead timestamp
* @param name Entity name
* @param zone Optional memory zone name, uses defaultZone if not specified
*/
async getEntityWithoutUpdatingLastRead(name: string, zone?: string): Promise<ESEntity | null> {
const actualZone = zone || this.defaultZone;
await this.initialize(actualZone);
try {
const indexName = this.getIndexForZone(actualZone);
const id = `entity:${name}`;
// Try direct get by ID first
try {
const response = await this.client.get({
index: indexName,
id: id
});
if (response && response._source) {
return response._source as ESEntity;
}
} catch (error) {
// If not found by ID, try search
if (error.statusCode === 404) {
// Fall through to search
} else {
throw error;
}
}
// If direct get fails, use search with explicit zone filter
const response = await this.client.search({
index: indexName,
body: {
query: {
bool: {
must: [
// Use term query for exact name matching
{ term: { name: name } },
{ term: { type: 'entity' } },
{ term: { zone: actualZone } }
]
}
},
size: 1
}
});
const typedResponse = response as unknown as ESSearchResponse<ESEntity>;
if (typedResponse.hits.total.value === 0) {
return null;
}
return typedResponse.hits.hits[0]._source;
} catch (error) {
if (error.statusCode === 404) {
return null;
}
throw error;
}
}
/**
* Get an entity by name and update lastRead timestamp and readCount
* @param name Entity name
* @param zone Optional memory zone name, uses defaultZone if not specified
*/
async getEntity(name: string, zone?: string): Promise<ESEntity | null> {
const actualZone = zone || this.defaultZone;
const entity = await this.getEntityWithoutUpdatingLastRead(name, actualZone);
if (!entity) {
return null;
}
// Update lastRead and readCount in memory (skip the database update if it fails)
const now = new Date().toISOString();
const updatedEntity = {
...entity,
lastRead: now,
readCount: entity.readCount + 1
};
try {
// Try to update in the database, but don't fail if it doesn't work
const indexName = this.getIndexForZone(actualZone);
// Search for the entity by name and zone to get the _id
const searchResponse = await this.client.search({
index: indexName,
body: {
query: {
bool: {
must: [
{ term: { name: name } },
{ term: { type: 'entity' } },
{ term: { zone: actualZone } }
]
}
},
size: 1
}
});
const typedResponse = searchResponse as unknown as ESSearchResponse<ESEntity>;
if (typedResponse.hits.total.value > 0) {
const docId = typedResponse.hits.hits[0]._id;
await this.client.update({
index: indexName,
id: docId,
doc: {
lastRead: now,
readCount: entity.readCount + 1
},
refresh: true
});
} else {
// This indicates the entity exists in memory but not in the index
// Instead of showing an error message, silently handle this condition
// The entity is still returned to the caller with updated timestamps
}
} catch (error) {
// If update fails, just log it and return the entity with updated timestamps
console.error(`Warning: Failed to update lastRead timestamp for entity ${name}: ${(error as Error).message}`);
}
return updatedEntity;
}
/**
* Delete an entity by name
* @param name Entity name
* @param zone Optional memory zone name, uses defaultZone if not specified
* @param options Optional configuration options
* @param options.cascadeRelations Whether to delete relations involving this entity (default: true)
*/
async deleteEntity(
name: string,
zone?: string,
options?: {
cascadeRelations?: boolean;
}
): Promise<boolean> {
// Validate entity name is not empty
if (!name || name.trim() === '') {
throw new Error('Entity name cannot be empty');
}
const actualZone = zone || this.defaultZone;
await this.initialize(actualZone);
// Default to true for cascading relation deletion
const cascadeRelations = options?.cascadeRelations !== false;
try {
// First, check if the entity exists
const entity = await this.getEntityWithoutUpdatingLastRead(name, actualZone);
if (!entity) {
return false;
}
const indexName = this.getIndexForZone(actualZone);
// Delete relations involving this entity if cascading is enabled
if (cascadeRelations) {
console.error(`Cascading relations for entity ${name} in zone ${actualZone}`);
// First, delete relations within the same zone
await this.client.deleteByQuery({
index: indexName,
body: {
query: {
bool: {
must: [
{ term: { type: 'relation' } },
{
bool: {
should: [
{ term: { from: name } },
{ term: { to: name } }
]
}
}
]
}
}
},
refresh: true
});
// Then, delete cross-zone relations where this entity is involved
// These are stored in the KG_RELATIONS_INDEX
await this.client.deleteByQuery({
index: KG_RELATIONS_INDEX,
body: {
query: {
bool: {
must: [
{
bool: {
should: [
{
bool: {
must: [
{ term: { from: name } },
{ term: { fromZone: actualZone } }
]
}
},
{
bool: {
must: [
{ term: { to: name } },
{ term: { toZone: actualZone } }
]
}
}
]
}
}
]
}
}
},
refresh: true
});
} else {
console.error(`Skipping relation cascade for entity ${name} in zone ${actualZone}`);
}
// Delete the entity
await this.client.delete({
index: indexName,
id: `entity:${name}`,
refresh: true
});
return true;
} catch (error) {
if (error.statusCode === 404) {
return false;
}
throw error;
}
}
/**
* Check if a memory zone exists
* @param zone Zone name to check
* @returns Promise<boolean> True if the zone exists, false otherwise
*
* This method uses a caching strategy to avoid repeated database queries:
* 1. Default zone always exists and is automatically added to the cache
* 2. If the requested zone is in the cache, returns the cached result
* 3. Otherwise, checks the zone metadata and updates the cache
* 4. The cache is maintained by addMemoryZone, deleteMemoryZone, and listMemoryZones
*/
async zoneExists(zone: string): Promise<boolean> {
if (!zone || zone === this.defaultZone) {
// Default zone always exists
this.existingZonesCache[this.defaultZone] = true;
return true;
}
// Check cache first
if (this.existingZonesCache[zone] !== undefined) {
return this.existingZonesCache[zone];
}
await this.initialize();
// Check metadata for zone
const metadata = await this.getZoneMetadata(zone);
if (metadata) {
// Cache the result
this.existingZonesCache[zone] = true;
return true;
}
// As a fallback, check if the index exists
const indexName = this.getIndexForZone(zone);
const indexExists = await this.client.indices.exists({ index: indexName });
// Cache the result
this.existingZonesCache[zone] = indexExists;
return indexExists;
}
/**
* Create or update a relation between entities
* @param relation Relation to create or update
* @param fromZone Optional zone for the source entity, uses defaultZone if not specified
* @param toZone Optional zone for the target entity, uses defaultZone if not specified
* @param options Optional configuration options
* @param options.autoCreateMissingEntities Whether to automatically create missing entities (default: true)
* @param options.validateZones Whether to validate that zones exist before creating entities or relations (default: true)
*/
async saveRelation(
relation: Omit<ESRelation, 'type' | 'fromZone' | 'toZone'>,
fromZone?: string,
toZone?: string,
options?: {
autoCreateMissingEntities?: boolean;
validateZones?: boolean;
}
): Promise<ESRelation> {
await this.initialize();
// Default to true for backwards compatibility
const autoCreateMissingEntities = options?.autoCreateMissingEntities ?? true;
// Default to true for zone validation
const validateZones = options?.validateZones ?? true;
const actualFromZone = fromZone || this.defaultZone;
const actualToZone = toZone || this.defaultZone;
// Validate that zones exist if required
if (validateZones) {
// Check fromZone
const fromZoneExists = await this.zoneExists(actualFromZone);
if (!fromZoneExists) {
throw new Error(`Cannot create relation: Source zone '${actualFromZone}' does not exist. Create the zone first.`);
}
// Check toZone
const toZoneExists = await this.zoneExists(actualToZone);
if (!toZoneExists) {
throw new Error(`Cannot create relation: Target zone '${actualToZone}' does not exist. Create the zone first.`);
}
}
// Check if both entities exist
const fromEntity = await this.getEntityWithoutUpdatingLastRead(relation.from, actualFromZone);
const toEntity = await this.getEntityWithoutUpdatingLastRead(relation.to, actualToZone);
// If either entity doesn't exist
if (!fromEntity || !toEntity) {
// If auto-creation is disabled, throw an error
if (!autoCreateMissingEntities) {
const missingEntities = [];
if (!fromEntity) {
missingEntities.push(`'${relation.from}' in zone '${actualFromZone}'`);
}
if (!toEntity) {
missingEntities.push(`'${relation.to}' in zone '${actualToZone}'`);
}
throw new Error(`Cannot create relation: Missing entities ${missingEntities.join(' and ')}`);
}
// Otherwise, auto-create the missing entities
if (!fromEntity) {
console.error(`Auto-creating missing entity: ${relation.from} in zone ${actualFromZone}`);
const newEntity = {
type: 'entity',
name: relation.from,
entityType: 'unknown',
observations: [],
readCount: 0,
lastRead: new Date().toISOString(),
lastWrite: new Date().toISOString(),
relevanceScore: 1.0,
zone: actualFromZone
};
// We've already validated the zone, so we can skip validation here
await this.saveEntity({
name: relation.from,
entityType: 'unknown',
observations: [],
relevanceScore: 1.0
}, actualFromZone, { validateZones: false });
}
if (!toEntity) {
console.error(`Auto-creating missing entity: ${relation.to} in zone ${actualToZone}`);
const newEntity = {
type: 'entity',
name: relation.to,
entityType: 'unknown',
observations: [],
readCount: 0,
lastRead: new Date().toISOString(),
lastWrite: new Date().toISOString(),
relevanceScore: 1.0,
zone: actualToZone
};
// We've already validated the zone, so we can skip validation here
await this.saveEntity({
name: relation.to,
entityType: 'unknown',
observations: [],
relevanceScore: 1.0
}, actualToZone, { validateZones: false });
}
}
// Create the relation
const newRelation: ESRelation = {
type: 'relation',
from: relation.from,
fromZone: actualFromZone,
to: relation.to,
toZone: actualToZone,
relationType: relation.relationType
};
const id = `relation:${actualFromZone}:${relation.from}:${relation.relationType}:${actualToZone}:${relation.to}`;
await this.client.index({
index: KG_RELATIONS_INDEX,
id,
document: newRelation,
refresh: true
});
return newRelation;
}
/**
* Delete a relation between entities
* @param from Source entity name
* @param to Target entity name
* @param relationType Relation type
* @param fromZone Optional zone for the source entity, uses defaultZone if not specified
* @param toZone Optional zone for the target entity, uses defaultZone if not specified
*/
async deleteRelation(
from: string,
to: string,
relationType: string,
fromZone?: string,
toZone?: string
): Promise<boolean> {
await this.initialize();
const actualFromZone = fromZone || this.defaultZone;
const actualToZone = toZone || this.defaultZone;
try {
const relationId = `relation:${actualFromZone}:${from}:${relationType}:${actualToZone}:${to}`;
await this.client.delete({
index: KG_RELATIONS_INDEX,
id: relationId,
refresh: true
});
return true;
} catch (error) {
if (error.statusCode === 404) {
return false;
}
throw error;
}
}
/**
* Search for entities and relations in the knowledge graph
* @param params Search parameters
*/
async search(params: ESSearchParams & { zone?: string }): Promise<ESHighlightResponse<ESEntity | ESRelation>> {
const actualZone = params.zone || this.defaultZone;
await this.initialize(actualZone);
const indexName = this.getIndexForZone(actualZone);
// Special handling for wildcard query
if (params.query === '*') {
// console.error(`Performing wildcard search in zone: ${actualZone}, index: ${indexName}`);
try {
// Use match_all query for wildcard
const response = await this.client.search({
index: indexName,
body: {
query: {
match_all: {}
},
sort: [{ lastRead: { order: 'desc' } }],
size: params.limit || 10,
from: params.offset || 0
}
});
// console.error(`Wildcard search results: ${JSON.stringify(response.hits)}`);
return response as unknown as ESHighlightResponse<ESEntity | ESRelation>;
} catch (error) {
console.error(`Error in wildcard search: ${error.message}`);
throw error;
}
}
// Special handling for exact entity name search
if (params.query && !params.query.includes(' ')) {
// console.error(`Performing exact entity name search for "${params.query}" in zone: ${actualZone}, index: ${indexName}`);
try {
// Use match query for exact entity name
const response = await this.client.search({
index: indexName,
body: {
query: {
match: {
name: params.query
}
},
sort: [{ lastRead: { order: 'desc' } }],
size: params.limit || 10,
from: params.offset || 0
}
});
// console.error(`Exact entity name search results: ${JSON.stringify(response.hits)}`);
return response as unknown as ESHighlightResponse<ESEntity | ESRelation>;
} catch (error) {
console.error(`Error in exact entity name search: ${error.message}`);
throw error;
}
}
// Build search query for non-wildcard searches
const query: any = {
bool: {
must: []
}
};
// Process the query to handle boolean operators or fuzzy search notation
if (params.query.includes(' AND ') ||
params.query.includes(' OR ') ||
params.query.includes(' NOT ') ||
params.query.includes('~')) {
// Use Elasticsearch's query_string query for boolean operators and fuzzy search
query.bool.must.push({
query_string: {
query: params.query,
fields: ['name^3', 'entityType^2', 'observations', 'relationType^2'],
default_operator: 'AND',
// Enable fuzzy matching by default
fuzziness: 'AUTO',
// Allow fuzzy matching on all terms unless explicitly disabled
fuzzy_max_expansions: 50,
// Support phrase slop for proximity searches
phrase_slop: 2
}
});
console.error(`Using query_string for advanced query: ${params.query}`);
} else {
// Use multi_match for simple queries without boolean operators
query.bool.must.push({
multi_match: {
query: params.query,
fields: ['name^3', 'entityType^2', 'observations', 'relationType^2'],
// Enable fuzzy matching by default
fuzziness: 'AUTO'
}
});
console.error(`Using multi_match for simple query: ${params.query}`);
}
// Add entityType filter if specified
if (params.entityTypes && params.entityTypes.length > 0) {
// Use a more flexible filter with "should" (equivalent to OR) to match any of the entity types
const entityTypeFilters = params.entityTypes.map(type => ({
match: {
entityType: {
query: type,
operator: "and"
}
}
}));
query.bool.must.push({
bool: {
should: entityTypeFilters,
minimum_should_match: 1
}
});
console.error(`Applied entity type filters: ${JSON.stringify(entityTypeFilters)}`);
}
// Log validation to ensure zone filter is being applied
console.error(`Searching in zone: ${actualZone}, index: ${indexName}, query: ${JSON.stringify(query)}`);
// Set up sort order
let sort: any[] = [];
if (params.sortBy === 'recent') {
sort = [{ lastRead: { order: 'desc' } }];
} else if (params.sortBy === 'importance') {
sort = [
// Always sort by relevanceScore in descending order (highest first)
{ relevanceScore: { order: 'desc' } }
];
} else {
// Default is by relevance (using _score)
sort = [{ _score: { order: 'desc' } }];
}
try {
// Execute search
const response = await this.client.search({
index: indexName,
body: {
query,
sort,
highlight: {
fields: {
name: {},
observations: {},
entityType: {}
}
},
size: params.limit || 10,
from: params.offset || 0
}
});
console.error(`Search results: ${JSON.stringify(response.hits)}`);
return response as unknown as ESHighlightResponse<ESEntity | ESRelation>;
} catch (error) {
console.error(`Error in search: ${error.message}`);
throw error;
}
}
/**
* Get all entities related to a given entity, up to a certain depth
* @param name Entity name
* @param maxDepth Maximum traversal depth
* @param zone Optional memory zone name, uses defaultZone if not specified
*/
async getRelatedEntities(
name: string,
maxDepth: number = 1,
zone?: string
): Promise<{
entities: ESEntity[],
relations: ESRelation[]
}> {
const actualZone = zone || this.defaultZone;
await this.initialize(actualZone);
// Start with the root entity
const rootEntity = await this.getEntity(name, actualZone);
if (!rootEntity) {
return { entities: [], relations: [] };
}
// Keep track of entities and relations we've found
const entitiesMap = new Map<string, ESEntity>();
entitiesMap.set(`${actualZone}:${rootEntity.name}`, rootEntity);
const relationsMap = new Map<string, ESRelation>();
const visitedEntities = new Set<string>();
// Queue of entities to process, with their depth
const queue: Array<{ entity: ESEntity, zone: string, depth: number }> = [
{ entity: rootEntity, zone: actualZone, depth: 0 }
];
while (queue.length > 0) {
const { entity, zone: entityZone, depth } = queue.shift()!;
const entityKey = `${entityZone}:${entity.name}`;
// Skip if we've already processed this entity or if we've reached max depth
if (visitedEntities.has(entityKey) || depth >= maxDepth) {
continue;
}
visitedEntities.add(entityKey);
// Find all relations involving this entity
const fromResponse = await this.client.search({
index: KG_RELATIONS_INDEX,
body: {
query: {
bool: {
must: [
{ term: { from: entity.name } },
{ term: { "fromZone.keyword": entityZone } }
]
}
},
size: 1000
}
});
const toResponse = await this.client.search({
index: KG_RELATIONS_INDEX,
body: {
query: {
bool: {
must: [
{ term: { to: entity.name } },
{ term: { "toZone.keyword": entityZone } }
]
}
},
size: 1000
}
});
// Process relations where this entity is the source
const fromHits = (fromResponse as unknown as ESSearchResponse<ESRelation>).hits.hits;
for (const hit of fromHits) {
const relation = hit._source;
const relationKey = `${relation.fromZone}:${relation.from}|${relation.relationType}|${relation.toZone}:${relation.to}`;
// Skip if we've already processed this relation
if (relationsMap.has(relationKey)) {
continue;
}
relationsMap.set(relationKey, relation);
// Process the target entity
const otherEntityKey = `${relation.toZone}:${relation.to}`;
if (!entitiesMap.has(otherEntityKey)) {
// Fetch the other entity
const otherEntity = await this.getEntity(relation.to, relation.toZone);
if (otherEntity) {
entitiesMap.set(otherEntityKey, otherEntity);
// Add the other entity to the queue if we haven't reached max depth
if (depth < maxDepth - 1) {
queue.push({ entity: otherEntity, zone: relation.toZone, depth: depth + 1 });
}
}
}
}
// Process relations where this entity is the target
const toHits = (toResponse as unknown as ESSearchResponse<ESRelation>).hits.hits;
for (const hit of toHits) {
const relation = hit._source;
const relationKey = `${relation.fromZone}:${relation.from}|${relation.relationType}|${relation.toZone}:${relation.to}`;
// Skip if we've already processed this relation
if (relationsMap.has(relationKey)) {
continue;
}
relationsMap.set(relationKey, relation);
// Process the source entity
const otherEntityKey = `${relation.fromZone}:${relation.from}`;
if (!entitiesMap.has(otherEntityKey)) {
// Fetch the other entity
const otherEntity = await this.getEntity(relation.from, relation.fromZone);
if (otherEntity) {
entitiesMap.set(otherEntityKey, otherEntity);
// Add the other entity to the queue if we haven't reached max depth
if (depth < maxDepth - 1) {
queue.push({ entity: otherEntity, zone: relation.fromZone, depth: depth + 1 });
}
}
}
}
}
return {
entities: Array.from(entitiesMap.values()),
relations: Array.from(relationsMap.values())
};
}
/**
* Import data into the knowledge graph
* @param data Array of entities and relations to import
* @param zone Optional memory zone for entities, uses defaultZone if not specified
* @param options Optional configuration options
* @param options.validateZones Whether to validate that zones exist before importing (default: true)
*/
async importData(
data: Array<ESEntity | ESRelation>,
zone?: string,
options?: {
validateZones?: boolean;
}
): Promise<{
entitiesAdded: number;
relationsAdded: number;
invalidRelations?: Array<{relation: ESRelation, reason: string}>;
}> {
const actualZone = zone || this.defaultZone;
await this.initialize(actualZone);
// Default to true for zone validation
const validateZones = options?.validateZones ?? true;
// Validate that zone exists if required
if (validateZones && actualZone !== this.defaultZone) {
const zoneExists = await this.zoneExists(actualZone);
if (!zoneExists) {
throw new Error(`Cannot import data: Zone '${actualZone}' does not exist. Create the zone first.`);
}
}
let entitiesAdded = 0;
let relationsAdded = 0;
const invalidRelations: Array<{relation: ESRelation, reason: string}> = [];
// Process entities first, since relations depend on them
const entities = data.filter(item => item.type === 'entity') as ESEntity[];
const entityOperations: any[] = [];
for (const entity of entities) {
// Add zone information if not already present
const entityWithZone = {
...entity,
zone: entity.zone || actualZone
};
const id = `entity:${entity.name}`;
entityOperations.push({ index: { _index: this.getIndexForZone(actualZone), _id: id } });
entityOperations.push(entityWithZone);
entitiesAdded++;
}
if (entityOperations.length > 0) {
await this.client.bulk({
operations: entityOperations,
refresh: true
});
}
// Now process relations
const relations = data.filter(item => item.type === 'relation') as ESRelation[];
const relationOperations: any[] = [];
for (const relation of relations) {
// For relations with explicit zones
if (relation.fromZone !== undefined && relation.toZone !== undefined) {
// If zone validation is enabled, check that both zones exist
if (validateZones) {
// Check fromZone if it's not the default zone
if (relation.fromZone !== this.defaultZone) {
const fromZoneExists = await this.zoneExists(relation.fromZone);
if (!fromZoneExists) {
invalidRelations.push({
relation,
reason: `Source zone '${relation.fromZone}' does not exist. Create the zone first.`
});
continue;
}
}
// Check toZone if it's not the default zone
if (relation.toZone !== this.defaultZone) {
const toZoneExists = await this.zoneExists(relation.toZone);
if (!toZoneExists) {
invalidRelations.push({
relation,
reason: `Target zone '${relation.toZone}' does not exist. Create the zone first.`
});
continue;
}
}
}
// Verify that both entities exist
const fromEntity = await this.getEntityWithoutUpdatingLastRead(relation.from, relation.fromZone);
const toEntity = await this.getEntityWithoutUpdatingLastRead(relation.to, relation.toZone);
if (!fromEntity) {
invalidRelations.push({
relation,
reason: `Source entity '${relation.from}' in zone '${relation.fromZone}' does not exist`
});
continue;
}
if (!toEntity) {
invalidRelations.push({
relation,
reason: `Target entity '${relation.to}' in zone '${relation.toZone}' does not exist`
});
continue;
}
const id = `relation:${relation.fromZone}:${relation.from}:${relation.relationType}:${relation.toZone}:${relation.to}`;
relationOperations.push({ index: { _index: KG_RELATIONS_INDEX, _id: id } });
relationOperations.push(relation);
relationsAdded++;
} else {
// Old format - needs to be converted
// For backward compatibility, assume both entities are in the specified zone
const fromZone = actualZone;
const toZone = actualZone;
// Verify that both entities exist
const fromEntity = await this.getEntityWithoutUpdatingLastRead(relation.from, fromZone);
const toEntity = await this.getEntityWithoutUpdatingLastRead(relation.to, toZone);
if (!fromEntity) {
invalidRelations.push({
relation,
reason: `Source entity '${relation.from}' in zone '${fromZone}' does not exist`
});
continue;
}
if (!toEntity) {
invalidRelations.push({
relation,
reason: `Target entity '${relation.to}' in zone '${toZone}' does not exist`
});
continue;
}
// Convert to new format
const newRelation: ESRelation = {
type: 'relation',
from: relation.from,
fromZone,
to: relation.to,
toZone,
relationType: relation.relationType
};
const id = `relation:${fromZone}:${relation.from}:${relation.relationType}:${toZone}:${relation.to}`;
relationOperations.push({ index: { _index: KG_RELATIONS_INDEX, _id: id } });
relationOperations.push(newRelation);
relationsAdded++;
}
}
if (relationOperations.length > 0) {
await this.client.bulk({
operations: relationOperations,
refresh: true
});
}
return {
entitiesAdded,
relationsAdded,
invalidRelations: invalidRelations.length ? invalidRelations : undefined
};
}
/**
* Import data into the knowledge graph, recreating zones as needed
* @param data Export data containing entities, relations, and zone metadata
*/
async importAllData(data: {
entities: ESEntity[],
relations: ESRelation[],
zones: ZoneMetadata[]
}): Promise<{
zonesAdded: number;
entitiesAdded: number;
relationsAdded: number;
}> {
await this.initialize();
let zonesAdded = 0;
let entitiesAdded = 0;
let relationsAdded = 0;
// First create all zones
for (const zone of data.zones) {
if (zone.name !== 'default') {
await this.addMemoryZone(zone.name, zone.description, zone.config);
// addMemoryZone already updates the cache
zonesAdded++;
} else {
// Make sure default zone is in the cache
this.existingZonesCache['default'] = true;
}
}
// Import entities by zone
const entitiesByZone: Record<string, ESEntity[]> = {};
for (const entity of data.entities) {
const zone = entity.zone || 'default';
if (!entitiesByZone[zone]) {
entitiesByZone[zone] = [];
}
entitiesByZone[zone].push(entity);
}
// Import entities for each zone
for (const [zone, entities] of Object.entries(entitiesByZone)) {
const result = await this.importData(entities, zone);
entitiesAdded += result.entitiesAdded;
}
// Import all relations
if (data.relations.length > 0) {
const result = await this.importData(data.relations);
relationsAdded = result.relationsAdded;
}
return {
zonesAdded,
entitiesAdded,
relationsAdded
};
}
/**
* Export all data from a knowledge graph
* @param zone Optional memory zone for entities, uses defaultZone if not specified
*/
async exportData(zone?: string): Promise<Array<ESEntity | ESRelation>> {
const actualZone = zone || this.defaultZone;
await this.initialize(actualZone);
// Fetch all entities from the specified zone
const indexName = this.getIndexForZone(actualZone);
const entityResponse = await this.client.search({
index: indexName,
body: {
query: { term: { type: 'entity' } },
size: 10000
}
});
const entities = entityResponse.hits.hits.map(hit => hit._source) as ESEntity[];
// Fetch all relations involving entities in this zone
const relationResponse = await this.client.search({
index: KG_RELATIONS_INDEX,
body: {
query: {
bool: {
should: [
{ term: { fromZone: actualZone } },
{ term: { toZone: actualZone } }
],
minimum_should_match: 1
}
},
size: 10000
}
});
const relations = relationResponse.hits.hits.map(hit => hit._source) as ESRelation[];
// Combine entities and relations
return [...entities, ...relations];
}
/**
* Get all relations involving a set of entities
* @param entityNames Array of entity names
* @param zone Optional memory zone for all entities, uses defaultZone if not specified
*/
async getRelationsForEntities(
entityNames: string[],
zone?: string
): Promise<{
relations: ESRelation[]
}> {
const actualZone = zone || this.defaultZone;
await this.initialize(actualZone);
if (entityNames.length === 0) {
return { relations: [] };
}
// Find all relations where any of these entities are involved
// We need to search for both directions - as source and as target
const fromQuery = entityNames.map(name => ({
bool: {
must: [
{ term: { from: name } },
{ term: { "fromZone.keyword": actualZone } }
]
}
}));
const toQuery = entityNames.map(name => ({
bool: {
must: [
{ term: { to: name } },
{ term: { "toZone.keyword": actualZone } }
]
}
}));
const response = await this.client.search({
index: KG_RELATIONS_INDEX,
body: {
query: {
bool: {
should: [...fromQuery, ...toQuery],
minimum_should_match: 1
}
},
size: 1000
}
});
const relations = (response as unknown as ESSearchResponse<ESRelation>)
.hits.hits
.map(hit => hit._source);
return { relations };
}
/**
* Save or update zone metadata
* @param name Zone name
* @param description Optional description
* @param config Optional configuration
*/
private async saveZoneMetadata(
name: string,
description?: string,
config?: Record<string, any>
): Promise<void> {
await this.initialize();
const now = new Date().toISOString();
// Check if zone metadata exists
let existing: ZoneMetadata | null = null;
try {
const response = await this.client.get({
index: KG_METADATA_INDEX,
id: `zone:${name}`
});
existing = response._source as ZoneMetadata;
} catch (error) {
// Zone doesn't exist yet
}
const metadata: ZoneMetadata = {
name,
description: description || existing?.description,
shortDescription: existing?.shortDescription,
createdAt: existing?.createdAt || now,
lastModified: now,
config: config || existing?.config
};
await this.client.index({
index: KG_METADATA_INDEX,
id: `zone:${name}`,
document: metadata,
refresh: true
});
}
/**
* List all available memory zones
* @param reason Optional reason for listing zones, used for AI filtering
*/
async listMemoryZones(reason?: string): Promise<ZoneMetadata[]> {
await this.initialize();
try {
// First try getting zones from metadata
const response = await this.client.search({
index: KG_METADATA_INDEX,
body: {
query: { match_all: {} },
size: 1000
}
});
const zones = response.hits.hits.map(hit => hit._source as ZoneMetadata);
if (zones.length > 0) {
// Update cache with all known zones
zones.forEach(zone => {
this.existingZonesCache[zone.name] = true;
});
return zones;
}
} catch (error) {
console.warn('Error getting zones from metadata, falling back to index detection:', error);
}
// Fallback to listing indices (for backward compatibility)
const indicesResponse = await this.client.indices.get({
index: `${KG_INDEX_PREFIX}@*`
});
// Extract zone names from index names
const zoneNames = Object.keys(indicesResponse)
.filter(indexName => indexName.startsWith(`${KG_INDEX_PREFIX}@`))
.map(indexName => indexName.substring(KG_INDEX_PREFIX.length + 1)); // +1 for the @ symbol
// Convert to metadata format
const now = new Date().toISOString();
const zones = zoneNames.map(name => ({
name,
createdAt: now,
lastModified: now
}));
// Update cache with all detected zones
zones.forEach(zone => {
this.existingZonesCache[zone.name] = true;
});
// Save detected zones to metadata for future
for (const zone of zones) {
await this.saveZoneMetadata(zone.name, `Zone detected from index: ${getIndexName(zone.name)}`);
}
return zones;
}
/**
* Add a new memory zone (creates the index if it doesn't exist)
* @param zone Zone name to add
* @param description Optional description of the zone
* @param config Optional configuration for the zone
*/
async addMemoryZone(
zone: string,
description?: string,
config?: Record<string, any>
): Promise<boolean> {
if (!zone || zone === 'default') {
throw new Error('Invalid zone name. Cannot be empty or "default".');
}
// Initialize the index for this zone
await this.initialize(zone);
// Add to metadata
await this.saveZoneMetadata(zone, description, config);
// Update the cache
this.existingZonesCache[zone] = true;
return true;
}
/**
* Get metadata for a specific zone
* @param zone Zone name
*/
async getZoneMetadata(zone: string): Promise<ZoneMetadata | null> {
await this.initialize();
try {
const response = await this.client.get({
index: KG_METADATA_INDEX,
id: `zone:${zone}`
});
return response._source as ZoneMetadata;
} catch (error) {
return null;
}
}
/**
* Delete a memory zone and all its data
* @param zone Zone name to delete
*/
async deleteMemoryZone(zone: string): Promise<boolean> {
if (zone === 'default') {
throw new Error('Cannot delete the default zone.');
}
await this.initialize();
try {
const indexName = this.getIndexForZone(zone);
// Check if index exists before trying to delete it
const indexExists = await this.client.indices.exists({
index: indexName
});
if (indexExists) {
// Delete the index
await this.client.indices.delete({
index: indexName
});
console.error(`Deleted index: ${indexName}`);
}
// Check if metadata exists before trying to delete it
try {
const metadataExists = await this.client.exists({
index: KG_METADATA_INDEX,
id: `zone:${zone}`
});
if (metadataExists) {
// Delete from metadata
await this.client.delete({
index: KG_METADATA_INDEX,
id: `zone:${zone}`
});
}
} catch (metadataError) {
// Log but continue even if metadata deletion fails
console.error(`Warning: Error checking/deleting metadata for zone ${zone}:`, metadataError.message);
}
// Remove from initialized indices cache
this.initializedIndices.delete(indexName);
// Update the zones cache
delete this.existingZonesCache[zone];
// Clean up relations for this zone
try {
await this.client.deleteByQuery({
index: KG_RELATIONS_INDEX,
body: {
query: {
bool: {
should: [
{ term: { fromZone: zone } },
{ term: { toZone: zone } }
],
minimum_should_match: 1
}
}
},
refresh: true
});
} catch (relationError) {
// Log but continue even if relation cleanup fails
console.error(`Warning: Error cleaning up relations for zone ${zone}:`, relationError.message);
}
return true;
} catch (error) {
console.error(`Error deleting zone ${zone}:`, error);
return false;
}
}
/**
* Get statistics for a memory zone
* @param zone Zone name, uses defaultZone if not specified
*/
async getMemoryZoneStats(zone?: string): Promise<{
zone: string;
entityCount: number;
relationCount: number;
entityTypes: Record<string, number>;
relationTypes: Record<string, number>;
}> {
const actualZone = zone || this.defaultZone;
await this.initialize(actualZone);
const indexName = this.getIndexForZone(actualZone);
// Get total counts
const countResponse = await this.client.count({
index: indexName,
body: {
query: {
term: { type: 'entity' }
}
}
});
const entityCount = countResponse.count;
const relationCountResponse = await this.client.count({
index: indexName,
body: {
query: {
term: { type: 'relation' }
}
}
});
const relationCount = relationCountResponse.count;
// Get entity type distribution
const entityTypesResponse = await this.client.search({
index: indexName,
body: {
size: 0,
query: {
term: { type: 'entity' }
},
aggs: {
entity_types: {
terms: {
field: 'entityType',
size: 100
}
}
}
}
});
const entityTypes: Record<string, number> = {};
const entityTypeAggs = entityTypesResponse.aggregations as any;
const entityTypeBuckets = entityTypeAggs?.entity_types?.buckets || [];
entityTypeBuckets.forEach((bucket: any) => {
entityTypes[bucket.key] = bucket.doc_count;
});
// Get relation type distribution
const relationTypesResponse = await this.client.search({
index: indexName,
body: {
size: 0,
query: {
term: { type: 'relation' }
},
aggs: {
relation_types: {
terms: {
field: 'relationType',
size: 100
}
}
}
}
});
const relationTypes: Record<string, number> = {};
const relationTypeAggs = relationTypesResponse.aggregations as any;
const relationTypeBuckets = relationTypeAggs?.relation_types?.buckets || [];
relationTypeBuckets.forEach((bucket: any) => {
relationTypes[bucket.key] = bucket.doc_count;
});
return {
zone: actualZone,
entityCount,
relationCount,
entityTypes,
relationTypes
};
}
/**
* Export all knowledge graph data, optionally limiting to specific zones
* @param zones Optional array of zone names to export, exports all zones if not specified
*/
async exportAllData(zones?: string[]): Promise<{
entities: ESEntity[],
relations: ESRelation[],
zones: ZoneMetadata[]
}> {
await this.initialize();
// Get all zones or filter to specified zones
const allZones = await this.listMemoryZones();
const zonesToExport = zones
? allZones.filter(zone => zones.includes(zone.name))
: allZones;
if (zonesToExport.length === 0) {
return { entities: [], relations: [], zones: [] };
}
// Collect all entities from each zone
const entities: ESEntity[] = [];
for (const zone of zonesToExport) {
const zoneData = await this.exportData(zone.name);
const zoneEntities = zoneData.filter(item => item.type === 'entity') as ESEntity[];
entities.push(...zoneEntities);
}
// Get all relations
let relations: ESRelation[] = [];
if (zones) {
// If specific zones are specified, only get relations involving those zones
const relationResponse = await this.client.search({
index: KG_RELATIONS_INDEX,
body: {
query: {
bool: {
should: [
...zonesToExport.map(zone => ({ term: { fromZone: zone.name } })),
...zonesToExport.map(zone => ({ term: { toZone: zone.name } }))
],
minimum_should_match: 1
}
},
size: 10000
}
});
relations = relationResponse.hits.hits.map(hit => hit._source) as ESRelation[];
} else {
// If no zones specified, get all relations
const relationResponse = await this.client.search({
index: KG_RELATIONS_INDEX,
body: {
query: { match_all: {} },
size: 10000
}
});
relations = relationResponse.hits.hits.map(hit => hit._source) as ESRelation[];
}
return {
entities,
relations,
zones: zonesToExport
};
}
/**
* Add observations to an existing entity
* @param name Entity name
* @param observations Array of observation strings to add
* @param zone Optional memory zone name, uses defaultZone if not specified
* @returns The updated entity
*/
async addObservations(name: string, observations: string[], zone?: string): Promise<ESEntity> {
const actualZone = zone || this.defaultZone;
// Get existing entity
const entity = await this.getEntity(name, actualZone);
if (!entity) {
throw new Error(`Entity "${name}" not found in zone "${actualZone}"`);
}
// Add new observations to the existing ones
const updatedObservations = [
...entity.observations,
...observations
];
// Update the entity
const updatedEntity = await this.saveEntity({
name: entity.name,
entityType: entity.entityType,
observations: updatedObservations,
relevanceScore: entity.relevanceScore
}, actualZone);
return updatedEntity;
}
/**
* Mark an entity as important or not important
* @param name Entity name
* @param important Whether the entity is important
* @param zone Optional memory zone name, uses defaultZone if not specified
* @param options Optional configuration options
* @param options.autoCreateMissingEntities Whether to automatically create missing entities (default: false)
* @returns The updated entity
*/
async markImportant(
name: string,
important: boolean,
zone?: string,
options?: {
autoCreateMissingEntities?: boolean;
}
): Promise<ESEntity> {
return this.updateEntityRelevanceScore(name, important ? 10 : 0.1, zone, options);
}
/**
* Mark an entity as important or not important
* @param name Entity name
* @param important Whether the entity is important
* @param zone Optional memory zone name, uses defaultZone if not specified
* @param options Optional configuration options
* @param options.autoCreateMissingEntities Whether to automatically create missing entities (default: false)
* @returns The updated entity
*/
async updateEntityRelevanceScore(
name: string,
ratio: number,
zone?: string,
options?: {
autoCreateMissingEntities?: boolean;
}
): Promise<ESEntity> {
const actualZone = zone || this.defaultZone;
// Default to false for auto-creation (different from saveRelation)
const autoCreateMissingEntities = options?.autoCreateMissingEntities ?? false;
// Get existing entity
// Get existing entity
let entity = await this.getEntity(name, actualZone);
// If entity doesn't exist
if (!entity) {
if (autoCreateMissingEntities) {
// Auto-create the entity with unknown type
entity = await this.saveEntity({
name: name,
entityType: 'unknown',
observations: [],
relevanceScore: 1.0
}, actualZone);
} else {
throw new Error(`Entity "${name}" not found in zone "${actualZone}"`);
}
}
// Calculate the new relevance score
// If marking as important, multiply by 10
// If removing importance, divide by 10
const baseRelevanceScore = entity.relevanceScore || 1.0;
const newRelevanceScore = ratio > 1.0
? Math.min(25, baseRelevanceScore * ratio)
: Math.max(0.01, baseRelevanceScore * ratio);
// Update entity with new relevance score
const updatedEntity = await this.saveEntity({
name: entity.name,
entityType: entity.entityType,
observations: entity.observations,
relevanceScore: newRelevanceScore
}, actualZone);
return updatedEntity;
}
/**
* Get recent entities
* @param limit Maximum number of entities to return
* @param includeObservations Whether to include observations
* @param zone Optional memory zone name, uses defaultZone if not specified
* @returns Array of recent entities
*/
async getRecentEntities(limit: number, includeObservations: boolean, zone?: string): Promise<ESEntity[]> {
const actualZone = zone || this.defaultZone;
// Search with empty query but sort by recency
const searchParams: ESSearchParams = {
query: "*", // Use wildcard instead of empty query to match all documents
limit: limit,
sortBy: 'recent', // Sort by recency
includeObservations
};
// Add zone if specified
if (actualZone) {
(searchParams as any).zone = actualZone;
}
const results = await this.search(searchParams);
// Filter to only include entities
return results.hits.hits
.filter((hit: any) => hit._source.type === 'entity')
.map((hit: any) => hit._source);
}
/**
* Copy entities from one zone to another
* @param entityNames Array of entity names to copy
* @param sourceZone Source zone to copy from
* @param targetZone Target zone to copy to
* @param options Optional configuration
* @param options.copyRelations Whether to copy relations involving these entities (default: true)
* @param options.overwrite Whether to overwrite entities if they already exist in target zone (default: false)
* @returns Result of the copy operation
*/
async copyEntitiesBetweenZones(
entityNames: string[],
sourceZone: string,
targetZone: string,
options?: {
copyRelations?: boolean;
overwrite?: boolean;
}
): Promise<{
entitiesCopied: string[];
entitiesSkipped: { name: string; reason: string }[];
relationsCopied: number;
}> {
if (sourceZone === targetZone) {
throw new Error('Source and target zones must be different');
}
// Default options
const copyRelations = options?.copyRelations !== false;
const overwrite = options?.overwrite === true;
await this.initialize(sourceZone);
await this.initialize(targetZone);
const result = {
entitiesCopied: [] as string[],
entitiesSkipped: [] as { name: string; reason: string }[],
relationsCopied: 0
};
// Get entities from source zone
for (const name of entityNames) {
// Get the entity from source zone
const entity = await this.getEntityWithoutUpdatingLastRead(name, sourceZone);
if (!entity) {
result.entitiesSkipped.push({
name,
reason: `Entity not found in source zone '${sourceZone}'`
});
continue;
}
// Check if entity exists in target zone
const existingEntity = await this.getEntityWithoutUpdatingLastRead(name, targetZone);
if (existingEntity && !overwrite) {
result.entitiesSkipped.push({
name,
reason: `Entity already exists in target zone '${targetZone}' and overwrite is disabled`
});
continue;
}
// Copy the entity to target zone
const { ...entityCopy } = entity;
delete entityCopy.zone; // Zone will be set by saveEntity
try {
await this.saveEntity(entityCopy, targetZone);
result.entitiesCopied.push(name);
} catch (error) {
result.entitiesSkipped.push({
name,
reason: `Error copying entity: ${(error as Error).message}`
});
continue;
}
}
// Copy relations if requested
if (copyRelations && result.entitiesCopied.length > 0) {
// Get all relations for these entities in source zone
const { relations } = await this.getRelationsForEntities(result.entitiesCopied, sourceZone);
// Filter to only include relations where both entities were copied
// or relations between copied entities and entities that already exist in target zone
const relationsToCreate: ESRelation[] = [];
for (const relation of relations) {
let fromExists = result.entitiesCopied.includes(relation.from);
let toExists = result.entitiesCopied.includes(relation.to);
// If one side of the relation wasn't copied, check if it exists in target zone
if (!fromExists) {
const fromEntityInTarget = await this.getEntityWithoutUpdatingLastRead(relation.from, targetZone);
fromExists = !!fromEntityInTarget;
}
if (!toExists) {
const toEntityInTarget = await this.getEntityWithoutUpdatingLastRead(relation.to, targetZone);
toExists = !!toEntityInTarget;
}
// Only create relations where both sides exist
if (fromExists && toExists) {
relationsToCreate.push(relation);
}
}
// Save the filtered relations
for (const relation of relationsToCreate) {
try {
await this.saveRelation({
from: relation.from,
to: relation.to,
relationType: relation.relationType
}, targetZone, targetZone);
result.relationsCopied++;
} catch (error) {
console.error(`Error copying relation from ${relation.from} to ${relation.to}: ${(error as Error).message}`);
}
}
}
return result;
}
/**
* Move entities from one zone to another (copy + delete from source)
* @param entityNames Array of entity names to move
* @param sourceZone Source zone to move from
* @param targetZone Target zone to move to
* @param options Optional configuration
* @param options.moveRelations Whether to move relations involving these entities (default: true)
* @param options.overwrite Whether to overwrite entities if they already exist in target zone (default: false)
* @returns Result of the move operation
*/
async moveEntitiesBetweenZones(
entityNames: string[],
sourceZone: string,
targetZone: string,
options?: {
moveRelations?: boolean;
overwrite?: boolean;
}
): Promise<{
entitiesMoved: string[];
entitiesSkipped: { name: string; reason: string }[];
relationsMoved: number;
}> {
if (sourceZone === targetZone) {
throw new Error('Source and target zones must be different');
}
// Default options
const moveRelations = options?.moveRelations !== false;
// First copy the entities
const copyResult = await this.copyEntitiesBetweenZones(
entityNames,
sourceZone,
targetZone,
{
copyRelations: moveRelations,
overwrite: options?.overwrite
}
);
const result = {
entitiesMoved: [] as string[],
entitiesSkipped: copyResult.entitiesSkipped,
relationsMoved: copyResult.relationsCopied
};
// Delete copied entities from source zone
for (const name of copyResult.entitiesCopied) {
try {
// Don't cascade relations when deleting from source, as we've already copied them
await this.deleteEntity(name, sourceZone, { cascadeRelations: false });
result.entitiesMoved.push(name);
} catch (error) {
// If deletion fails, add to skipped list but keep the entity in the moved list
// since it was successfully copied
result.entitiesSkipped.push({
name,
reason: `Entity was copied but could not be deleted from source: ${(error as Error).message}`
});
}
}
return result;
}
/**
* Merge two or more zones into a target zone
* @param sourceZones Array of source zone names to merge from
* @param targetZone Target zone to merge into
* @param options Optional configuration
* @param options.deleteSourceZones Whether to delete source zones after merging (default: false)
* @param options.overwriteConflicts How to handle entity name conflicts (default: 'skip')
* @returns Result of the merge operation
*/
async mergeZones(
sourceZones: string[],
targetZone: string,
options?: {
deleteSourceZones?: boolean;
overwriteConflicts?: 'skip' | 'overwrite' | 'rename';
}
): Promise<{
mergedZones: string[];
failedZones: { zone: string; reason: string }[];
entitiesCopied: number;
entitiesSkipped: number;
relationsCopied: number;
}> {
// Validate parameters
if (sourceZones.includes(targetZone)) {
throw new Error('Target zone cannot be included in source zones');
}
if (sourceZones.length === 0) {
throw new Error('At least one source zone must be specified');
}
// Default options
const deleteSourceZones = options?.deleteSourceZones === true;
const overwriteConflicts = options?.overwriteConflicts || 'skip';
// Initialize target zone
await this.initialize(targetZone);
const result = {
mergedZones: [] as string[],
failedZones: [] as { zone: string; reason: string }[],
entitiesCopied: 0,
entitiesSkipped: 0,
relationsCopied: 0
};
// Process each source zone
for (const sourceZone of sourceZones) {
try {
// Get all entities from source zone
const allEntities = await this.searchEntities({
query: '*',
limit: 10000,
zone: sourceZone
});
if (allEntities.length === 0) {
result.failedZones.push({
zone: sourceZone,
reason: 'Zone has no entities'
});
continue;
}
// Extract entity names
const entityNames = allEntities.map(entity => entity.name);
// Process according to conflict resolution strategy
if (overwriteConflicts === 'rename') {
// For 'rename' strategy, we need to check each entity and rename if necessary
for (const entity of allEntities) {
const existingEntity = await this.getEntityWithoutUpdatingLastRead(entity.name, targetZone);
if (existingEntity) {
// Entity exists in target zone, generate a new name
const newName = `${entity.name}_from_${sourceZone}`;
// Create a copy with the new name
const entityCopy = { ...entity, name: newName };
delete entityCopy.zone; // Zone will be set by saveEntity
try {
await this.saveEntity(entityCopy, targetZone);
result.entitiesCopied++;
} catch (error) {
result.entitiesSkipped++;
console.error(`Error copying entity ${entity.name} with new name ${newName}: ${(error as Error).message}`);
}
} else {
// Entity doesn't exist, copy as is
const entityCopy = { ...entity };
delete entityCopy.zone; // Zone will be set by saveEntity
try {
await this.saveEntity(entityCopy, targetZone);
result.entitiesCopied++;
} catch (error) {
result.entitiesSkipped++;
console.error(`Error copying entity ${entity.name}: ${(error as Error).message}`);
}
}
}
// Now copy relations, adjusting for renamed entities
const { relations } = await this.getRelationsForEntities(entityNames, sourceZone);
for (const relation of relations) {
try {
// Check if entities were renamed
let fromName = relation.from;
let toName = relation.to;
const fromEntityInTarget = await this.getEntityWithoutUpdatingLastRead(fromName, targetZone);
if (!fromEntityInTarget) {
// Check if it was renamed
const renamedFromName = `${fromName}_from_${sourceZone}`;
const renamedFromEntityInTarget = await this.getEntityWithoutUpdatingLastRead(renamedFromName, targetZone);
if (renamedFromEntityInTarget) {
fromName = renamedFromName;
}
}
const toEntityInTarget = await this.getEntityWithoutUpdatingLastRead(toName, targetZone);
if (!toEntityInTarget) {
// Check if it was renamed
const renamedToName = `${toName}_from_${sourceZone}`;
const renamedToEntityInTarget = await this.getEntityWithoutUpdatingLastRead(renamedToName, targetZone);
if (renamedToEntityInTarget) {
toName = renamedToName;
}
}
// Only create relation if both entities exist
if (await this.getEntityWithoutUpdatingLastRead(fromName, targetZone) &&
await this.getEntityWithoutUpdatingLastRead(toName, targetZone)) {
await this.saveRelation({
from: fromName,
to: toName,
relationType: relation.relationType
}, targetZone, targetZone);
result.relationsCopied++;
}
} catch (error) {
console.error(`Error copying relation from ${relation.from} to ${relation.to}: ${(error as Error).message}`);
}
}
} else {
// For 'skip' or 'overwrite' strategy, use copyEntitiesBetweenZones
const copyResult = await this.copyEntitiesBetweenZones(
entityNames,
sourceZone,
targetZone,
{
copyRelations: true,
overwrite: overwriteConflicts === 'overwrite'
}
);
result.entitiesCopied += copyResult.entitiesCopied.length;
result.entitiesSkipped += copyResult.entitiesSkipped.length;
result.relationsCopied += copyResult.relationsCopied;
}
// Mark as successfully merged
result.mergedZones.push(sourceZone);
// Delete source zone if requested
if (deleteSourceZones) {
await this.deleteMemoryZone(sourceZone);
}
} catch (error) {
result.failedZones.push({
zone: sourceZone,
reason: (error as Error).message
});
}
}
return result;
}
/**
* Search for entities by name or other attributes
* @param params Search parameters
* @returns Array of matching entities
*/
async searchEntities(params: {
query: string;
entityTypes?: string[];
limit?: number;
includeObservations?: boolean;
zone?: string;
}): Promise<ESEntity[]> {
// Use existing search method with appropriate parameters
const searchResponse = await this.search({
query: params.query,
entityTypes: params.entityTypes,
limit: params.limit,
offset: 0,
zone: params.zone
});
// Extract entities from the search response
const entities: ESEntity[] = [];
if (searchResponse && searchResponse.hits && searchResponse.hits.hits) {
for (const hit of searchResponse.hits.hits) {
if (hit._source && hit._source.type === 'entity') {
entities.push(hit._source as ESEntity);
}
}
}
return entities;
}
/**
* Update zone metadata with new descriptions
* @param name Zone name
* @param description Full description
* @param shortDescription Short description
* @param config Optional configuration
*/
async updateZoneDescriptions(
name: string,
description: string,
shortDescription: string,
config?: Record<string, any>
): Promise<void> {
await this.initialize();
const now = new Date().toISOString();
// Check if zone metadata exists
let existing: ZoneMetadata | null = null;
try {
const response = await this.client.get({
index: KG_METADATA_INDEX,
id: `zone:${name}`
});
existing = response._source as ZoneMetadata;
} catch (error) {
// Zone doesn't exist yet, create it first
if (!await this.zoneExists(name)) {
await this.addMemoryZone(name);
}
}
const metadata: ZoneMetadata = {
name,
description,
shortDescription,
createdAt: existing?.createdAt || now,
lastModified: now,
config: config || existing?.config
};
await this.client.index({
index: KG_METADATA_INDEX,
id: `zone:${name}`,
body: metadata,
refresh: true
});
console.log(`Updated descriptions for zone: ${name}`);
}
/**
* High-level search method that returns clean entity data for user-facing applications
* This method acts as a wrapper around the raw search, with additional processing and AI filtering
*
* @param params Search parameters including query, filters, and AI-related fields
* @returns Clean entity and relation data, filtered by AI if informationNeeded is provided
*/
async userSearch(params: {
query: string;
entityTypes?: string[];
limit?: number;
includeObservations?: boolean;
sortBy?: 'relevance' | 'recent' | 'importance';
zone?: string;
informationNeeded?: string;
reason?: string;
}): Promise<{
entities: Array<{
name: string;
entityType: string;
observations?: string[];
lastRead?: string;
lastWrite?: string;
}>;
relations: Array<{
from: string;
to: string;
type: string;
fromZone: string;
toZone: string;
}>;
}> {
// Set default values
const includeObservations = params.includeObservations ?? false;
const defaultLimit = includeObservations ? 5 : 20;
const zone = params.zone || this.defaultZone;
const informationNeeded = params.informationNeeded;
const reason = params.reason;
// If informationNeeded is provided, increase the limit to get more results
// that will be filtered later by the AI
const searchLimit = informationNeeded ?
(params.limit ? params.limit * 4 : defaultLimit * 4) :
(params.limit || defaultLimit);
// Prepare search parameters for the raw search
const searchParams: ESSearchParams = {
query: params.query,
entityTypes: params.entityTypes,
limit: searchLimit,
sortBy: params.sortBy,
includeObservations,
zone,
informationNeeded,
reason
};
// Perform the raw search
const results = await this.search(searchParams);
// Transform the results to a clean format, removing unnecessary fields
const entities = results.hits.hits
.filter(hit => hit._source.type === 'entity')
.map(hit => {
const entity: {
name: string;
entityType: string;
observations?: string[];
lastRead?: string;
lastWrite?: string;
} = {
name: (hit._source as ESEntity).name,
entityType: (hit._source as ESEntity).entityType,
};
// Only include observations and timestamps if requested
if (includeObservations) {
entity.observations = (hit._source as ESEntity).observations;
entity.lastWrite = (hit._source as ESEntity).lastWrite;
entity.lastRead = (hit._source as ESEntity).lastRead;
}
return entity;
});
// Apply AI filtering if informationNeeded is provided and AI is available
let filteredEntities = entities;
if (informationNeeded && GroqAI.isEnabled && entities.length > 0) {
try {
// Get relevant entity names using AI filtering
const usefulness = await GroqAI.filterSearchResults(entities, informationNeeded, reason);
// If AI filtering returned null (error case), use original entities
if (usefulness === null) {
console.warn('AI filtering returned null, using original results');
filteredEntities = entities.slice(0, params.limit || defaultLimit);
} else {
// Filter entities to only include those with a usefulness score
filteredEntities = entities.filter(entity =>
usefulness[entity.name] !== undefined
);
// Sort entities by their relevance score from highest to lowest
filteredEntities.sort((a, b) => {
const scoreA = usefulness[a.name] || 0;
const scoreB = usefulness[b.name] || 0;
return scoreB - scoreA;
});
const usefulEntities = filteredEntities.filter(entity => usefulness[entity.name] >= 60);
const definatelyNotUsefulEntities = filteredEntities.filter(entity => usefulness[entity.name] < 20);
// for each useful entities, increase the relevanceScore
for (const entity of usefulEntities) {
this.updateEntityRelevanceScore(entity.name, (usefulness[entity.name] + 45) * 0.01, zone);
}
// for each definately not useful entities, decrease the relevanceScore
for (const entity of definatelyNotUsefulEntities) {
this.updateEntityRelevanceScore(entity.name, 0.8 + usefulness[entity.name] * 0.01, zone);
}
// If no entities were found relevant, fall back to the original results
if (filteredEntities.length === 0) {
filteredEntities = entities.slice(0, params.limit || defaultLimit);
} else {
// Limit the filtered results to the requested amount
filteredEntities = filteredEntities.slice(0, params.limit || defaultLimit);
}
}
} catch (error) {
console.error('Error applying AI filtering:', error);
// Fall back to the original results but limit to the requested amount
filteredEntities = entities.slice(0, params.limit || defaultLimit);
}
} else if (entities.length > (params.limit || defaultLimit)) {
// If we're not using AI filtering but retrieved more results due to the doubled limit,
// limit the results to the originally requested amount
filteredEntities = entities.slice(0, params.limit || defaultLimit);
}
// Get relations between these entities
const entityNames = filteredEntities.map(e => e.name);
const { relations } = await this.getRelationsForEntities(entityNames, zone);
// Map relations to a clean format
const formattedRelations = relations.map(r => ({
from: r.from,
to: r.to,
type: r.relationType,
fromZone: r.fromZone,
toZone: r.toZone
}));
return {
entities: filteredEntities,
relations: formattedRelations
};
}
}
```