main
1package memory
2
3import (
4 "encoding/json"
5 "fmt"
6 "os"
7 "strings"
8 "sync"
9
10 "github.com/xlgmokha/mcp/pkg/mcp"
11)
12
13// KnowledgeGraph represents the in-memory knowledge graph
14type KnowledgeGraph struct {
15 Entities map[string]*Entity `json:"entities"`
16 Relations map[string]Relation `json:"relations"`
17}
18
19// Entity represents an entity in the knowledge graph
20type Entity struct {
21 Name string `json:"name"`
22 EntityType string `json:"entityType"`
23 Observations []string `json:"observations"`
24}
25
26// Relation represents a relationship between entities
27type Relation struct {
28 From string `json:"from"`
29 To string `json:"to"`
30 RelationType string `json:"relationType"`
31}
32
33// MemoryOperations provides memory graph operations
34type MemoryOperations struct {
35 memoryFile string
36 graph *KnowledgeGraph
37 mu sync.RWMutex
38 loaded bool
39}
40
41// NewMemoryOperations creates a new MemoryOperations helper
42func NewMemoryOperations(memoryFile string) *MemoryOperations {
43 return &MemoryOperations{
44 memoryFile: memoryFile,
45 graph: &KnowledgeGraph{
46 Entities: make(map[string]*Entity),
47 Relations: make(map[string]Relation),
48 },
49 }
50}
51
52// New creates a new Memory MCP server
53func New(memoryFile string) *mcp.Server {
54 memory := NewMemoryOperations(memoryFile)
55 builder := mcp.NewServerBuilder("mcp-memory", "1.0.0")
56
57 // Add create_entities tool
58 builder.AddTool(mcp.NewTool("create_entities", "Create multiple new entities in the knowledge graph", map[string]interface{}{
59 "type": "object",
60 "properties": map[string]interface{}{
61 "entities": map[string]interface{}{
62 "type": "array",
63 "items": map[string]interface{}{
64 "type": "object",
65 "properties": map[string]interface{}{
66 "name": map[string]interface{}{
67 "type": "string",
68 "description": "The name of the entity",
69 },
70 "entityType": map[string]interface{}{
71 "type": "string",
72 "description": "The type of the entity",
73 },
74 "observations": map[string]interface{}{
75 "type": "array",
76 "items": map[string]interface{}{
77 "type": "string",
78 },
79 "description": "An array of observation contents associated with the entity",
80 },
81 },
82 "required": []string{"name", "entityType", "observations"},
83 },
84 },
85 },
86 "required": []string{"entities"},
87 }, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
88 if err := memory.ensureGraphLoaded(); err != nil {
89 return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
90 }
91
92 memory.mu.Lock()
93 defer memory.mu.Unlock()
94
95 entitiesArg, ok := req.Arguments["entities"]
96 if !ok {
97 return mcp.NewToolError("entities parameter is required"), nil
98 }
99
100 entitiesSlice, ok := entitiesArg.([]interface{})
101 if !ok {
102 return mcp.NewToolError("entities must be an array"), nil
103 }
104
105 var createdEntities []string
106
107 for _, entityArg := range entitiesSlice {
108 entityMap, ok := entityArg.(map[string]interface{})
109 if !ok {
110 return mcp.NewToolError("each entity must be an object"), nil
111 }
112
113 name, ok := entityMap["name"].(string)
114 if !ok {
115 return mcp.NewToolError("entity name must be a string"), nil
116 }
117
118 entityType, ok := entityMap["entityType"].(string)
119 if !ok {
120 return mcp.NewToolError("entity type must be a string"), nil
121 }
122
123 observationsArg, ok := entityMap["observations"]
124 if !ok {
125 return mcp.NewToolError("entity observations are required"), nil
126 }
127
128 observationsSlice, ok := observationsArg.([]interface{})
129 if !ok {
130 return mcp.NewToolError("entity observations must be an array"), nil
131 }
132
133 var observations []string
134 for _, obs := range observationsSlice {
135 obsStr, ok := obs.(string)
136 if !ok {
137 return mcp.NewToolError("each observation must be a string"), nil
138 }
139 observations = append(observations, obsStr)
140 }
141
142 memory.graph.Entities[name] = &Entity{
143 Name: name,
144 EntityType: entityType,
145 Observations: observations,
146 }
147 createdEntities = append(createdEntities, name)
148 }
149
150 if err := memory.saveGraph(); err != nil {
151 return mcp.NewToolError(fmt.Sprintf("Failed to save graph: %v", err)), nil
152 }
153
154 return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Created %d entities: %s", len(createdEntities), strings.Join(createdEntities, ", ")))), nil
155 }))
156
157 // Add create_relations tool
158 builder.AddTool(mcp.NewTool("create_relations", "Create multiple new relations between entities in the knowledge graph. Relations should be in active voice", map[string]interface{}{
159 "type": "object",
160 "properties": map[string]interface{}{
161 "relations": map[string]interface{}{
162 "type": "array",
163 "items": map[string]interface{}{
164 "type": "object",
165 "properties": map[string]interface{}{
166 "from": map[string]interface{}{
167 "type": "string",
168 "description": "The name of the entity where the relation starts",
169 },
170 "to": map[string]interface{}{
171 "type": "string",
172 "description": "The name of the entity where the relation ends",
173 },
174 "relationType": map[string]interface{}{
175 "type": "string",
176 "description": "The type of the relation",
177 },
178 },
179 "required": []string{"from", "to", "relationType"},
180 },
181 },
182 },
183 "required": []string{"relations"},
184 }, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
185 if err := memory.ensureGraphLoaded(); err != nil {
186 return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
187 }
188
189 memory.mu.Lock()
190 defer memory.mu.Unlock()
191
192 relationsArg, ok := req.Arguments["relations"]
193 if !ok {
194 return mcp.NewToolError("relations parameter is required"), nil
195 }
196
197 relationsSlice, ok := relationsArg.([]interface{})
198 if !ok {
199 return mcp.NewToolError("relations must be an array"), nil
200 }
201
202 var createdRelations []string
203
204 for _, relationArg := range relationsSlice {
205 relationMap, ok := relationArg.(map[string]interface{})
206 if !ok {
207 return mcp.NewToolError("each relation must be an object"), nil
208 }
209
210 from, ok := relationMap["from"].(string)
211 if !ok {
212 return mcp.NewToolError("relation 'from' must be a string"), nil
213 }
214
215 to, ok := relationMap["to"].(string)
216 if !ok {
217 return mcp.NewToolError("relation 'to' must be a string"), nil
218 }
219
220 relationType, ok := relationMap["relationType"].(string)
221 if !ok {
222 return mcp.NewToolError("relation type must be a string"), nil
223 }
224
225 if _, exists := memory.graph.Entities[from]; !exists {
226 return mcp.NewToolError(fmt.Sprintf("entity '%s' does not exist", from)), nil
227 }
228
229 if _, exists := memory.graph.Entities[to]; !exists {
230 return mcp.NewToolError(fmt.Sprintf("entity '%s' does not exist", to)), nil
231 }
232
233 relationKey := fmt.Sprintf("%s-%s-%s", from, relationType, to)
234 memory.graph.Relations[relationKey] = Relation{
235 From: from,
236 To: to,
237 RelationType: relationType,
238 }
239 createdRelations = append(createdRelations, fmt.Sprintf("%s %s %s", from, relationType, to))
240 }
241
242 if err := memory.saveGraph(); err != nil {
243 return mcp.NewToolError(fmt.Sprintf("Failed to save graph: %v", err)), nil
244 }
245
246 return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Created %d relations: %s", len(createdRelations), strings.Join(createdRelations, ", ")))), nil
247 }))
248
249 // Add add_observations tool
250 builder.AddTool(mcp.NewTool("add_observations", "Add new observations to existing entities in the knowledge graph", map[string]interface{}{
251 "type": "object",
252 "properties": map[string]interface{}{
253 "observations": map[string]interface{}{
254 "type": "array",
255 "items": map[string]interface{}{
256 "type": "object",
257 "properties": map[string]interface{}{
258 "entityName": map[string]interface{}{
259 "type": "string",
260 "description": "The name of the entity to add the observations to",
261 },
262 "contents": map[string]interface{}{
263 "type": "array",
264 "items": map[string]interface{}{
265 "type": "string",
266 },
267 "description": "An array of observation contents to add",
268 },
269 },
270 "required": []string{"entityName", "contents"},
271 },
272 },
273 },
274 "required": []string{"observations"},
275 }, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
276 if err := memory.ensureGraphLoaded(); err != nil {
277 return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
278 }
279
280 memory.mu.Lock()
281 defer memory.mu.Unlock()
282
283 observationsArg, ok := req.Arguments["observations"]
284 if !ok {
285 return mcp.NewToolError("observations parameter is required"), nil
286 }
287
288 observationsSlice, ok := observationsArg.([]interface{})
289 if !ok {
290 return mcp.NewToolError("observations must be an array"), nil
291 }
292
293 var addedCount int
294
295 for _, obsArg := range observationsSlice {
296 obsMap, ok := obsArg.(map[string]interface{})
297 if !ok {
298 return mcp.NewToolError("each observation must be an object"), nil
299 }
300
301 entityName, ok := obsMap["entityName"].(string)
302 if !ok {
303 return mcp.NewToolError("entity name must be a string"), nil
304 }
305
306 contentsArg, ok := obsMap["contents"]
307 if !ok {
308 return mcp.NewToolError("observation contents are required"), nil
309 }
310
311 contentsSlice, ok := contentsArg.([]interface{})
312 if !ok {
313 return mcp.NewToolError("observation contents must be an array"), nil
314 }
315
316 entity, exists := memory.graph.Entities[entityName]
317 if !exists {
318 return mcp.NewToolError(fmt.Sprintf("entity '%s' does not exist", entityName)), nil
319 }
320
321 for _, content := range contentsSlice {
322 contentStr, ok := content.(string)
323 if !ok {
324 return mcp.NewToolError("each observation content must be a string"), nil
325 }
326 entity.Observations = append(entity.Observations, contentStr)
327 addedCount++
328 }
329 }
330
331 if err := memory.saveGraph(); err != nil {
332 return mcp.NewToolError(fmt.Sprintf("Failed to save graph: %v", err)), nil
333 }
334
335 return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Added %d observations", addedCount))), nil
336 }))
337
338 // Add delete_entities tool
339 builder.AddTool(mcp.NewTool("delete_entities", "Delete multiple entities and their associated relations from the knowledge graph", map[string]interface{}{
340 "type": "object",
341 "properties": map[string]interface{}{
342 "entityNames": map[string]interface{}{
343 "type": "array",
344 "items": map[string]interface{}{
345 "type": "string",
346 },
347 "description": "An array of entity names to delete",
348 },
349 },
350 "required": []string{"entityNames"},
351 }, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
352 if err := memory.ensureGraphLoaded(); err != nil {
353 return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
354 }
355
356 memory.mu.Lock()
357 defer memory.mu.Unlock()
358
359 entityNamesArg, ok := req.Arguments["entityNames"]
360 if !ok {
361 return mcp.NewToolError("entityNames parameter is required"), nil
362 }
363
364 entityNamesSlice, ok := entityNamesArg.([]interface{})
365 if !ok {
366 return mcp.NewToolError("entityNames must be an array"), nil
367 }
368
369 var deletedEntities []string
370 var deletedRelations int
371
372 for _, nameArg := range entityNamesSlice {
373 name, ok := nameArg.(string)
374 if !ok {
375 return mcp.NewToolError("each entity name must be a string"), nil
376 }
377
378 if _, exists := memory.graph.Entities[name]; !exists {
379 continue
380 }
381
382 delete(memory.graph.Entities, name)
383 deletedEntities = append(deletedEntities, name)
384
385 for key, relation := range memory.graph.Relations {
386 if relation.From == name || relation.To == name {
387 delete(memory.graph.Relations, key)
388 deletedRelations++
389 }
390 }
391 }
392
393 if err := memory.saveGraph(); err != nil {
394 return mcp.NewToolError(fmt.Sprintf("Failed to save graph: %v", err)), nil
395 }
396
397 return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Deleted %d entities and %d relations", len(deletedEntities), deletedRelations))), nil
398 }))
399
400 // Add delete_observations tool
401 builder.AddTool(mcp.NewTool("delete_observations", "Delete specific observations from entities in the knowledge graph", map[string]interface{}{
402 "type": "object",
403 "properties": map[string]interface{}{
404 "deletions": map[string]interface{}{
405 "type": "array",
406 "items": map[string]interface{}{
407 "type": "object",
408 "properties": map[string]interface{}{
409 "entityName": map[string]interface{}{
410 "type": "string",
411 "description": "The name of the entity containing the observations",
412 },
413 "observations": map[string]interface{}{
414 "type": "array",
415 "items": map[string]interface{}{
416 "type": "string",
417 },
418 "description": "An array of observations to delete",
419 },
420 },
421 "required": []string{"entityName", "observations"},
422 },
423 },
424 },
425 "required": []string{"deletions"},
426 }, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
427 if err := memory.ensureGraphLoaded(); err != nil {
428 return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
429 }
430
431 memory.mu.Lock()
432 defer memory.mu.Unlock()
433
434 deletionsArg, ok := req.Arguments["deletions"]
435 if !ok {
436 return mcp.NewToolError("deletions parameter is required"), nil
437 }
438
439 deletionsSlice, ok := deletionsArg.([]interface{})
440 if !ok {
441 return mcp.NewToolError("deletions must be an array"), nil
442 }
443
444 var deletedCount int
445
446 for _, delArg := range deletionsSlice {
447 delMap, ok := delArg.(map[string]interface{})
448 if !ok {
449 return mcp.NewToolError("each deletion must be an object"), nil
450 }
451
452 entityName, ok := delMap["entityName"].(string)
453 if !ok {
454 return mcp.NewToolError("entity name must be a string"), nil
455 }
456
457 observationsArg, ok := delMap["observations"]
458 if !ok {
459 return mcp.NewToolError("observations are required"), nil
460 }
461
462 observationsSlice, ok := observationsArg.([]interface{})
463 if !ok {
464 return mcp.NewToolError("observations must be an array"), nil
465 }
466
467 entity, exists := memory.graph.Entities[entityName]
468 if !exists {
469 continue
470 }
471
472 for _, obsArg := range observationsSlice {
473 obsStr, ok := obsArg.(string)
474 if !ok {
475 continue
476 }
477
478 for i, existingObs := range entity.Observations {
479 if existingObs == obsStr {
480 entity.Observations = append(entity.Observations[:i], entity.Observations[i+1:]...)
481 deletedCount++
482 break
483 }
484 }
485 }
486 }
487
488 if err := memory.saveGraph(); err != nil {
489 return mcp.NewToolError(fmt.Sprintf("Failed to save graph: %v", err)), nil
490 }
491
492 return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Deleted %d observations", deletedCount))), nil
493 }))
494
495 // Add delete_relations tool
496 builder.AddTool(mcp.NewTool("delete_relations", "Delete multiple relations from the knowledge graph", map[string]interface{}{
497 "type": "object",
498 "properties": map[string]interface{}{
499 "relations": map[string]interface{}{
500 "type": "array",
501 "items": map[string]interface{}{
502 "type": "object",
503 "properties": map[string]interface{}{
504 "from": map[string]interface{}{
505 "type": "string",
506 "description": "The name of the entity where the relation starts",
507 },
508 "to": map[string]interface{}{
509 "type": "string",
510 "description": "The name of the entity where the relation ends",
511 },
512 "relationType": map[string]interface{}{
513 "type": "string",
514 "description": "The type of the relation",
515 },
516 },
517 "required": []string{"from", "to", "relationType"},
518 },
519 },
520 },
521 "required": []string{"relations"},
522 }, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
523 if err := memory.ensureGraphLoaded(); err != nil {
524 return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
525 }
526
527 memory.mu.Lock()
528 defer memory.mu.Unlock()
529
530 relationsArg, ok := req.Arguments["relations"]
531 if !ok {
532 return mcp.NewToolError("relations parameter is required"), nil
533 }
534
535 relationsSlice, ok := relationsArg.([]interface{})
536 if !ok {
537 return mcp.NewToolError("relations must be an array"), nil
538 }
539
540 var deletedCount int
541
542 for _, relationArg := range relationsSlice {
543 relationMap, ok := relationArg.(map[string]interface{})
544 if !ok {
545 continue
546 }
547
548 from, ok := relationMap["from"].(string)
549 if !ok {
550 continue
551 }
552
553 to, ok := relationMap["to"].(string)
554 if !ok {
555 continue
556 }
557
558 relationType, ok := relationMap["relationType"].(string)
559 if !ok {
560 continue
561 }
562
563 relationKey := fmt.Sprintf("%s-%s-%s", from, relationType, to)
564 if _, exists := memory.graph.Relations[relationKey]; exists {
565 delete(memory.graph.Relations, relationKey)
566 deletedCount++
567 }
568 }
569
570 if err := memory.saveGraph(); err != nil {
571 return mcp.NewToolError(fmt.Sprintf("Failed to save graph: %v", err)), nil
572 }
573
574 return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Deleted %d relations", deletedCount))), nil
575 }))
576
577 // Add read_graph tool
578 builder.AddTool(mcp.NewTool("read_graph", "Read the entire knowledge graph", map[string]interface{}{
579 "type": "object",
580 "properties": map[string]interface{}{},
581 }, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
582 if err := memory.ensureGraphLoaded(); err != nil {
583 return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
584 }
585
586 memory.mu.RLock()
587 defer memory.mu.RUnlock()
588
589 graphData, err := json.MarshalIndent(memory.graph, "", " ")
590 if err != nil {
591 return mcp.NewToolError(fmt.Sprintf("Failed to serialize graph: %v", err)), nil
592 }
593
594 return mcp.NewToolResult(mcp.NewTextContent(string(graphData))), nil
595 }))
596
597 // Add search_nodes tool
598 builder.AddTool(mcp.NewTool("search_nodes", "Search for nodes in the knowledge graph based on a query", map[string]interface{}{
599 "type": "object",
600 "properties": map[string]interface{}{
601 "query": map[string]interface{}{
602 "type": "string",
603 "description": "The search query to match against entity names, types, and observation content",
604 },
605 },
606 "required": []string{"query"},
607 }, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
608 if err := memory.ensureGraphLoaded(); err != nil {
609 return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
610 }
611
612 memory.mu.RLock()
613 defer memory.mu.RUnlock()
614
615 query, ok := req.Arguments["query"].(string)
616 if !ok {
617 return mcp.NewToolError("query parameter is required"), nil
618 }
619
620 query = strings.ToLower(query)
621 var matchedEntities []*Entity
622
623 for _, entity := range memory.graph.Entities {
624 if strings.Contains(strings.ToLower(entity.Name), query) ||
625 strings.Contains(strings.ToLower(entity.EntityType), query) {
626 matchedEntities = append(matchedEntities, entity)
627 continue
628 }
629
630 for _, observation := range entity.Observations {
631 if strings.Contains(strings.ToLower(observation), query) {
632 matchedEntities = append(matchedEntities, entity)
633 break
634 }
635 }
636 }
637
638 resultData, err := json.MarshalIndent(matchedEntities, "", " ")
639 if err != nil {
640 return mcp.NewToolError(fmt.Sprintf("Failed to serialize results: %v", err)), nil
641 }
642
643 return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Found %d matching entities:\n%s", len(matchedEntities), string(resultData)))), nil
644 }))
645
646 // Add open_nodes tool
647 builder.AddTool(mcp.NewTool("open_nodes", "Open specific nodes in the knowledge graph by their names", map[string]interface{}{
648 "type": "object",
649 "properties": map[string]interface{}{
650 "names": map[string]interface{}{
651 "type": "array",
652 "items": map[string]interface{}{
653 "type": "string",
654 },
655 "description": "An array of entity names to retrieve",
656 },
657 },
658 "required": []string{"names"},
659 }, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
660 if err := memory.ensureGraphLoaded(); err != nil {
661 return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
662 }
663
664 memory.mu.RLock()
665 defer memory.mu.RUnlock()
666
667 namesArg, ok := req.Arguments["names"]
668 if !ok {
669 return mcp.NewToolError("names parameter is required"), nil
670 }
671
672 namesSlice, ok := namesArg.([]interface{})
673 if !ok {
674 return mcp.NewToolError("names must be an array"), nil
675 }
676
677 var foundEntities []*Entity
678
679 for _, nameArg := range namesSlice {
680 name, ok := nameArg.(string)
681 if !ok {
682 continue
683 }
684
685 if entity, exists := memory.graph.Entities[name]; exists {
686 foundEntities = append(foundEntities, entity)
687 }
688 }
689
690 resultData, err := json.MarshalIndent(foundEntities, "", " ")
691 if err != nil {
692 return mcp.NewToolError(fmt.Sprintf("Failed to serialize results: %v", err)), nil
693 }
694
695 return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Found %d entities:\n%s", len(foundEntities), string(resultData)))), nil
696 }))
697
698 // Add knowledge-query prompt
699 builder.AddPrompt(mcp.NewPrompt("knowledge-query", "Prompt for querying and exploring the knowledge graph", []mcp.PromptArgument{
700 {
701 Name: "query",
702 Description: "What you want to search for or ask about in the knowledge graph",
703 Required: true,
704 },
705 {
706 Name: "context",
707 Description: "Additional context about your question (optional)",
708 Required: false,
709 },
710 }, func(req mcp.GetPromptRequest) (mcp.GetPromptResult, error) {
711 query, hasQuery := req.Arguments["query"].(string)
712 context, hasContext := req.Arguments["context"].(string)
713
714 if !hasQuery || query == "" {
715 return mcp.GetPromptResult{}, fmt.Errorf("query argument is required")
716 }
717
718 var messages []mcp.PromptMessage
719
720 userContent := fmt.Sprintf(`I want to search the knowledge graph for: %s`, query)
721 if hasContext && context != "" {
722 userContent += fmt.Sprintf("\n\nAdditional context: %s", context)
723 }
724
725 messages = append(messages, mcp.PromptMessage{
726 Role: "user",
727 Content: mcp.NewTextContent(userContent),
728 })
729
730 assistantContent := fmt.Sprintf(`I'll help you search the knowledge graph for "%s". Here are some strategies you can use:
731
732**Search Commands:**
733- Use "search_nodes" tool with query: "%s"
734- Use "read_graph" tool to see the entire knowledge structure
735- Use "open_nodes" tool if you know specific entity names
736
737**What to look for:**
738- Entities with names containing "%s"
739- Entity types that match your query
740- Observations that mention "%s"
741- Related entities through relationships
742
743**Next steps:**
7441. Start with a broad search using search_nodes
7452. Examine the results to find relevant entities
7463. Use open_nodes to get detailed information about specific entities
7474. Look at relationships to find connected information
748
749Would you like me to search for this information in the knowledge graph?`, query, query, query, query)
750
751 messages = append(messages, mcp.PromptMessage{
752 Role: "assistant",
753 Content: mcp.NewTextContent(assistantContent),
754 })
755
756 description := fmt.Sprintf("Knowledge graph search guidance for: %s", query)
757
758 return mcp.GetPromptResult{
759 Description: description,
760 Messages: messages,
761 }, nil
762 }))
763
764 // Add memory:// pattern resource
765 builder.AddResource(mcp.NewResource(
766 "memory://graph",
767 "Knowledge Graph",
768 "application/json",
769 func(req mcp.ReadResourceRequest) (mcp.ReadResourceResult, error) {
770 if err := memory.ensureGraphLoaded(); err != nil {
771 return mcp.ReadResourceResult{}, fmt.Errorf("failed to load graph: %v", err)
772 }
773
774 memory.mu.RLock()
775 defer memory.mu.RUnlock()
776
777 graphData, err := json.MarshalIndent(memory.graph, "", " ")
778 if err != nil {
779 return mcp.ReadResourceResult{}, fmt.Errorf("failed to serialize graph: %v", err)
780 }
781
782 return mcp.ReadResourceResult{
783 Contents: []mcp.Content{
784 mcp.NewTextContent(string(graphData)),
785 },
786 }, nil
787 },
788 ))
789
790 // Add knowledge graph root
791 rootName := fmt.Sprintf("Knowledge Graph (%d entities, %d relations)", len(memory.graph.Entities), len(memory.graph.Relations))
792 builder.AddRoot(mcp.NewRoot("memory://graph", rootName))
793
794 return builder.Build()
795}
796
797// Helper methods for MemoryOperations
798
799func (memory *MemoryOperations) ensureGraphLoaded() error {
800 memory.mu.Lock()
801 defer memory.mu.Unlock()
802
803 if !memory.loaded {
804 memory.loaded = true
805 return memory.loadGraphInternal()
806 }
807 return nil
808}
809
810func (memory *MemoryOperations) loadGraphInternal() error {
811 if memory.memoryFile == "" {
812 return nil
813 }
814
815 data, err := os.ReadFile(memory.memoryFile)
816 if os.IsNotExist(err) {
817 return nil
818 }
819 if err != nil {
820 return fmt.Errorf("failed to read memory file: %v", err)
821 }
822
823 if len(data) == 0 {
824 return nil
825 }
826
827 var loadedGraph KnowledgeGraph
828 if err := json.Unmarshal(data, &loadedGraph); err != nil {
829 return fmt.Errorf("failed to parse memory file: %v", err)
830 }
831
832 if loadedGraph.Entities != nil {
833 memory.graph.Entities = loadedGraph.Entities
834 }
835 if loadedGraph.Relations != nil {
836 memory.graph.Relations = loadedGraph.Relations
837 }
838
839 return nil
840}
841
842func (memory *MemoryOperations) saveGraph() error {
843 if memory.memoryFile == "" {
844 return nil
845 }
846
847 data, err := json.MarshalIndent(memory.graph, "", " ")
848 if err != nil {
849 return fmt.Errorf("failed to serialize graph: %v", err)
850 }
851
852 if err := os.WriteFile(memory.memoryFile, data, 0644); err != nil {
853 return fmt.Errorf("failed to write memory file: %v", err)
854 }
855
856 return nil
857}