This is page 3 of 3. Use http://codebase.md/j3k0/mcp-elastic-memory?lines=true&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/kg-client.ts:
--------------------------------------------------------------------------------
```typescript
1 | // WARNING: no console.log in this file, it will break MCP server. Use console.error instead
2 |
3 | import { Client } from '@elastic/elasticsearch';
4 | import {
5 | KG_INDEX_CONFIG,
6 | KG_INDEX_PREFIX,
7 | KG_RELATIONS_INDEX,
8 | KG_METADATA_INDEX,
9 | getIndexName,
10 | ESEntity,
11 | ESRelation,
12 | ESHighlightResponse,
13 | ESSearchParams,
14 | ESSearchResponse
15 | } from './es-types.js';
16 |
17 | // Zone metadata document structure
18 | interface ZoneMetadata {
19 | name: string;
20 | description?: string;
21 | shortDescription?: string;
22 | createdAt: string;
23 | lastModified: string;
24 | config?: Record<string, any>;
25 | }
26 |
27 | // Import the AI service
28 | import GroqAI from './ai-service.js';
29 |
30 | /**
31 | * Knowledge Graph Client
32 | *
33 | * Core library for interacting with the Elasticsearch-backed knowledge graph
34 | */
35 | export class KnowledgeGraphClient {
36 | private client: Client;
37 | private initialized: boolean = false;
38 | private defaultZone: string;
39 |
40 | // Cache of initialized indices to avoid repeated checks
41 | private initializedIndices: Set<string> = new Set();
42 | // Cache of existing zones to avoid repeated database queries when checking zone existence
43 | // This improves performance for operations that check the same zones multiple times
44 | private existingZonesCache: Record<string, boolean> = {};
45 |
46 | /**
47 | * Create a new KnowledgeGraphClient
48 | * @param options Connection options
49 | */
50 | constructor(private options: {
51 | node: string;
52 | auth?: { username: string; password: string };
53 | defaultZone?: string;
54 | }) {
55 | this.client = new Client({
56 | node: options.node,
57 | auth: options.auth,
58 | });
59 | this.defaultZone = options.defaultZone || process.env.KG_DEFAULT_ZONE || 'default';
60 | }
61 |
62 | private getIndexForZone(zone?: string): string {
63 | return getIndexName(zone || this.defaultZone);
64 | }
65 |
66 | /**
67 | * Initialize the knowledge graph (create index if needed)
68 | */
69 | async initialize(zone?: string): Promise<void> {
70 | if (!this.initialized) {
71 | this.client = new Client({
72 | node: this.options.node,
73 | auth: this.options.auth,
74 | });
75 | this.initialized = true;
76 |
77 | // Initialize the metadata index if it doesn't exist yet
78 | const metadataIndexExists = await this.client.indices.exists({
79 | index: KG_METADATA_INDEX
80 | });
81 |
82 | if (!metadataIndexExists) {
83 | await this.client.indices.create({
84 | index: KG_METADATA_INDEX,
85 | mappings: {
86 | properties: {
87 | name: { type: 'keyword' },
88 | description: { type: 'text' },
89 | createdAt: { type: 'date' },
90 | lastModified: { type: 'date' },
91 | config: { type: 'object', enabled: false }
92 | }
93 | }
94 | });
95 | console.error(`Created metadata index: ${KG_METADATA_INDEX}`);
96 |
97 | // Add default zone metadata
98 | await this.saveZoneMetadata('default', 'Default knowledge zone');
99 | }
100 |
101 | // Initialize the relations index if it doesn't exist yet
102 | const relationsIndexExists = await this.client.indices.exists({
103 | index: KG_RELATIONS_INDEX
104 | });
105 |
106 | if (!relationsIndexExists) {
107 | await this.client.indices.create({
108 | index: KG_RELATIONS_INDEX,
109 | ...KG_INDEX_CONFIG
110 | });
111 | console.error(`Created relations index: ${KG_RELATIONS_INDEX}`);
112 | }
113 | }
114 |
115 | // Continue with zone-specific index initialization
116 | const indexName = this.getIndexForZone(zone);
117 |
118 | // If we've already initialized this index, skip
119 | if (this.initializedIndices.has(indexName)) {
120 | return;
121 | }
122 |
123 | const indexExists = await this.client.indices.exists({ index: indexName });
124 |
125 | if (!indexExists) {
126 | await this.client.indices.create({
127 | index: indexName,
128 | ...KG_INDEX_CONFIG
129 | });
130 | console.error(`Created index: ${indexName}`);
131 | }
132 |
133 | this.initializedIndices.add(indexName);
134 | }
135 |
136 | /**
137 | * Create or update an entity
138 | * @param entity Entity to create or update
139 | * @param zone Optional memory zone name, uses defaultZone if not specified
140 | * @param options Optional configuration options
141 | * @param options.validateZones Whether to validate that zones exist before creating entities (default: true)
142 | */
143 | async saveEntity(
144 | entity: Omit<ESEntity, 'type' | 'readCount' | 'lastRead' | 'lastWrite' | 'zone'>,
145 | zone?: string,
146 | options?: {
147 | validateZones?: boolean;
148 | }
149 | ): Promise<ESEntity> {
150 | // Validate entity name is not empty
151 | if (!entity.name || entity.name.trim() === '') {
152 | throw new Error('Entity name cannot be empty');
153 | }
154 |
155 | const actualZone = zone || this.defaultZone;
156 | await this.initialize(actualZone);
157 |
158 | // Default to true for zone validation
159 | const validateZones = options?.validateZones ?? true;
160 |
161 | // Validate that zone exists if required
162 | if (validateZones && actualZone !== this.defaultZone) {
163 | const zoneExists = await this.zoneExists(actualZone);
164 | if (!zoneExists) {
165 | throw new Error(`Cannot create entity: Zone '${actualZone}' does not exist. Create the zone first.`);
166 | }
167 | }
168 |
169 | const now = new Date().toISOString();
170 | const existingEntity = await this.getEntity(entity.name, actualZone);
171 |
172 | const newEntity: ESEntity = {
173 | type: 'entity',
174 | name: entity.name,
175 | entityType: entity.entityType,
176 | observations: entity.observations || [],
177 | // If entity exists, preserve its readCount and lastRead, but update lastWrite
178 | readCount: existingEntity?.readCount ?? 0,
179 | lastRead: existingEntity?.lastRead ?? now,
180 | lastWrite: now,
181 | relevanceScore: entity.relevanceScore ?? (existingEntity?.relevanceScore ?? 1.0),
182 | zone: actualZone
183 | };
184 |
185 | const indexName = this.getIndexForZone(actualZone);
186 | await this.client.index({
187 | index: indexName,
188 | id: `entity:${entity.name}`,
189 | document: newEntity,
190 | refresh: true // Make sure it's immediately available for search
191 | });
192 |
193 | return newEntity;
194 | }
195 |
196 | /**
197 | * Get an entity by name without updating lastRead timestamp
198 | * @param name Entity name
199 | * @param zone Optional memory zone name, uses defaultZone if not specified
200 | */
201 | async getEntityWithoutUpdatingLastRead(name: string, zone?: string): Promise<ESEntity | null> {
202 | const actualZone = zone || this.defaultZone;
203 | await this.initialize(actualZone);
204 |
205 | try {
206 | const indexName = this.getIndexForZone(actualZone);
207 | const id = `entity:${name}`;
208 |
209 | // Try direct get by ID first
210 | try {
211 | const response = await this.client.get({
212 | index: indexName,
213 | id: id
214 | });
215 |
216 | if (response && response._source) {
217 | return response._source as ESEntity;
218 | }
219 | } catch (error) {
220 | // If not found by ID, try search
221 | if (error.statusCode === 404) {
222 | // Fall through to search
223 | } else {
224 | throw error;
225 | }
226 | }
227 |
228 | // If direct get fails, use search with explicit zone filter
229 | const response = await this.client.search({
230 | index: indexName,
231 | body: {
232 | query: {
233 | bool: {
234 | must: [
235 | // Use term query for exact name matching
236 | { term: { name: name } },
237 | { term: { type: 'entity' } },
238 | { term: { zone: actualZone } }
239 | ]
240 | }
241 | },
242 | size: 1
243 | }
244 | });
245 |
246 | const typedResponse = response as unknown as ESSearchResponse<ESEntity>;
247 |
248 | if (typedResponse.hits.total.value === 0) {
249 | return null;
250 | }
251 |
252 | return typedResponse.hits.hits[0]._source;
253 | } catch (error) {
254 | if (error.statusCode === 404) {
255 | return null;
256 | }
257 | throw error;
258 | }
259 | }
260 |
261 | /**
262 | * Get an entity by name and update lastRead timestamp and readCount
263 | * @param name Entity name
264 | * @param zone Optional memory zone name, uses defaultZone if not specified
265 | */
266 | async getEntity(name: string, zone?: string): Promise<ESEntity | null> {
267 | const actualZone = zone || this.defaultZone;
268 | const entity = await this.getEntityWithoutUpdatingLastRead(name, actualZone);
269 |
270 | if (!entity) {
271 | return null;
272 | }
273 |
274 | // Update lastRead and readCount in memory (skip the database update if it fails)
275 | const now = new Date().toISOString();
276 | const updatedEntity = {
277 | ...entity,
278 | lastRead: now,
279 | readCount: entity.readCount + 1
280 | };
281 |
282 | try {
283 | // Try to update in the database, but don't fail if it doesn't work
284 | const indexName = this.getIndexForZone(actualZone);
285 |
286 | // Search for the entity by name and zone to get the _id
287 | const searchResponse = await this.client.search({
288 | index: indexName,
289 | body: {
290 | query: {
291 | bool: {
292 | must: [
293 | { term: { name: name } },
294 | { term: { type: 'entity' } },
295 | { term: { zone: actualZone } }
296 | ]
297 | }
298 | },
299 | size: 1
300 | }
301 | });
302 |
303 | const typedResponse = searchResponse as unknown as ESSearchResponse<ESEntity>;
304 |
305 | if (typedResponse.hits.total.value > 0) {
306 | const docId = typedResponse.hits.hits[0]._id;
307 |
308 | await this.client.update({
309 | index: indexName,
310 | id: docId,
311 | doc: {
312 | lastRead: now,
313 | readCount: entity.readCount + 1
314 | },
315 | refresh: true
316 | });
317 | } else {
318 | // This indicates the entity exists in memory but not in the index
319 | // Instead of showing an error message, silently handle this condition
320 | // The entity is still returned to the caller with updated timestamps
321 | }
322 | } catch (error) {
323 | // If update fails, just log it and return the entity with updated timestamps
324 | console.error(`Warning: Failed to update lastRead timestamp for entity ${name}: ${(error as Error).message}`);
325 | }
326 |
327 | return updatedEntity;
328 | }
329 |
330 | /**
331 | * Delete an entity by name
332 | * @param name Entity name
333 | * @param zone Optional memory zone name, uses defaultZone if not specified
334 | * @param options Optional configuration options
335 | * @param options.cascadeRelations Whether to delete relations involving this entity (default: true)
336 | */
337 | async deleteEntity(
338 | name: string,
339 | zone?: string,
340 | options?: {
341 | cascadeRelations?: boolean;
342 | }
343 | ): Promise<boolean> {
344 | // Validate entity name is not empty
345 | if (!name || name.trim() === '') {
346 | throw new Error('Entity name cannot be empty');
347 | }
348 |
349 | const actualZone = zone || this.defaultZone;
350 | await this.initialize(actualZone);
351 |
352 | // Default to true for cascading relation deletion
353 | const cascadeRelations = options?.cascadeRelations !== false;
354 |
355 | try {
356 | // First, check if the entity exists
357 | const entity = await this.getEntityWithoutUpdatingLastRead(name, actualZone);
358 | if (!entity) {
359 | return false;
360 | }
361 |
362 | const indexName = this.getIndexForZone(actualZone);
363 |
364 | // Delete relations involving this entity if cascading is enabled
365 | if (cascadeRelations) {
366 | console.error(`Cascading relations for entity ${name} in zone ${actualZone}`);
367 |
368 | // First, delete relations within the same zone
369 | await this.client.deleteByQuery({
370 | index: indexName,
371 | body: {
372 | query: {
373 | bool: {
374 | must: [
375 | { term: { type: 'relation' } },
376 | {
377 | bool: {
378 | should: [
379 | { term: { from: name } },
380 | { term: { to: name } }
381 | ]
382 | }
383 | }
384 | ]
385 | }
386 | }
387 | },
388 | refresh: true
389 | });
390 |
391 | // Then, delete cross-zone relations where this entity is involved
392 | // These are stored in the KG_RELATIONS_INDEX
393 | await this.client.deleteByQuery({
394 | index: KG_RELATIONS_INDEX,
395 | body: {
396 | query: {
397 | bool: {
398 | must: [
399 | {
400 | bool: {
401 | should: [
402 | {
403 | bool: {
404 | must: [
405 | { term: { from: name } },
406 | { term: { fromZone: actualZone } }
407 | ]
408 | }
409 | },
410 | {
411 | bool: {
412 | must: [
413 | { term: { to: name } },
414 | { term: { toZone: actualZone } }
415 | ]
416 | }
417 | }
418 | ]
419 | }
420 | }
421 | ]
422 | }
423 | }
424 | },
425 | refresh: true
426 | });
427 | } else {
428 | console.error(`Skipping relation cascade for entity ${name} in zone ${actualZone}`);
429 | }
430 |
431 | // Delete the entity
432 | await this.client.delete({
433 | index: indexName,
434 | id: `entity:${name}`,
435 | refresh: true
436 | });
437 |
438 | return true;
439 | } catch (error) {
440 | if (error.statusCode === 404) {
441 | return false;
442 | }
443 | throw error;
444 | }
445 | }
446 |
447 | /**
448 | * Check if a memory zone exists
449 | * @param zone Zone name to check
450 | * @returns Promise<boolean> True if the zone exists, false otherwise
451 | *
452 | * This method uses a caching strategy to avoid repeated database queries:
453 | * 1. Default zone always exists and is automatically added to the cache
454 | * 2. If the requested zone is in the cache, returns the cached result
455 | * 3. Otherwise, checks the zone metadata and updates the cache
456 | * 4. The cache is maintained by addMemoryZone, deleteMemoryZone, and listMemoryZones
457 | */
458 | async zoneExists(zone: string): Promise<boolean> {
459 | if (!zone || zone === this.defaultZone) {
460 | // Default zone always exists
461 | this.existingZonesCache[this.defaultZone] = true;
462 | return true;
463 | }
464 |
465 | // Check cache first
466 | if (this.existingZonesCache[zone] !== undefined) {
467 | return this.existingZonesCache[zone];
468 | }
469 |
470 | await this.initialize();
471 |
472 | // Check metadata for zone
473 | const metadata = await this.getZoneMetadata(zone);
474 | if (metadata) {
475 | // Cache the result
476 | this.existingZonesCache[zone] = true;
477 | return true;
478 | }
479 |
480 | // As a fallback, check if the index exists
481 | const indexName = this.getIndexForZone(zone);
482 | const indexExists = await this.client.indices.exists({ index: indexName });
483 |
484 | // Cache the result
485 | this.existingZonesCache[zone] = indexExists;
486 |
487 | return indexExists;
488 | }
489 |
490 | /**
491 | * Create or update a relation between entities
492 | * @param relation Relation to create or update
493 | * @param fromZone Optional zone for the source entity, uses defaultZone if not specified
494 | * @param toZone Optional zone for the target entity, uses defaultZone if not specified
495 | * @param options Optional configuration options
496 | * @param options.autoCreateMissingEntities Whether to automatically create missing entities (default: true)
497 | * @param options.validateZones Whether to validate that zones exist before creating entities or relations (default: true)
498 | */
499 | async saveRelation(
500 | relation: Omit<ESRelation, 'type' | 'fromZone' | 'toZone'>,
501 | fromZone?: string,
502 | toZone?: string,
503 | options?: {
504 | autoCreateMissingEntities?: boolean;
505 | validateZones?: boolean;
506 | }
507 | ): Promise<ESRelation> {
508 | await this.initialize();
509 |
510 | // Default to true for backwards compatibility
511 | const autoCreateMissingEntities = options?.autoCreateMissingEntities ?? true;
512 | // Default to true for zone validation
513 | const validateZones = options?.validateZones ?? true;
514 |
515 | const actualFromZone = fromZone || this.defaultZone;
516 | const actualToZone = toZone || this.defaultZone;
517 |
518 | // Validate that zones exist if required
519 | if (validateZones) {
520 | // Check fromZone
521 | const fromZoneExists = await this.zoneExists(actualFromZone);
522 | if (!fromZoneExists) {
523 | throw new Error(`Cannot create relation: Source zone '${actualFromZone}' does not exist. Create the zone first.`);
524 | }
525 |
526 | // Check toZone
527 | const toZoneExists = await this.zoneExists(actualToZone);
528 | if (!toZoneExists) {
529 | throw new Error(`Cannot create relation: Target zone '${actualToZone}' does not exist. Create the zone first.`);
530 | }
531 | }
532 |
533 | // Check if both entities exist
534 | const fromEntity = await this.getEntityWithoutUpdatingLastRead(relation.from, actualFromZone);
535 | const toEntity = await this.getEntityWithoutUpdatingLastRead(relation.to, actualToZone);
536 |
537 | // If either entity doesn't exist
538 | if (!fromEntity || !toEntity) {
539 | // If auto-creation is disabled, throw an error
540 | if (!autoCreateMissingEntities) {
541 | const missingEntities = [];
542 | if (!fromEntity) {
543 | missingEntities.push(`'${relation.from}' in zone '${actualFromZone}'`);
544 | }
545 | if (!toEntity) {
546 | missingEntities.push(`'${relation.to}' in zone '${actualToZone}'`);
547 | }
548 |
549 | throw new Error(`Cannot create relation: Missing entities ${missingEntities.join(' and ')}`);
550 | }
551 |
552 | // Otherwise, auto-create the missing entities
553 | if (!fromEntity) {
554 | console.error(`Auto-creating missing entity: ${relation.from} in zone ${actualFromZone}`);
555 | const newEntity = {
556 | type: 'entity',
557 | name: relation.from,
558 | entityType: 'unknown',
559 | observations: [],
560 | readCount: 0,
561 | lastRead: new Date().toISOString(),
562 | lastWrite: new Date().toISOString(),
563 | relevanceScore: 1.0,
564 | zone: actualFromZone
565 | };
566 |
567 | // We've already validated the zone, so we can skip validation here
568 | await this.saveEntity({
569 | name: relation.from,
570 | entityType: 'unknown',
571 | observations: [],
572 | relevanceScore: 1.0
573 | }, actualFromZone, { validateZones: false });
574 | }
575 |
576 | if (!toEntity) {
577 | console.error(`Auto-creating missing entity: ${relation.to} in zone ${actualToZone}`);
578 | const newEntity = {
579 | type: 'entity',
580 | name: relation.to,
581 | entityType: 'unknown',
582 | observations: [],
583 | readCount: 0,
584 | lastRead: new Date().toISOString(),
585 | lastWrite: new Date().toISOString(),
586 | relevanceScore: 1.0,
587 | zone: actualToZone
588 | };
589 |
590 | // We've already validated the zone, so we can skip validation here
591 | await this.saveEntity({
592 | name: relation.to,
593 | entityType: 'unknown',
594 | observations: [],
595 | relevanceScore: 1.0
596 | }, actualToZone, { validateZones: false });
597 | }
598 | }
599 |
600 | // Create the relation
601 | const newRelation: ESRelation = {
602 | type: 'relation',
603 | from: relation.from,
604 | fromZone: actualFromZone,
605 | to: relation.to,
606 | toZone: actualToZone,
607 | relationType: relation.relationType
608 | };
609 |
610 | const id = `relation:${actualFromZone}:${relation.from}:${relation.relationType}:${actualToZone}:${relation.to}`;
611 |
612 | await this.client.index({
613 | index: KG_RELATIONS_INDEX,
614 | id,
615 | document: newRelation,
616 | refresh: true
617 | });
618 |
619 | return newRelation;
620 | }
621 |
622 | /**
623 | * Delete a relation between entities
624 | * @param from Source entity name
625 | * @param to Target entity name
626 | * @param relationType Relation type
627 | * @param fromZone Optional zone for the source entity, uses defaultZone if not specified
628 | * @param toZone Optional zone for the target entity, uses defaultZone if not specified
629 | */
630 | async deleteRelation(
631 | from: string,
632 | to: string,
633 | relationType: string,
634 | fromZone?: string,
635 | toZone?: string
636 | ): Promise<boolean> {
637 | await this.initialize();
638 |
639 | const actualFromZone = fromZone || this.defaultZone;
640 | const actualToZone = toZone || this.defaultZone;
641 |
642 | try {
643 | const relationId = `relation:${actualFromZone}:${from}:${relationType}:${actualToZone}:${to}`;
644 |
645 | await this.client.delete({
646 | index: KG_RELATIONS_INDEX,
647 | id: relationId,
648 | refresh: true
649 | });
650 |
651 | return true;
652 | } catch (error) {
653 | if (error.statusCode === 404) {
654 | return false;
655 | }
656 | throw error;
657 | }
658 | }
659 |
660 | /**
661 | * Search for entities and relations in the knowledge graph
662 | * @param params Search parameters
663 | */
664 | async search(params: ESSearchParams & { zone?: string }): Promise<ESHighlightResponse<ESEntity | ESRelation>> {
665 | const actualZone = params.zone || this.defaultZone;
666 | await this.initialize(actualZone);
667 |
668 | const indexName = this.getIndexForZone(actualZone);
669 |
670 | // Special handling for wildcard query
671 | if (params.query === '*') {
672 | // console.error(`Performing wildcard search in zone: ${actualZone}, index: ${indexName}`);
673 |
674 | try {
675 | // Use match_all query for wildcard
676 | const response = await this.client.search({
677 | index: indexName,
678 | body: {
679 | query: {
680 | match_all: {}
681 | },
682 | sort: [{ lastRead: { order: 'desc' } }],
683 | size: params.limit || 10,
684 | from: params.offset || 0
685 | }
686 | });
687 |
688 | // console.error(`Wildcard search results: ${JSON.stringify(response.hits)}`);
689 |
690 | return response as unknown as ESHighlightResponse<ESEntity | ESRelation>;
691 | } catch (error) {
692 | console.error(`Error in wildcard search: ${error.message}`);
693 | throw error;
694 | }
695 | }
696 |
697 | // Special handling for exact entity name search
698 | if (params.query && !params.query.includes(' ')) {
699 | // console.error(`Performing exact entity name search for "${params.query}" in zone: ${actualZone}, index: ${indexName}`);
700 |
701 | try {
702 | // Use match query for exact entity name
703 | const response = await this.client.search({
704 | index: indexName,
705 | body: {
706 | query: {
707 | match: {
708 | name: params.query
709 | }
710 | },
711 | sort: [{ lastRead: { order: 'desc' } }],
712 | size: params.limit || 10,
713 | from: params.offset || 0
714 | }
715 | });
716 |
717 | // console.error(`Exact entity name search results: ${JSON.stringify(response.hits)}`);
718 |
719 | return response as unknown as ESHighlightResponse<ESEntity | ESRelation>;
720 | } catch (error) {
721 | console.error(`Error in exact entity name search: ${error.message}`);
722 | throw error;
723 | }
724 | }
725 |
726 | // Build search query for non-wildcard searches
727 | const query: any = {
728 | bool: {
729 | must: []
730 | }
731 | };
732 |
733 | // Process the query to handle boolean operators or fuzzy search notation
734 | if (params.query.includes(' AND ') ||
735 | params.query.includes(' OR ') ||
736 | params.query.includes(' NOT ') ||
737 | params.query.includes('~')) {
738 | // Use Elasticsearch's query_string query for boolean operators and fuzzy search
739 | query.bool.must.push({
740 | query_string: {
741 | query: params.query,
742 | fields: ['name^3', 'entityType^2', 'observations', 'relationType^2'],
743 | default_operator: 'AND',
744 | // Enable fuzzy matching by default
745 | fuzziness: 'AUTO',
746 | // Allow fuzzy matching on all terms unless explicitly disabled
747 | fuzzy_max_expansions: 50,
748 | // Support phrase slop for proximity searches
749 | phrase_slop: 2
750 | }
751 | });
752 |
753 | console.error(`Using query_string for advanced query: ${params.query}`);
754 | } else {
755 | // Use multi_match for simple queries without boolean operators
756 | query.bool.must.push({
757 | multi_match: {
758 | query: params.query,
759 | fields: ['name^3', 'entityType^2', 'observations', 'relationType^2'],
760 | // Enable fuzzy matching by default
761 | fuzziness: 'AUTO'
762 | }
763 | });
764 |
765 | console.error(`Using multi_match for simple query: ${params.query}`);
766 | }
767 |
768 | // Add entityType filter if specified
769 | if (params.entityTypes && params.entityTypes.length > 0) {
770 | // Use a more flexible filter with "should" (equivalent to OR) to match any of the entity types
771 | const entityTypeFilters = params.entityTypes.map(type => ({
772 | match: {
773 | entityType: {
774 | query: type,
775 | operator: "and"
776 | }
777 | }
778 | }));
779 |
780 | query.bool.must.push({
781 | bool: {
782 | should: entityTypeFilters,
783 | minimum_should_match: 1
784 | }
785 | });
786 |
787 | console.error(`Applied entity type filters: ${JSON.stringify(entityTypeFilters)}`);
788 | }
789 |
790 | // Log validation to ensure zone filter is being applied
791 | console.error(`Searching in zone: ${actualZone}, index: ${indexName}, query: ${JSON.stringify(query)}`);
792 |
793 | // Set up sort order
794 | let sort: any[] = [];
795 | if (params.sortBy === 'recent') {
796 | sort = [{ lastRead: { order: 'desc' } }];
797 | } else if (params.sortBy === 'importance') {
798 | sort = [
799 | // Always sort by relevanceScore in descending order (highest first)
800 | { relevanceScore: { order: 'desc' } }
801 | ];
802 | } else {
803 | // Default is by relevance (using _score)
804 | sort = [{ _score: { order: 'desc' } }];
805 | }
806 |
807 | try {
808 | // Execute search
809 | const response = await this.client.search({
810 | index: indexName,
811 | body: {
812 | query,
813 | sort,
814 | highlight: {
815 | fields: {
816 | name: {},
817 | observations: {},
818 | entityType: {}
819 | }
820 | },
821 | size: params.limit || 10,
822 | from: params.offset || 0
823 | }
824 | });
825 |
826 | console.error(`Search results: ${JSON.stringify(response.hits)}`);
827 |
828 | return response as unknown as ESHighlightResponse<ESEntity | ESRelation>;
829 | } catch (error) {
830 | console.error(`Error in search: ${error.message}`);
831 | throw error;
832 | }
833 | }
834 |
835 | /**
836 | * Get all entities related to a given entity, up to a certain depth
837 | * @param name Entity name
838 | * @param maxDepth Maximum traversal depth
839 | * @param zone Optional memory zone name, uses defaultZone if not specified
840 | */
841 | async getRelatedEntities(
842 | name: string,
843 | maxDepth: number = 1,
844 | zone?: string
845 | ): Promise<{
846 | entities: ESEntity[],
847 | relations: ESRelation[]
848 | }> {
849 | const actualZone = zone || this.defaultZone;
850 | await this.initialize(actualZone);
851 |
852 | // Start with the root entity
853 | const rootEntity = await this.getEntity(name, actualZone);
854 | if (!rootEntity) {
855 | return { entities: [], relations: [] };
856 | }
857 |
858 | // Keep track of entities and relations we've found
859 | const entitiesMap = new Map<string, ESEntity>();
860 | entitiesMap.set(`${actualZone}:${rootEntity.name}`, rootEntity);
861 |
862 | const relationsMap = new Map<string, ESRelation>();
863 | const visitedEntities = new Set<string>();
864 |
865 | // Queue of entities to process, with their depth
866 | const queue: Array<{ entity: ESEntity, zone: string, depth: number }> = [
867 | { entity: rootEntity, zone: actualZone, depth: 0 }
868 | ];
869 |
870 | while (queue.length > 0) {
871 | const { entity, zone: entityZone, depth } = queue.shift()!;
872 | const entityKey = `${entityZone}:${entity.name}`;
873 |
874 | // Skip if we've already processed this entity or if we've reached max depth
875 | if (visitedEntities.has(entityKey) || depth >= maxDepth) {
876 | continue;
877 | }
878 |
879 | visitedEntities.add(entityKey);
880 |
881 | // Find all relations involving this entity
882 | const fromResponse = await this.client.search({
883 | index: KG_RELATIONS_INDEX,
884 | body: {
885 | query: {
886 | bool: {
887 | must: [
888 | { term: { from: entity.name } },
889 | { term: { "fromZone.keyword": entityZone } }
890 | ]
891 | }
892 | },
893 | size: 1000
894 | }
895 | });
896 |
897 | const toResponse = await this.client.search({
898 | index: KG_RELATIONS_INDEX,
899 | body: {
900 | query: {
901 | bool: {
902 | must: [
903 | { term: { to: entity.name } },
904 | { term: { "toZone.keyword": entityZone } }
905 | ]
906 | }
907 | },
908 | size: 1000
909 | }
910 | });
911 |
912 | // Process relations where this entity is the source
913 | const fromHits = (fromResponse as unknown as ESSearchResponse<ESRelation>).hits.hits;
914 | for (const hit of fromHits) {
915 | const relation = hit._source;
916 | const relationKey = `${relation.fromZone}:${relation.from}|${relation.relationType}|${relation.toZone}:${relation.to}`;
917 |
918 | // Skip if we've already processed this relation
919 | if (relationsMap.has(relationKey)) {
920 | continue;
921 | }
922 |
923 | relationsMap.set(relationKey, relation);
924 |
925 | // Process the target entity
926 | const otherEntityKey = `${relation.toZone}:${relation.to}`;
927 | if (!entitiesMap.has(otherEntityKey)) {
928 | // Fetch the other entity
929 | const otherEntity = await this.getEntity(relation.to, relation.toZone);
930 |
931 | if (otherEntity) {
932 | entitiesMap.set(otherEntityKey, otherEntity);
933 |
934 | // Add the other entity to the queue if we haven't reached max depth
935 | if (depth < maxDepth - 1) {
936 | queue.push({ entity: otherEntity, zone: relation.toZone, depth: depth + 1 });
937 | }
938 | }
939 | }
940 | }
941 |
942 | // Process relations where this entity is the target
943 | const toHits = (toResponse as unknown as ESSearchResponse<ESRelation>).hits.hits;
944 | for (const hit of toHits) {
945 | const relation = hit._source;
946 | const relationKey = `${relation.fromZone}:${relation.from}|${relation.relationType}|${relation.toZone}:${relation.to}`;
947 |
948 | // Skip if we've already processed this relation
949 | if (relationsMap.has(relationKey)) {
950 | continue;
951 | }
952 |
953 | relationsMap.set(relationKey, relation);
954 |
955 | // Process the source entity
956 | const otherEntityKey = `${relation.fromZone}:${relation.from}`;
957 | if (!entitiesMap.has(otherEntityKey)) {
958 | // Fetch the other entity
959 | const otherEntity = await this.getEntity(relation.from, relation.fromZone);
960 |
961 | if (otherEntity) {
962 | entitiesMap.set(otherEntityKey, otherEntity);
963 |
964 | // Add the other entity to the queue if we haven't reached max depth
965 | if (depth < maxDepth - 1) {
966 | queue.push({ entity: otherEntity, zone: relation.fromZone, depth: depth + 1 });
967 | }
968 | }
969 | }
970 | }
971 | }
972 |
973 | return {
974 | entities: Array.from(entitiesMap.values()),
975 | relations: Array.from(relationsMap.values())
976 | };
977 | }
978 |
979 | /**
980 | * Import data into the knowledge graph
981 | * @param data Array of entities and relations to import
982 | * @param zone Optional memory zone for entities, uses defaultZone if not specified
983 | * @param options Optional configuration options
984 | * @param options.validateZones Whether to validate that zones exist before importing (default: true)
985 | */
986 | async importData(
987 | data: Array<ESEntity | ESRelation>,
988 | zone?: string,
989 | options?: {
990 | validateZones?: boolean;
991 | }
992 | ): Promise<{
993 | entitiesAdded: number;
994 | relationsAdded: number;
995 | invalidRelations?: Array<{relation: ESRelation, reason: string}>;
996 | }> {
997 | const actualZone = zone || this.defaultZone;
998 | await this.initialize(actualZone);
999 |
1000 | // Default to true for zone validation
1001 | const validateZones = options?.validateZones ?? true;
1002 |
1003 | // Validate that zone exists if required
1004 | if (validateZones && actualZone !== this.defaultZone) {
1005 | const zoneExists = await this.zoneExists(actualZone);
1006 | if (!zoneExists) {
1007 | throw new Error(`Cannot import data: Zone '${actualZone}' does not exist. Create the zone first.`);
1008 | }
1009 | }
1010 |
1011 | let entitiesAdded = 0;
1012 | let relationsAdded = 0;
1013 | const invalidRelations: Array<{relation: ESRelation, reason: string}> = [];
1014 |
1015 | // Process entities first, since relations depend on them
1016 | const entities = data.filter(item => item.type === 'entity') as ESEntity[];
1017 | const entityOperations: any[] = [];
1018 |
1019 | for (const entity of entities) {
1020 | // Add zone information if not already present
1021 | const entityWithZone = {
1022 | ...entity,
1023 | zone: entity.zone || actualZone
1024 | };
1025 |
1026 | const id = `entity:${entity.name}`;
1027 | entityOperations.push({ index: { _index: this.getIndexForZone(actualZone), _id: id } });
1028 | entityOperations.push(entityWithZone);
1029 | entitiesAdded++;
1030 | }
1031 |
1032 | if (entityOperations.length > 0) {
1033 | await this.client.bulk({
1034 | operations: entityOperations,
1035 | refresh: true
1036 | });
1037 | }
1038 |
1039 | // Now process relations
1040 | const relations = data.filter(item => item.type === 'relation') as ESRelation[];
1041 | const relationOperations: any[] = [];
1042 |
1043 | for (const relation of relations) {
1044 | // For relations with explicit zones
1045 | if (relation.fromZone !== undefined && relation.toZone !== undefined) {
1046 | // If zone validation is enabled, check that both zones exist
1047 | if (validateZones) {
1048 | // Check fromZone if it's not the default zone
1049 | if (relation.fromZone !== this.defaultZone) {
1050 | const fromZoneExists = await this.zoneExists(relation.fromZone);
1051 | if (!fromZoneExists) {
1052 | invalidRelations.push({
1053 | relation,
1054 | reason: `Source zone '${relation.fromZone}' does not exist. Create the zone first.`
1055 | });
1056 | continue;
1057 | }
1058 | }
1059 |
1060 | // Check toZone if it's not the default zone
1061 | if (relation.toZone !== this.defaultZone) {
1062 | const toZoneExists = await this.zoneExists(relation.toZone);
1063 | if (!toZoneExists) {
1064 | invalidRelations.push({
1065 | relation,
1066 | reason: `Target zone '${relation.toZone}' does not exist. Create the zone first.`
1067 | });
1068 | continue;
1069 | }
1070 | }
1071 | }
1072 |
1073 | // Verify that both entities exist
1074 | const fromEntity = await this.getEntityWithoutUpdatingLastRead(relation.from, relation.fromZone);
1075 | const toEntity = await this.getEntityWithoutUpdatingLastRead(relation.to, relation.toZone);
1076 |
1077 | if (!fromEntity) {
1078 | invalidRelations.push({
1079 | relation,
1080 | reason: `Source entity '${relation.from}' in zone '${relation.fromZone}' does not exist`
1081 | });
1082 | continue;
1083 | }
1084 |
1085 | if (!toEntity) {
1086 | invalidRelations.push({
1087 | relation,
1088 | reason: `Target entity '${relation.to}' in zone '${relation.toZone}' does not exist`
1089 | });
1090 | continue;
1091 | }
1092 |
1093 | const id = `relation:${relation.fromZone}:${relation.from}:${relation.relationType}:${relation.toZone}:${relation.to}`;
1094 | relationOperations.push({ index: { _index: KG_RELATIONS_INDEX, _id: id } });
1095 | relationOperations.push(relation);
1096 | relationsAdded++;
1097 | } else {
1098 | // Old format - needs to be converted
1099 | // For backward compatibility, assume both entities are in the specified zone
1100 | const fromZone = actualZone;
1101 | const toZone = actualZone;
1102 |
1103 | // Verify that both entities exist
1104 | const fromEntity = await this.getEntityWithoutUpdatingLastRead(relation.from, fromZone);
1105 | const toEntity = await this.getEntityWithoutUpdatingLastRead(relation.to, toZone);
1106 |
1107 | if (!fromEntity) {
1108 | invalidRelations.push({
1109 | relation,
1110 | reason: `Source entity '${relation.from}' in zone '${fromZone}' does not exist`
1111 | });
1112 | continue;
1113 | }
1114 |
1115 | if (!toEntity) {
1116 | invalidRelations.push({
1117 | relation,
1118 | reason: `Target entity '${relation.to}' in zone '${toZone}' does not exist`
1119 | });
1120 | continue;
1121 | }
1122 |
1123 | // Convert to new format
1124 | const newRelation: ESRelation = {
1125 | type: 'relation',
1126 | from: relation.from,
1127 | fromZone,
1128 | to: relation.to,
1129 | toZone,
1130 | relationType: relation.relationType
1131 | };
1132 |
1133 | const id = `relation:${fromZone}:${relation.from}:${relation.relationType}:${toZone}:${relation.to}`;
1134 | relationOperations.push({ index: { _index: KG_RELATIONS_INDEX, _id: id } });
1135 | relationOperations.push(newRelation);
1136 | relationsAdded++;
1137 | }
1138 | }
1139 |
1140 | if (relationOperations.length > 0) {
1141 | await this.client.bulk({
1142 | operations: relationOperations,
1143 | refresh: true
1144 | });
1145 | }
1146 |
1147 | return {
1148 | entitiesAdded,
1149 | relationsAdded,
1150 | invalidRelations: invalidRelations.length ? invalidRelations : undefined
1151 | };
1152 | }
1153 |
1154 | /**
1155 | * Import data into the knowledge graph, recreating zones as needed
1156 | * @param data Export data containing entities, relations, and zone metadata
1157 | */
1158 | async importAllData(data: {
1159 | entities: ESEntity[],
1160 | relations: ESRelation[],
1161 | zones: ZoneMetadata[]
1162 | }): Promise<{
1163 | zonesAdded: number;
1164 | entitiesAdded: number;
1165 | relationsAdded: number;
1166 | }> {
1167 | await this.initialize();
1168 |
1169 | let zonesAdded = 0;
1170 | let entitiesAdded = 0;
1171 | let relationsAdded = 0;
1172 |
1173 | // First create all zones
1174 | for (const zone of data.zones) {
1175 | if (zone.name !== 'default') {
1176 | await this.addMemoryZone(zone.name, zone.description, zone.config);
1177 | // addMemoryZone already updates the cache
1178 | zonesAdded++;
1179 | } else {
1180 | // Make sure default zone is in the cache
1181 | this.existingZonesCache['default'] = true;
1182 | }
1183 | }
1184 |
1185 | // Import entities by zone
1186 | const entitiesByZone: Record<string, ESEntity[]> = {};
1187 | for (const entity of data.entities) {
1188 | const zone = entity.zone || 'default';
1189 | if (!entitiesByZone[zone]) {
1190 | entitiesByZone[zone] = [];
1191 | }
1192 | entitiesByZone[zone].push(entity);
1193 | }
1194 |
1195 | // Import entities for each zone
1196 | for (const [zone, entities] of Object.entries(entitiesByZone)) {
1197 | const result = await this.importData(entities, zone);
1198 | entitiesAdded += result.entitiesAdded;
1199 | }
1200 |
1201 | // Import all relations
1202 | if (data.relations.length > 0) {
1203 | const result = await this.importData(data.relations);
1204 | relationsAdded = result.relationsAdded;
1205 | }
1206 |
1207 | return {
1208 | zonesAdded,
1209 | entitiesAdded,
1210 | relationsAdded
1211 | };
1212 | }
1213 |
1214 | /**
1215 | * Export all data from a knowledge graph
1216 | * @param zone Optional memory zone for entities, uses defaultZone if not specified
1217 | */
1218 | async exportData(zone?: string): Promise<Array<ESEntity | ESRelation>> {
1219 | const actualZone = zone || this.defaultZone;
1220 | await this.initialize(actualZone);
1221 |
1222 | // Fetch all entities from the specified zone
1223 | const indexName = this.getIndexForZone(actualZone);
1224 | const entityResponse = await this.client.search({
1225 | index: indexName,
1226 | body: {
1227 | query: { term: { type: 'entity' } },
1228 | size: 10000
1229 | }
1230 | });
1231 |
1232 | const entities = entityResponse.hits.hits.map(hit => hit._source) as ESEntity[];
1233 |
1234 | // Fetch all relations involving entities in this zone
1235 | const relationResponse = await this.client.search({
1236 | index: KG_RELATIONS_INDEX,
1237 | body: {
1238 | query: {
1239 | bool: {
1240 | should: [
1241 | { term: { fromZone: actualZone } },
1242 | { term: { toZone: actualZone } }
1243 | ],
1244 | minimum_should_match: 1
1245 | }
1246 | },
1247 | size: 10000
1248 | }
1249 | });
1250 |
1251 | const relations = relationResponse.hits.hits.map(hit => hit._source) as ESRelation[];
1252 |
1253 | // Combine entities and relations
1254 | return [...entities, ...relations];
1255 | }
1256 |
1257 | /**
1258 | * Get all relations involving a set of entities
1259 | * @param entityNames Array of entity names
1260 | * @param zone Optional memory zone for all entities, uses defaultZone if not specified
1261 | */
1262 | async getRelationsForEntities(
1263 | entityNames: string[],
1264 | zone?: string
1265 | ): Promise<{
1266 | relations: ESRelation[]
1267 | }> {
1268 | const actualZone = zone || this.defaultZone;
1269 | await this.initialize(actualZone);
1270 |
1271 | if (entityNames.length === 0) {
1272 | return { relations: [] };
1273 | }
1274 |
1275 | // Find all relations where any of these entities are involved
1276 | // We need to search for both directions - as source and as target
1277 | const fromQuery = entityNames.map(name => ({
1278 | bool: {
1279 | must: [
1280 | { term: { from: name } },
1281 | { term: { "fromZone.keyword": actualZone } }
1282 | ]
1283 | }
1284 | }));
1285 |
1286 | const toQuery = entityNames.map(name => ({
1287 | bool: {
1288 | must: [
1289 | { term: { to: name } },
1290 | { term: { "toZone.keyword": actualZone } }
1291 | ]
1292 | }
1293 | }));
1294 |
1295 | const response = await this.client.search({
1296 | index: KG_RELATIONS_INDEX,
1297 | body: {
1298 | query: {
1299 | bool: {
1300 | should: [...fromQuery, ...toQuery],
1301 | minimum_should_match: 1
1302 | }
1303 | },
1304 | size: 1000
1305 | }
1306 | });
1307 |
1308 | const relations = (response as unknown as ESSearchResponse<ESRelation>)
1309 | .hits.hits
1310 | .map(hit => hit._source);
1311 |
1312 | return { relations };
1313 | }
1314 |
1315 | /**
1316 | * Save or update zone metadata
1317 | * @param name Zone name
1318 | * @param description Optional description
1319 | * @param config Optional configuration
1320 | */
1321 | private async saveZoneMetadata(
1322 | name: string,
1323 | description?: string,
1324 | config?: Record<string, any>
1325 | ): Promise<void> {
1326 | await this.initialize();
1327 |
1328 | const now = new Date().toISOString();
1329 |
1330 | // Check if zone metadata exists
1331 | let existing: ZoneMetadata | null = null;
1332 | try {
1333 | const response = await this.client.get({
1334 | index: KG_METADATA_INDEX,
1335 | id: `zone:${name}`
1336 | });
1337 | existing = response._source as ZoneMetadata;
1338 | } catch (error) {
1339 | // Zone doesn't exist yet
1340 | }
1341 |
1342 | const metadata: ZoneMetadata = {
1343 | name,
1344 | description: description || existing?.description,
1345 | shortDescription: existing?.shortDescription,
1346 | createdAt: existing?.createdAt || now,
1347 | lastModified: now,
1348 | config: config || existing?.config
1349 | };
1350 |
1351 | await this.client.index({
1352 | index: KG_METADATA_INDEX,
1353 | id: `zone:${name}`,
1354 | document: metadata,
1355 | refresh: true
1356 | });
1357 | }
1358 |
1359 | /**
1360 | * List all available memory zones
1361 | * @param reason Optional reason for listing zones, used for AI filtering
1362 | */
1363 | async listMemoryZones(reason?: string): Promise<ZoneMetadata[]> {
1364 | await this.initialize();
1365 |
1366 | try {
1367 | // First try getting zones from metadata
1368 | const response = await this.client.search({
1369 | index: KG_METADATA_INDEX,
1370 | body: {
1371 | query: { match_all: {} },
1372 | size: 1000
1373 | }
1374 | });
1375 |
1376 | const zones = response.hits.hits.map(hit => hit._source as ZoneMetadata);
1377 |
1378 | if (zones.length > 0) {
1379 | // Update cache with all known zones
1380 | zones.forEach(zone => {
1381 | this.existingZonesCache[zone.name] = true;
1382 | });
1383 |
1384 | return zones;
1385 | }
1386 | } catch (error) {
1387 | console.warn('Error getting zones from metadata, falling back to index detection:', error);
1388 | }
1389 |
1390 | // Fallback to listing indices (for backward compatibility)
1391 | const indicesResponse = await this.client.indices.get({
1392 | index: `${KG_INDEX_PREFIX}@*`
1393 | });
1394 |
1395 | // Extract zone names from index names
1396 | const zoneNames = Object.keys(indicesResponse)
1397 | .filter(indexName => indexName.startsWith(`${KG_INDEX_PREFIX}@`))
1398 | .map(indexName => indexName.substring(KG_INDEX_PREFIX.length + 1)); // +1 for the @ symbol
1399 |
1400 | // Convert to metadata format
1401 | const now = new Date().toISOString();
1402 | const zones = zoneNames.map(name => ({
1403 | name,
1404 | createdAt: now,
1405 | lastModified: now
1406 | }));
1407 |
1408 | // Update cache with all detected zones
1409 | zones.forEach(zone => {
1410 | this.existingZonesCache[zone.name] = true;
1411 | });
1412 |
1413 | // Save detected zones to metadata for future
1414 | for (const zone of zones) {
1415 | await this.saveZoneMetadata(zone.name, `Zone detected from index: ${getIndexName(zone.name)}`);
1416 | }
1417 |
1418 | return zones;
1419 | }
1420 |
1421 | /**
1422 | * Add a new memory zone (creates the index if it doesn't exist)
1423 | * @param zone Zone name to add
1424 | * @param description Optional description of the zone
1425 | * @param config Optional configuration for the zone
1426 | */
1427 | async addMemoryZone(
1428 | zone: string,
1429 | description?: string,
1430 | config?: Record<string, any>
1431 | ): Promise<boolean> {
1432 | if (!zone || zone === 'default') {
1433 | throw new Error('Invalid zone name. Cannot be empty or "default".');
1434 | }
1435 |
1436 | // Initialize the index for this zone
1437 | await this.initialize(zone);
1438 |
1439 | // Add to metadata
1440 | await this.saveZoneMetadata(zone, description, config);
1441 |
1442 | // Update the cache
1443 | this.existingZonesCache[zone] = true;
1444 |
1445 | return true;
1446 | }
1447 |
1448 | /**
1449 | * Get metadata for a specific zone
1450 | * @param zone Zone name
1451 | */
1452 | async getZoneMetadata(zone: string): Promise<ZoneMetadata | null> {
1453 | await this.initialize();
1454 |
1455 | try {
1456 | const response = await this.client.get({
1457 | index: KG_METADATA_INDEX,
1458 | id: `zone:${zone}`
1459 | });
1460 | return response._source as ZoneMetadata;
1461 | } catch (error) {
1462 | return null;
1463 | }
1464 | }
1465 |
1466 | /**
1467 | * Delete a memory zone and all its data
1468 | * @param zone Zone name to delete
1469 | */
1470 | async deleteMemoryZone(zone: string): Promise<boolean> {
1471 | if (zone === 'default') {
1472 | throw new Error('Cannot delete the default zone.');
1473 | }
1474 |
1475 | await this.initialize();
1476 |
1477 | try {
1478 | const indexName = this.getIndexForZone(zone);
1479 |
1480 | // Check if index exists before trying to delete it
1481 | const indexExists = await this.client.indices.exists({
1482 | index: indexName
1483 | });
1484 |
1485 | if (indexExists) {
1486 | // Delete the index
1487 | await this.client.indices.delete({
1488 | index: indexName
1489 | });
1490 | console.error(`Deleted index: ${indexName}`);
1491 | }
1492 |
1493 | // Check if metadata exists before trying to delete it
1494 | try {
1495 | const metadataExists = await this.client.exists({
1496 | index: KG_METADATA_INDEX,
1497 | id: `zone:${zone}`
1498 | });
1499 |
1500 | if (metadataExists) {
1501 | // Delete from metadata
1502 | await this.client.delete({
1503 | index: KG_METADATA_INDEX,
1504 | id: `zone:${zone}`
1505 | });
1506 | }
1507 | } catch (metadataError) {
1508 | // Log but continue even if metadata deletion fails
1509 | console.error(`Warning: Error checking/deleting metadata for zone ${zone}:`, metadataError.message);
1510 | }
1511 |
1512 | // Remove from initialized indices cache
1513 | this.initializedIndices.delete(indexName);
1514 |
1515 | // Update the zones cache
1516 | delete this.existingZonesCache[zone];
1517 |
1518 | // Clean up relations for this zone
1519 | try {
1520 | await this.client.deleteByQuery({
1521 | index: KG_RELATIONS_INDEX,
1522 | body: {
1523 | query: {
1524 | bool: {
1525 | should: [
1526 | { term: { fromZone: zone } },
1527 | { term: { toZone: zone } }
1528 | ],
1529 | minimum_should_match: 1
1530 | }
1531 | }
1532 | },
1533 | refresh: true
1534 | });
1535 | } catch (relationError) {
1536 | // Log but continue even if relation cleanup fails
1537 | console.error(`Warning: Error cleaning up relations for zone ${zone}:`, relationError.message);
1538 | }
1539 |
1540 | return true;
1541 | } catch (error) {
1542 | console.error(`Error deleting zone ${zone}:`, error);
1543 | return false;
1544 | }
1545 | }
1546 |
1547 | /**
1548 | * Get statistics for a memory zone
1549 | * @param zone Zone name, uses defaultZone if not specified
1550 | */
1551 | async getMemoryZoneStats(zone?: string): Promise<{
1552 | zone: string;
1553 | entityCount: number;
1554 | relationCount: number;
1555 | entityTypes: Record<string, number>;
1556 | relationTypes: Record<string, number>;
1557 | }> {
1558 | const actualZone = zone || this.defaultZone;
1559 | await this.initialize(actualZone);
1560 |
1561 | const indexName = this.getIndexForZone(actualZone);
1562 |
1563 | // Get total counts
1564 | const countResponse = await this.client.count({
1565 | index: indexName,
1566 | body: {
1567 | query: {
1568 | term: { type: 'entity' }
1569 | }
1570 | }
1571 | });
1572 | const entityCount = countResponse.count;
1573 |
1574 | const relationCountResponse = await this.client.count({
1575 | index: indexName,
1576 | body: {
1577 | query: {
1578 | term: { type: 'relation' }
1579 | }
1580 | }
1581 | });
1582 | const relationCount = relationCountResponse.count;
1583 |
1584 | // Get entity type distribution
1585 | const entityTypesResponse = await this.client.search({
1586 | index: indexName,
1587 | body: {
1588 | size: 0,
1589 | query: {
1590 | term: { type: 'entity' }
1591 | },
1592 | aggs: {
1593 | entity_types: {
1594 | terms: {
1595 | field: 'entityType',
1596 | size: 100
1597 | }
1598 | }
1599 | }
1600 | }
1601 | });
1602 |
1603 | const entityTypes: Record<string, number> = {};
1604 | const entityTypeAggs = entityTypesResponse.aggregations as any;
1605 | const entityTypeBuckets = entityTypeAggs?.entity_types?.buckets || [];
1606 | entityTypeBuckets.forEach((bucket: any) => {
1607 | entityTypes[bucket.key] = bucket.doc_count;
1608 | });
1609 |
1610 | // Get relation type distribution
1611 | const relationTypesResponse = await this.client.search({
1612 | index: indexName,
1613 | body: {
1614 | size: 0,
1615 | query: {
1616 | term: { type: 'relation' }
1617 | },
1618 | aggs: {
1619 | relation_types: {
1620 | terms: {
1621 | field: 'relationType',
1622 | size: 100
1623 | }
1624 | }
1625 | }
1626 | }
1627 | });
1628 |
1629 | const relationTypes: Record<string, number> = {};
1630 | const relationTypeAggs = relationTypesResponse.aggregations as any;
1631 | const relationTypeBuckets = relationTypeAggs?.relation_types?.buckets || [];
1632 | relationTypeBuckets.forEach((bucket: any) => {
1633 | relationTypes[bucket.key] = bucket.doc_count;
1634 | });
1635 |
1636 | return {
1637 | zone: actualZone,
1638 | entityCount,
1639 | relationCount,
1640 | entityTypes,
1641 | relationTypes
1642 | };
1643 | }
1644 |
1645 | /**
1646 | * Export all knowledge graph data, optionally limiting to specific zones
1647 | * @param zones Optional array of zone names to export, exports all zones if not specified
1648 | */
1649 | async exportAllData(zones?: string[]): Promise<{
1650 | entities: ESEntity[],
1651 | relations: ESRelation[],
1652 | zones: ZoneMetadata[]
1653 | }> {
1654 | await this.initialize();
1655 |
1656 | // Get all zones or filter to specified zones
1657 | const allZones = await this.listMemoryZones();
1658 | const zonesToExport = zones
1659 | ? allZones.filter(zone => zones.includes(zone.name))
1660 | : allZones;
1661 |
1662 | if (zonesToExport.length === 0) {
1663 | return { entities: [], relations: [], zones: [] };
1664 | }
1665 |
1666 | // Collect all entities from each zone
1667 | const entities: ESEntity[] = [];
1668 | for (const zone of zonesToExport) {
1669 | const zoneData = await this.exportData(zone.name);
1670 | const zoneEntities = zoneData.filter(item => item.type === 'entity') as ESEntity[];
1671 | entities.push(...zoneEntities);
1672 | }
1673 |
1674 | // Get all relations
1675 | let relations: ESRelation[] = [];
1676 | if (zones) {
1677 | // If specific zones are specified, only get relations involving those zones
1678 | const relationResponse = await this.client.search({
1679 | index: KG_RELATIONS_INDEX,
1680 | body: {
1681 | query: {
1682 | bool: {
1683 | should: [
1684 | ...zonesToExport.map(zone => ({ term: { fromZone: zone.name } })),
1685 | ...zonesToExport.map(zone => ({ term: { toZone: zone.name } }))
1686 | ],
1687 | minimum_should_match: 1
1688 | }
1689 | },
1690 | size: 10000
1691 | }
1692 | });
1693 |
1694 | relations = relationResponse.hits.hits.map(hit => hit._source) as ESRelation[];
1695 | } else {
1696 | // If no zones specified, get all relations
1697 | const relationResponse = await this.client.search({
1698 | index: KG_RELATIONS_INDEX,
1699 | body: {
1700 | query: { match_all: {} },
1701 | size: 10000
1702 | }
1703 | });
1704 |
1705 | relations = relationResponse.hits.hits.map(hit => hit._source) as ESRelation[];
1706 | }
1707 |
1708 | return {
1709 | entities,
1710 | relations,
1711 | zones: zonesToExport
1712 | };
1713 | }
1714 |
1715 | /**
1716 | * Add observations to an existing entity
1717 | * @param name Entity name
1718 | * @param observations Array of observation strings to add
1719 | * @param zone Optional memory zone name, uses defaultZone if not specified
1720 | * @returns The updated entity
1721 | */
1722 | async addObservations(name: string, observations: string[], zone?: string): Promise<ESEntity> {
1723 | const actualZone = zone || this.defaultZone;
1724 |
1725 | // Get existing entity
1726 | const entity = await this.getEntity(name, actualZone);
1727 | if (!entity) {
1728 | throw new Error(`Entity "${name}" not found in zone "${actualZone}"`);
1729 | }
1730 |
1731 | // Add new observations to the existing ones
1732 | const updatedObservations = [
1733 | ...entity.observations,
1734 | ...observations
1735 | ];
1736 |
1737 | // Update the entity
1738 | const updatedEntity = await this.saveEntity({
1739 | name: entity.name,
1740 | entityType: entity.entityType,
1741 | observations: updatedObservations,
1742 | relevanceScore: entity.relevanceScore
1743 | }, actualZone);
1744 |
1745 | return updatedEntity;
1746 | }
1747 |
1748 | /**
1749 | * Mark an entity as important or not important
1750 | * @param name Entity name
1751 | * @param important Whether the entity is important
1752 | * @param zone Optional memory zone name, uses defaultZone if not specified
1753 | * @param options Optional configuration options
1754 | * @param options.autoCreateMissingEntities Whether to automatically create missing entities (default: false)
1755 | * @returns The updated entity
1756 | */
1757 | async markImportant(
1758 | name: string,
1759 | important: boolean,
1760 | zone?: string,
1761 | options?: {
1762 | autoCreateMissingEntities?: boolean;
1763 | }
1764 | ): Promise<ESEntity> {
1765 | return this.updateEntityRelevanceScore(name, important ? 10 : 0.1, zone, options);
1766 | }
1767 |
1768 | /**
1769 | * Mark an entity as important or not important
1770 | * @param name Entity name
1771 | * @param important Whether the entity is important
1772 | * @param zone Optional memory zone name, uses defaultZone if not specified
1773 | * @param options Optional configuration options
1774 | * @param options.autoCreateMissingEntities Whether to automatically create missing entities (default: false)
1775 | * @returns The updated entity
1776 | */
1777 | async updateEntityRelevanceScore(
1778 | name: string,
1779 | ratio: number,
1780 | zone?: string,
1781 | options?: {
1782 | autoCreateMissingEntities?: boolean;
1783 | }
1784 | ): Promise<ESEntity> {
1785 | const actualZone = zone || this.defaultZone;
1786 |
1787 | // Default to false for auto-creation (different from saveRelation)
1788 | const autoCreateMissingEntities = options?.autoCreateMissingEntities ?? false;
1789 |
1790 | // Get existing entity
1791 |
1792 | // Get existing entity
1793 | let entity = await this.getEntity(name, actualZone);
1794 |
1795 | // If entity doesn't exist
1796 | if (!entity) {
1797 | if (autoCreateMissingEntities) {
1798 | // Auto-create the entity with unknown type
1799 | entity = await this.saveEntity({
1800 | name: name,
1801 | entityType: 'unknown',
1802 | observations: [],
1803 | relevanceScore: 1.0
1804 | }, actualZone);
1805 | } else {
1806 | throw new Error(`Entity "${name}" not found in zone "${actualZone}"`);
1807 | }
1808 | }
1809 |
1810 | // Calculate the new relevance score
1811 | // If marking as important, multiply by 10
1812 | // If removing importance, divide by 10
1813 | const baseRelevanceScore = entity.relevanceScore || 1.0;
1814 | const newRelevanceScore = ratio > 1.0
1815 | ? Math.min(25, baseRelevanceScore * ratio)
1816 | : Math.max(0.01, baseRelevanceScore * ratio);
1817 |
1818 | // Update entity with new relevance score
1819 | const updatedEntity = await this.saveEntity({
1820 | name: entity.name,
1821 | entityType: entity.entityType,
1822 | observations: entity.observations,
1823 | relevanceScore: newRelevanceScore
1824 | }, actualZone);
1825 |
1826 | return updatedEntity;
1827 | }
1828 |
1829 | /**
1830 | * Get recent entities
1831 | * @param limit Maximum number of entities to return
1832 | * @param includeObservations Whether to include observations
1833 | * @param zone Optional memory zone name, uses defaultZone if not specified
1834 | * @returns Array of recent entities
1835 | */
1836 | async getRecentEntities(limit: number, includeObservations: boolean, zone?: string): Promise<ESEntity[]> {
1837 | const actualZone = zone || this.defaultZone;
1838 |
1839 | // Search with empty query but sort by recency
1840 | const searchParams: ESSearchParams = {
1841 | query: "*", // Use wildcard instead of empty query to match all documents
1842 | limit: limit,
1843 | sortBy: 'recent', // Sort by recency
1844 | includeObservations
1845 | };
1846 |
1847 | // Add zone if specified
1848 | if (actualZone) {
1849 | (searchParams as any).zone = actualZone;
1850 | }
1851 |
1852 | const results = await this.search(searchParams);
1853 |
1854 | // Filter to only include entities
1855 | return results.hits.hits
1856 | .filter((hit: any) => hit._source.type === 'entity')
1857 | .map((hit: any) => hit._source);
1858 | }
1859 |
1860 | /**
1861 | * Copy entities from one zone to another
1862 | * @param entityNames Array of entity names to copy
1863 | * @param sourceZone Source zone to copy from
1864 | * @param targetZone Target zone to copy to
1865 | * @param options Optional configuration
1866 | * @param options.copyRelations Whether to copy relations involving these entities (default: true)
1867 | * @param options.overwrite Whether to overwrite entities if they already exist in target zone (default: false)
1868 | * @returns Result of the copy operation
1869 | */
1870 | async copyEntitiesBetweenZones(
1871 | entityNames: string[],
1872 | sourceZone: string,
1873 | targetZone: string,
1874 | options?: {
1875 | copyRelations?: boolean;
1876 | overwrite?: boolean;
1877 | }
1878 | ): Promise<{
1879 | entitiesCopied: string[];
1880 | entitiesSkipped: { name: string; reason: string }[];
1881 | relationsCopied: number;
1882 | }> {
1883 | if (sourceZone === targetZone) {
1884 | throw new Error('Source and target zones must be different');
1885 | }
1886 |
1887 | // Default options
1888 | const copyRelations = options?.copyRelations !== false;
1889 | const overwrite = options?.overwrite === true;
1890 |
1891 | await this.initialize(sourceZone);
1892 | await this.initialize(targetZone);
1893 |
1894 | const result = {
1895 | entitiesCopied: [] as string[],
1896 | entitiesSkipped: [] as { name: string; reason: string }[],
1897 | relationsCopied: 0
1898 | };
1899 |
1900 | // Get entities from source zone
1901 | for (const name of entityNames) {
1902 | // Get the entity from source zone
1903 | const entity = await this.getEntityWithoutUpdatingLastRead(name, sourceZone);
1904 | if (!entity) {
1905 | result.entitiesSkipped.push({
1906 | name,
1907 | reason: `Entity not found in source zone '${sourceZone}'`
1908 | });
1909 | continue;
1910 | }
1911 |
1912 | // Check if entity exists in target zone
1913 | const existingEntity = await this.getEntityWithoutUpdatingLastRead(name, targetZone);
1914 | if (existingEntity && !overwrite) {
1915 | result.entitiesSkipped.push({
1916 | name,
1917 | reason: `Entity already exists in target zone '${targetZone}' and overwrite is disabled`
1918 | });
1919 | continue;
1920 | }
1921 |
1922 | // Copy the entity to target zone
1923 | const { ...entityCopy } = entity;
1924 | delete entityCopy.zone; // Zone will be set by saveEntity
1925 |
1926 | try {
1927 | await this.saveEntity(entityCopy, targetZone);
1928 | result.entitiesCopied.push(name);
1929 | } catch (error) {
1930 | result.entitiesSkipped.push({
1931 | name,
1932 | reason: `Error copying entity: ${(error as Error).message}`
1933 | });
1934 | continue;
1935 | }
1936 | }
1937 |
1938 | // Copy relations if requested
1939 | if (copyRelations && result.entitiesCopied.length > 0) {
1940 | // Get all relations for these entities in source zone
1941 | const { relations } = await this.getRelationsForEntities(result.entitiesCopied, sourceZone);
1942 |
1943 | // Filter to only include relations where both entities were copied
1944 | // or relations between copied entities and entities that already exist in target zone
1945 | const relationsToCreate: ESRelation[] = [];
1946 |
1947 | for (const relation of relations) {
1948 | let fromExists = result.entitiesCopied.includes(relation.from);
1949 | let toExists = result.entitiesCopied.includes(relation.to);
1950 |
1951 | // If one side of the relation wasn't copied, check if it exists in target zone
1952 | if (!fromExists) {
1953 | const fromEntityInTarget = await this.getEntityWithoutUpdatingLastRead(relation.from, targetZone);
1954 | fromExists = !!fromEntityInTarget;
1955 | }
1956 |
1957 | if (!toExists) {
1958 | const toEntityInTarget = await this.getEntityWithoutUpdatingLastRead(relation.to, targetZone);
1959 | toExists = !!toEntityInTarget;
1960 | }
1961 |
1962 | // Only create relations where both sides exist
1963 | if (fromExists && toExists) {
1964 | relationsToCreate.push(relation);
1965 | }
1966 | }
1967 |
1968 | // Save the filtered relations
1969 | for (const relation of relationsToCreate) {
1970 | try {
1971 | await this.saveRelation({
1972 | from: relation.from,
1973 | to: relation.to,
1974 | relationType: relation.relationType
1975 | }, targetZone, targetZone);
1976 |
1977 | result.relationsCopied++;
1978 | } catch (error) {
1979 | console.error(`Error copying relation from ${relation.from} to ${relation.to}: ${(error as Error).message}`);
1980 | }
1981 | }
1982 | }
1983 |
1984 | return result;
1985 | }
1986 |
1987 | /**
1988 | * Move entities from one zone to another (copy + delete from source)
1989 | * @param entityNames Array of entity names to move
1990 | * @param sourceZone Source zone to move from
1991 | * @param targetZone Target zone to move to
1992 | * @param options Optional configuration
1993 | * @param options.moveRelations Whether to move relations involving these entities (default: true)
1994 | * @param options.overwrite Whether to overwrite entities if they already exist in target zone (default: false)
1995 | * @returns Result of the move operation
1996 | */
1997 | async moveEntitiesBetweenZones(
1998 | entityNames: string[],
1999 | sourceZone: string,
2000 | targetZone: string,
2001 | options?: {
2002 | moveRelations?: boolean;
2003 | overwrite?: boolean;
2004 | }
2005 | ): Promise<{
2006 | entitiesMoved: string[];
2007 | entitiesSkipped: { name: string; reason: string }[];
2008 | relationsMoved: number;
2009 | }> {
2010 | if (sourceZone === targetZone) {
2011 | throw new Error('Source and target zones must be different');
2012 | }
2013 |
2014 | // Default options
2015 | const moveRelations = options?.moveRelations !== false;
2016 |
2017 | // First copy the entities
2018 | const copyResult = await this.copyEntitiesBetweenZones(
2019 | entityNames,
2020 | sourceZone,
2021 | targetZone,
2022 | {
2023 | copyRelations: moveRelations,
2024 | overwrite: options?.overwrite
2025 | }
2026 | );
2027 |
2028 | const result = {
2029 | entitiesMoved: [] as string[],
2030 | entitiesSkipped: copyResult.entitiesSkipped,
2031 | relationsMoved: copyResult.relationsCopied
2032 | };
2033 |
2034 | // Delete copied entities from source zone
2035 | for (const name of copyResult.entitiesCopied) {
2036 | try {
2037 | // Don't cascade relations when deleting from source, as we've already copied them
2038 | await this.deleteEntity(name, sourceZone, { cascadeRelations: false });
2039 | result.entitiesMoved.push(name);
2040 | } catch (error) {
2041 | // If deletion fails, add to skipped list but keep the entity in the moved list
2042 | // since it was successfully copied
2043 | result.entitiesSkipped.push({
2044 | name,
2045 | reason: `Entity was copied but could not be deleted from source: ${(error as Error).message}`
2046 | });
2047 | }
2048 | }
2049 |
2050 | return result;
2051 | }
2052 |
2053 | /**
2054 | * Merge two or more zones into a target zone
2055 | * @param sourceZones Array of source zone names to merge from
2056 | * @param targetZone Target zone to merge into
2057 | * @param options Optional configuration
2058 | * @param options.deleteSourceZones Whether to delete source zones after merging (default: false)
2059 | * @param options.overwriteConflicts How to handle entity name conflicts (default: 'skip')
2060 | * @returns Result of the merge operation
2061 | */
2062 | async mergeZones(
2063 | sourceZones: string[],
2064 | targetZone: string,
2065 | options?: {
2066 | deleteSourceZones?: boolean;
2067 | overwriteConflicts?: 'skip' | 'overwrite' | 'rename';
2068 | }
2069 | ): Promise<{
2070 | mergedZones: string[];
2071 | failedZones: { zone: string; reason: string }[];
2072 | entitiesCopied: number;
2073 | entitiesSkipped: number;
2074 | relationsCopied: number;
2075 | }> {
2076 | // Validate parameters
2077 | if (sourceZones.includes(targetZone)) {
2078 | throw new Error('Target zone cannot be included in source zones');
2079 | }
2080 |
2081 | if (sourceZones.length === 0) {
2082 | throw new Error('At least one source zone must be specified');
2083 | }
2084 |
2085 | // Default options
2086 | const deleteSourceZones = options?.deleteSourceZones === true;
2087 | const overwriteConflicts = options?.overwriteConflicts || 'skip';
2088 |
2089 | // Initialize target zone
2090 | await this.initialize(targetZone);
2091 |
2092 | const result = {
2093 | mergedZones: [] as string[],
2094 | failedZones: [] as { zone: string; reason: string }[],
2095 | entitiesCopied: 0,
2096 | entitiesSkipped: 0,
2097 | relationsCopied: 0
2098 | };
2099 |
2100 | // Process each source zone
2101 | for (const sourceZone of sourceZones) {
2102 | try {
2103 | // Get all entities from source zone
2104 | const allEntities = await this.searchEntities({
2105 | query: '*',
2106 | limit: 10000,
2107 | zone: sourceZone
2108 | });
2109 |
2110 | if (allEntities.length === 0) {
2111 | result.failedZones.push({
2112 | zone: sourceZone,
2113 | reason: 'Zone has no entities'
2114 | });
2115 | continue;
2116 | }
2117 |
2118 | // Extract entity names
2119 | const entityNames = allEntities.map(entity => entity.name);
2120 |
2121 | // Process according to conflict resolution strategy
2122 | if (overwriteConflicts === 'rename') {
2123 | // For 'rename' strategy, we need to check each entity and rename if necessary
2124 | for (const entity of allEntities) {
2125 | const existingEntity = await this.getEntityWithoutUpdatingLastRead(entity.name, targetZone);
2126 |
2127 | if (existingEntity) {
2128 | // Entity exists in target zone, generate a new name
2129 | const newName = `${entity.name}_from_${sourceZone}`;
2130 |
2131 | // Create a copy with the new name
2132 | const entityCopy = { ...entity, name: newName };
2133 | delete entityCopy.zone; // Zone will be set by saveEntity
2134 |
2135 | try {
2136 | await this.saveEntity(entityCopy, targetZone);
2137 | result.entitiesCopied++;
2138 | } catch (error) {
2139 | result.entitiesSkipped++;
2140 | console.error(`Error copying entity ${entity.name} with new name ${newName}: ${(error as Error).message}`);
2141 | }
2142 | } else {
2143 | // Entity doesn't exist, copy as is
2144 | const entityCopy = { ...entity };
2145 | delete entityCopy.zone; // Zone will be set by saveEntity
2146 |
2147 | try {
2148 | await this.saveEntity(entityCopy, targetZone);
2149 | result.entitiesCopied++;
2150 | } catch (error) {
2151 | result.entitiesSkipped++;
2152 | console.error(`Error copying entity ${entity.name}: ${(error as Error).message}`);
2153 | }
2154 | }
2155 | }
2156 |
2157 | // Now copy relations, adjusting for renamed entities
2158 | const { relations } = await this.getRelationsForEntities(entityNames, sourceZone);
2159 |
2160 | for (const relation of relations) {
2161 | try {
2162 | // Check if entities were renamed
2163 | let fromName = relation.from;
2164 | let toName = relation.to;
2165 |
2166 | const fromEntityInTarget = await this.getEntityWithoutUpdatingLastRead(fromName, targetZone);
2167 | if (!fromEntityInTarget) {
2168 | // Check if it was renamed
2169 | const renamedFromName = `${fromName}_from_${sourceZone}`;
2170 | const renamedFromEntityInTarget = await this.getEntityWithoutUpdatingLastRead(renamedFromName, targetZone);
2171 | if (renamedFromEntityInTarget) {
2172 | fromName = renamedFromName;
2173 | }
2174 | }
2175 |
2176 | const toEntityInTarget = await this.getEntityWithoutUpdatingLastRead(toName, targetZone);
2177 | if (!toEntityInTarget) {
2178 | // Check if it was renamed
2179 | const renamedToName = `${toName}_from_${sourceZone}`;
2180 | const renamedToEntityInTarget = await this.getEntityWithoutUpdatingLastRead(renamedToName, targetZone);
2181 | if (renamedToEntityInTarget) {
2182 | toName = renamedToName;
2183 | }
2184 | }
2185 |
2186 | // Only create relation if both entities exist
2187 | if (await this.getEntityWithoutUpdatingLastRead(fromName, targetZone) &&
2188 | await this.getEntityWithoutUpdatingLastRead(toName, targetZone)) {
2189 | await this.saveRelation({
2190 | from: fromName,
2191 | to: toName,
2192 | relationType: relation.relationType
2193 | }, targetZone, targetZone);
2194 |
2195 | result.relationsCopied++;
2196 | }
2197 | } catch (error) {
2198 | console.error(`Error copying relation from ${relation.from} to ${relation.to}: ${(error as Error).message}`);
2199 | }
2200 | }
2201 | } else {
2202 | // For 'skip' or 'overwrite' strategy, use copyEntitiesBetweenZones
2203 | const copyResult = await this.copyEntitiesBetweenZones(
2204 | entityNames,
2205 | sourceZone,
2206 | targetZone,
2207 | {
2208 | copyRelations: true,
2209 | overwrite: overwriteConflicts === 'overwrite'
2210 | }
2211 | );
2212 |
2213 | result.entitiesCopied += copyResult.entitiesCopied.length;
2214 | result.entitiesSkipped += copyResult.entitiesSkipped.length;
2215 | result.relationsCopied += copyResult.relationsCopied;
2216 | }
2217 |
2218 | // Mark as successfully merged
2219 | result.mergedZones.push(sourceZone);
2220 |
2221 | // Delete source zone if requested
2222 | if (deleteSourceZones) {
2223 | await this.deleteMemoryZone(sourceZone);
2224 | }
2225 | } catch (error) {
2226 | result.failedZones.push({
2227 | zone: sourceZone,
2228 | reason: (error as Error).message
2229 | });
2230 | }
2231 | }
2232 |
2233 | return result;
2234 | }
2235 |
2236 | /**
2237 | * Search for entities by name or other attributes
2238 | * @param params Search parameters
2239 | * @returns Array of matching entities
2240 | */
2241 | async searchEntities(params: {
2242 | query: string;
2243 | entityTypes?: string[];
2244 | limit?: number;
2245 | includeObservations?: boolean;
2246 | zone?: string;
2247 | }): Promise<ESEntity[]> {
2248 | // Use existing search method with appropriate parameters
2249 | const searchResponse = await this.search({
2250 | query: params.query,
2251 | entityTypes: params.entityTypes,
2252 | limit: params.limit,
2253 | offset: 0,
2254 | zone: params.zone
2255 | });
2256 |
2257 | // Extract entities from the search response
2258 | const entities: ESEntity[] = [];
2259 | if (searchResponse && searchResponse.hits && searchResponse.hits.hits) {
2260 | for (const hit of searchResponse.hits.hits) {
2261 | if (hit._source && hit._source.type === 'entity') {
2262 | entities.push(hit._source as ESEntity);
2263 | }
2264 | }
2265 | }
2266 |
2267 | return entities;
2268 | }
2269 |
2270 | /**
2271 | * Update zone metadata with new descriptions
2272 | * @param name Zone name
2273 | * @param description Full description
2274 | * @param shortDescription Short description
2275 | * @param config Optional configuration
2276 | */
2277 | async updateZoneDescriptions(
2278 | name: string,
2279 | description: string,
2280 | shortDescription: string,
2281 | config?: Record<string, any>
2282 | ): Promise<void> {
2283 | await this.initialize();
2284 |
2285 | const now = new Date().toISOString();
2286 |
2287 | // Check if zone metadata exists
2288 | let existing: ZoneMetadata | null = null;
2289 | try {
2290 | const response = await this.client.get({
2291 | index: KG_METADATA_INDEX,
2292 | id: `zone:${name}`
2293 | });
2294 | existing = response._source as ZoneMetadata;
2295 | } catch (error) {
2296 | // Zone doesn't exist yet, create it first
2297 | if (!await this.zoneExists(name)) {
2298 | await this.addMemoryZone(name);
2299 | }
2300 | }
2301 |
2302 | const metadata: ZoneMetadata = {
2303 | name,
2304 | description,
2305 | shortDescription,
2306 | createdAt: existing?.createdAt || now,
2307 | lastModified: now,
2308 | config: config || existing?.config
2309 | };
2310 |
2311 | await this.client.index({
2312 | index: KG_METADATA_INDEX,
2313 | id: `zone:${name}`,
2314 | body: metadata,
2315 | refresh: true
2316 | });
2317 |
2318 | console.log(`Updated descriptions for zone: ${name}`);
2319 | }
2320 |
2321 | /**
2322 | * High-level search method that returns clean entity data for user-facing applications
2323 | * This method acts as a wrapper around the raw search, with additional processing and AI filtering
2324 | *
2325 | * @param params Search parameters including query, filters, and AI-related fields
2326 | * @returns Clean entity and relation data, filtered by AI if informationNeeded is provided
2327 | */
2328 | async userSearch(params: {
2329 | query: string;
2330 | entityTypes?: string[];
2331 | limit?: number;
2332 | includeObservations?: boolean;
2333 | sortBy?: 'relevance' | 'recent' | 'importance';
2334 | zone?: string;
2335 | informationNeeded?: string;
2336 | reason?: string;
2337 | }): Promise<{
2338 | entities: Array<{
2339 | name: string;
2340 | entityType: string;
2341 | observations?: string[];
2342 | lastRead?: string;
2343 | lastWrite?: string;
2344 | }>;
2345 | relations: Array<{
2346 | from: string;
2347 | to: string;
2348 | type: string;
2349 | fromZone: string;
2350 | toZone: string;
2351 | }>;
2352 | }> {
2353 | // Set default values
2354 | const includeObservations = params.includeObservations ?? false;
2355 | const defaultLimit = includeObservations ? 5 : 20;
2356 | const zone = params.zone || this.defaultZone;
2357 | const informationNeeded = params.informationNeeded;
2358 | const reason = params.reason;
2359 |
2360 | // If informationNeeded is provided, increase the limit to get more results
2361 | // that will be filtered later by the AI
2362 | const searchLimit = informationNeeded ?
2363 | (params.limit ? params.limit * 4 : defaultLimit * 4) :
2364 | (params.limit || defaultLimit);
2365 |
2366 | // Prepare search parameters for the raw search
2367 | const searchParams: ESSearchParams = {
2368 | query: params.query,
2369 | entityTypes: params.entityTypes,
2370 | limit: searchLimit,
2371 | sortBy: params.sortBy,
2372 | includeObservations,
2373 | zone,
2374 | informationNeeded,
2375 | reason
2376 | };
2377 |
2378 | // Perform the raw search
2379 | const results = await this.search(searchParams);
2380 |
2381 | // Transform the results to a clean format, removing unnecessary fields
2382 | const entities = results.hits.hits
2383 | .filter(hit => hit._source.type === 'entity')
2384 | .map(hit => {
2385 | const entity: {
2386 | name: string;
2387 | entityType: string;
2388 | observations?: string[];
2389 | lastRead?: string;
2390 | lastWrite?: string;
2391 | } = {
2392 | name: (hit._source as ESEntity).name,
2393 | entityType: (hit._source as ESEntity).entityType,
2394 | };
2395 |
2396 | // Only include observations and timestamps if requested
2397 | if (includeObservations) {
2398 | entity.observations = (hit._source as ESEntity).observations;
2399 | entity.lastWrite = (hit._source as ESEntity).lastWrite;
2400 | entity.lastRead = (hit._source as ESEntity).lastRead;
2401 | }
2402 |
2403 | return entity;
2404 | });
2405 |
2406 | // Apply AI filtering if informationNeeded is provided and AI is available
2407 | let filteredEntities = entities;
2408 | if (informationNeeded && GroqAI.isEnabled && entities.length > 0) {
2409 | try {
2410 | // Get relevant entity names using AI filtering
2411 | const usefulness = await GroqAI.filterSearchResults(entities, informationNeeded, reason);
2412 |
2413 | // If AI filtering returned null (error case), use original entities
2414 | if (usefulness === null) {
2415 | console.warn('AI filtering returned null, using original results');
2416 | filteredEntities = entities.slice(0, params.limit || defaultLimit);
2417 | } else {
2418 | // Filter entities to only include those with a usefulness score
2419 | filteredEntities = entities.filter(entity =>
2420 | usefulness[entity.name] !== undefined
2421 | );
2422 |
2423 | // Sort entities by their relevance score from highest to lowest
2424 | filteredEntities.sort((a, b) => {
2425 | const scoreA = usefulness[a.name] || 0;
2426 | const scoreB = usefulness[b.name] || 0;
2427 | return scoreB - scoreA;
2428 | });
2429 |
2430 | const usefulEntities = filteredEntities.filter(entity => usefulness[entity.name] >= 60);
2431 | const definatelyNotUsefulEntities = filteredEntities.filter(entity => usefulness[entity.name] < 20);
2432 |
2433 | // for each useful entities, increase the relevanceScore
2434 | for (const entity of usefulEntities) {
2435 | this.updateEntityRelevanceScore(entity.name, (usefulness[entity.name] + 45) * 0.01, zone);
2436 | }
2437 |
2438 | // for each definately not useful entities, decrease the relevanceScore
2439 | for (const entity of definatelyNotUsefulEntities) {
2440 | this.updateEntityRelevanceScore(entity.name, 0.8 + usefulness[entity.name] * 0.01, zone);
2441 | }
2442 |
2443 | // If no entities were found relevant, fall back to the original results
2444 | if (filteredEntities.length === 0) {
2445 | filteredEntities = entities.slice(0, params.limit || defaultLimit);
2446 | } else {
2447 | // Limit the filtered results to the requested amount
2448 | filteredEntities = filteredEntities.slice(0, params.limit || defaultLimit);
2449 | }
2450 | }
2451 | } catch (error) {
2452 | console.error('Error applying AI filtering:', error);
2453 | // Fall back to the original results but limit to the requested amount
2454 | filteredEntities = entities.slice(0, params.limit || defaultLimit);
2455 | }
2456 | } else if (entities.length > (params.limit || defaultLimit)) {
2457 | // If we're not using AI filtering but retrieved more results due to the doubled limit,
2458 | // limit the results to the originally requested amount
2459 | filteredEntities = entities.slice(0, params.limit || defaultLimit);
2460 | }
2461 |
2462 | // Get relations between these entities
2463 | const entityNames = filteredEntities.map(e => e.name);
2464 | const { relations } = await this.getRelationsForEntities(entityNames, zone);
2465 |
2466 | // Map relations to a clean format
2467 | const formattedRelations = relations.map(r => ({
2468 | from: r.from,
2469 | to: r.to,
2470 | type: r.relationType,
2471 | fromZone: r.fromZone,
2472 | toZone: r.toZone
2473 | }));
2474 |
2475 | return {
2476 | entities: filteredEntities,
2477 | relations: formattedRelations
2478 | };
2479 | }
2480 | }
```