path: root/pkg/memory/server.go

package memory

import (
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"sync"

	"github.com/xlgmokha/mcp/pkg/mcp"
)

// KnowledgeGraph represents the in-memory knowledge graph
type KnowledgeGraph struct {
	Entities  map[string]*Entity  `json:"entities"`
	Relations map[string]Relation `json:"relations"`
}

// Entity represents an entity in the knowledge graph
type Entity struct {
	Name         string   `json:"name"`
	EntityType   string   `json:"entityType"`
	Observations []string `json:"observations"`
}

// Relation represents a relationship between entities
type Relation struct {
	From         string `json:"from"`
	To           string `json:"to"`
	RelationType string `json:"relationType"`
}

// MemoryOperations provides memory graph operations
type MemoryOperations struct {
	memoryFile string
	graph      *KnowledgeGraph
	mu         sync.RWMutex
	loaded     bool
}

// NewMemoryOperations creates a new MemoryOperations helper
func NewMemoryOperations(memoryFile string) *MemoryOperations {
	return &MemoryOperations{
		memoryFile: memoryFile,
		graph: &KnowledgeGraph{
			Entities:  make(map[string]*Entity),
			Relations: make(map[string]Relation),
		},
	}
}

// New creates a new Memory MCP server
func New(memoryFile string) *mcp.Server {
	memory := NewMemoryOperations(memoryFile)
	builder := mcp.NewServerBuilder("mcp-memory", "1.0.0")

	// Add create_entities tool
	builder.AddTool(mcp.NewTool("create_entities", "Create multiple new entities in the knowledge graph", map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"entities": map[string]interface{}{
				"type": "array",
				"items": map[string]interface{}{
					"type": "object",
					"properties": map[string]interface{}{
						"name": map[string]interface{}{
							"type":        "string",
							"description": "The name of the entity",
						},
						"entityType": map[string]interface{}{
							"type":        "string",
							"description": "The type of the entity",
						},
						"observations": map[string]interface{}{
							"type": "array",
							"items": map[string]interface{}{
								"type": "string",
							},
							"description": "An array of observation contents associated with the entity",
						},
					},
					"required": []string{"name", "entityType", "observations"},
				},
			},
		},
		"required": []string{"entities"},
	}, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
		if err := memory.ensureGraphLoaded(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
		}

		memory.mu.Lock()
		defer memory.mu.Unlock()

		entitiesArg, ok := req.Arguments["entities"]
		if !ok {
			return mcp.NewToolError("entities parameter is required"), nil
		}

		entitiesSlice, ok := entitiesArg.([]interface{})
		if !ok {
			return mcp.NewToolError("entities must be an array"), nil
		}

		var createdEntities []string

		for _, entityArg := range entitiesSlice {
			entityMap, ok := entityArg.(map[string]interface{})
			if !ok {
				return mcp.NewToolError("each entity must be an object"), nil
			}

			name, ok := entityMap["name"].(string)
			if !ok {
				return mcp.NewToolError("entity name must be a string"), nil
			}

			entityType, ok := entityMap["entityType"].(string)
			if !ok {
				return mcp.NewToolError("entity type must be a string"), nil
			}

			observationsArg, ok := entityMap["observations"]
			if !ok {
				return mcp.NewToolError("entity observations are required"), nil
			}

			observationsSlice, ok := observationsArg.([]interface{})
			if !ok {
				return mcp.NewToolError("entity observations must be an array"), nil
			}

			var observations []string
			for _, obs := range observationsSlice {
				obsStr, ok := obs.(string)
				if !ok {
					return mcp.NewToolError("each observation must be a string"), nil
				}
				observations = append(observations, obsStr)
			}

			memory.graph.Entities[name] = &Entity{
				Name:         name,
				EntityType:   entityType,
				Observations: observations,
			}
			createdEntities = append(createdEntities, name)
		}

		if err := memory.saveGraph(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to save graph: %v", err)), nil
		}

		return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Created %d entities: %s", len(createdEntities), strings.Join(createdEntities, ", ")))), nil
	}))

	// Add create_relations tool
	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{}{
		"type": "object",
		"properties": map[string]interface{}{
			"relations": map[string]interface{}{
				"type": "array",
				"items": map[string]interface{}{
					"type": "object",
					"properties": map[string]interface{}{
						"from": map[string]interface{}{
							"type":        "string",
							"description": "The name of the entity where the relation starts",
						},
						"to": map[string]interface{}{
							"type":        "string",
							"description": "The name of the entity where the relation ends",
						},
						"relationType": map[string]interface{}{
							"type":        "string",
							"description": "The type of the relation",
						},
					},
					"required": []string{"from", "to", "relationType"},
				},
			},
		},
		"required": []string{"relations"},
	}, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
		if err := memory.ensureGraphLoaded(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
		}

		memory.mu.Lock()
		defer memory.mu.Unlock()

		relationsArg, ok := req.Arguments["relations"]
		if !ok {
			return mcp.NewToolError("relations parameter is required"), nil
		}

		relationsSlice, ok := relationsArg.([]interface{})
		if !ok {
			return mcp.NewToolError("relations must be an array"), nil
		}

		var createdRelations []string

		for _, relationArg := range relationsSlice {
			relationMap, ok := relationArg.(map[string]interface{})
			if !ok {
				return mcp.NewToolError("each relation must be an object"), nil
			}

			from, ok := relationMap["from"].(string)
			if !ok {
				return mcp.NewToolError("relation 'from' must be a string"), nil
			}

			to, ok := relationMap["to"].(string)
			if !ok {
				return mcp.NewToolError("relation 'to' must be a string"), nil
			}

			relationType, ok := relationMap["relationType"].(string)
			if !ok {
				return mcp.NewToolError("relation type must be a string"), nil
			}

			if _, exists := memory.graph.Entities[from]; !exists {
				return mcp.NewToolError(fmt.Sprintf("entity '%s' does not exist", from)), nil
			}

			if _, exists := memory.graph.Entities[to]; !exists {
				return mcp.NewToolError(fmt.Sprintf("entity '%s' does not exist", to)), nil
			}

			relationKey := fmt.Sprintf("%s-%s-%s", from, relationType, to)
			memory.graph.Relations[relationKey] = Relation{
				From:         from,
				To:           to,
				RelationType: relationType,
			}
			createdRelations = append(createdRelations, fmt.Sprintf("%s %s %s", from, relationType, to))
		}

		if err := memory.saveGraph(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to save graph: %v", err)), nil
		}

		return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Created %d relations: %s", len(createdRelations), strings.Join(createdRelations, ", ")))), nil
	}))

	// Add add_observations tool
	builder.AddTool(mcp.NewTool("add_observations", "Add new observations to existing entities in the knowledge graph", map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"observations": map[string]interface{}{
				"type": "array",
				"items": map[string]interface{}{
					"type": "object",
					"properties": map[string]interface{}{
						"entityName": map[string]interface{}{
							"type":        "string",
							"description": "The name of the entity to add the observations to",
						},
						"contents": map[string]interface{}{
							"type": "array",
							"items": map[string]interface{}{
								"type": "string",
							},
							"description": "An array of observation contents to add",
						},
					},
					"required": []string{"entityName", "contents"},
				},
			},
		},
		"required": []string{"observations"},
	}, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
		if err := memory.ensureGraphLoaded(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
		}

		memory.mu.Lock()
		defer memory.mu.Unlock()

		observationsArg, ok := req.Arguments["observations"]
		if !ok {
			return mcp.NewToolError("observations parameter is required"), nil
		}

		observationsSlice, ok := observationsArg.([]interface{})
		if !ok {
			return mcp.NewToolError("observations must be an array"), nil
		}

		var addedCount int

		for _, obsArg := range observationsSlice {
			obsMap, ok := obsArg.(map[string]interface{})
			if !ok {
				return mcp.NewToolError("each observation must be an object"), nil
			}

			entityName, ok := obsMap["entityName"].(string)
			if !ok {
				return mcp.NewToolError("entity name must be a string"), nil
			}

			contentsArg, ok := obsMap["contents"]
			if !ok {
				return mcp.NewToolError("observation contents are required"), nil
			}

			contentsSlice, ok := contentsArg.([]interface{})
			if !ok {
				return mcp.NewToolError("observation contents must be an array"), nil
			}

			entity, exists := memory.graph.Entities[entityName]
			if !exists {
				return mcp.NewToolError(fmt.Sprintf("entity '%s' does not exist", entityName)), nil
			}

			for _, content := range contentsSlice {
				contentStr, ok := content.(string)
				if !ok {
					return mcp.NewToolError("each observation content must be a string"), nil
				}
				entity.Observations = append(entity.Observations, contentStr)
				addedCount++
			}
		}

		if err := memory.saveGraph(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to save graph: %v", err)), nil
		}

		return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Added %d observations", addedCount))), nil
	}))

	// Add delete_entities tool
	builder.AddTool(mcp.NewTool("delete_entities", "Delete multiple entities and their associated relations from the knowledge graph", map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"entityNames": map[string]interface{}{
				"type": "array",
				"items": map[string]interface{}{
					"type": "string",
				},
				"description": "An array of entity names to delete",
			},
		},
		"required": []string{"entityNames"},
	}, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
		if err := memory.ensureGraphLoaded(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
		}

		memory.mu.Lock()
		defer memory.mu.Unlock()

		entityNamesArg, ok := req.Arguments["entityNames"]
		if !ok {
			return mcp.NewToolError("entityNames parameter is required"), nil
		}

		entityNamesSlice, ok := entityNamesArg.([]interface{})
		if !ok {
			return mcp.NewToolError("entityNames must be an array"), nil
		}

		var deletedEntities []string
		var deletedRelations int

		for _, nameArg := range entityNamesSlice {
			name, ok := nameArg.(string)
			if !ok {
				return mcp.NewToolError("each entity name must be a string"), nil
			}

			if _, exists := memory.graph.Entities[name]; !exists {
				continue
			}

			delete(memory.graph.Entities, name)
			deletedEntities = append(deletedEntities, name)

			for key, relation := range memory.graph.Relations {
				if relation.From == name || relation.To == name {
					delete(memory.graph.Relations, key)
					deletedRelations++
				}
			}
		}

		if err := memory.saveGraph(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to save graph: %v", err)), nil
		}

		return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Deleted %d entities and %d relations", len(deletedEntities), deletedRelations))), nil
	}))

	// Add delete_observations tool
	builder.AddTool(mcp.NewTool("delete_observations", "Delete specific observations from entities in the knowledge graph", map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"deletions": map[string]interface{}{
				"type": "array",
				"items": map[string]interface{}{
					"type": "object",
					"properties": map[string]interface{}{
						"entityName": map[string]interface{}{
							"type":        "string",
							"description": "The name of the entity containing the observations",
						},
						"observations": map[string]interface{}{
							"type": "array",
							"items": map[string]interface{}{
								"type": "string",
							},
							"description": "An array of observations to delete",
						},
					},
					"required": []string{"entityName", "observations"},
				},
			},
		},
		"required": []string{"deletions"},
	}, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
		if err := memory.ensureGraphLoaded(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
		}

		memory.mu.Lock()
		defer memory.mu.Unlock()

		deletionsArg, ok := req.Arguments["deletions"]
		if !ok {
			return mcp.NewToolError("deletions parameter is required"), nil
		}

		deletionsSlice, ok := deletionsArg.([]interface{})
		if !ok {
			return mcp.NewToolError("deletions must be an array"), nil
		}

		var deletedCount int

		for _, delArg := range deletionsSlice {
			delMap, ok := delArg.(map[string]interface{})
			if !ok {
				return mcp.NewToolError("each deletion must be an object"), nil
			}

			entityName, ok := delMap["entityName"].(string)
			if !ok {
				return mcp.NewToolError("entity name must be a string"), nil
			}

			observationsArg, ok := delMap["observations"]
			if !ok {
				return mcp.NewToolError("observations are required"), nil
			}

			observationsSlice, ok := observationsArg.([]interface{})
			if !ok {
				return mcp.NewToolError("observations must be an array"), nil
			}

			entity, exists := memory.graph.Entities[entityName]
			if !exists {
				continue
			}

			for _, obsArg := range observationsSlice {
				obsStr, ok := obsArg.(string)
				if !ok {
					continue
				}

				for i, existingObs := range entity.Observations {
					if existingObs == obsStr {
						entity.Observations = append(entity.Observations[:i], entity.Observations[i+1:]...)
						deletedCount++
						break
					}
				}
			}
		}

		if err := memory.saveGraph(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to save graph: %v", err)), nil
		}

		return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Deleted %d observations", deletedCount))), nil
	}))

	// Add delete_relations tool
	builder.AddTool(mcp.NewTool("delete_relations", "Delete multiple relations from the knowledge graph", map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"relations": map[string]interface{}{
				"type": "array",
				"items": map[string]interface{}{
					"type": "object",
					"properties": map[string]interface{}{
						"from": map[string]interface{}{
							"type":        "string",
							"description": "The name of the entity where the relation starts",
						},
						"to": map[string]interface{}{
							"type":        "string",
							"description": "The name of the entity where the relation ends",
						},
						"relationType": map[string]interface{}{
							"type":        "string",
							"description": "The type of the relation",
						},
					},
					"required": []string{"from", "to", "relationType"},
				},
			},
		},
		"required": []string{"relations"},
	}, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
		if err := memory.ensureGraphLoaded(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
		}

		memory.mu.Lock()
		defer memory.mu.Unlock()

		relationsArg, ok := req.Arguments["relations"]
		if !ok {
			return mcp.NewToolError("relations parameter is required"), nil
		}

		relationsSlice, ok := relationsArg.([]interface{})
		if !ok {
			return mcp.NewToolError("relations must be an array"), nil
		}

		var deletedCount int

		for _, relationArg := range relationsSlice {
			relationMap, ok := relationArg.(map[string]interface{})
			if !ok {
				continue
			}

			from, ok := relationMap["from"].(string)
			if !ok {
				continue
			}

			to, ok := relationMap["to"].(string)
			if !ok {
				continue
			}

			relationType, ok := relationMap["relationType"].(string)
			if !ok {
				continue
			}

			relationKey := fmt.Sprintf("%s-%s-%s", from, relationType, to)
			if _, exists := memory.graph.Relations[relationKey]; exists {
				delete(memory.graph.Relations, relationKey)
				deletedCount++
			}
		}

		if err := memory.saveGraph(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to save graph: %v", err)), nil
		}

		return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Deleted %d relations", deletedCount))), nil
	}))

	// Add read_graph tool
	builder.AddTool(mcp.NewTool("read_graph", "Read the entire knowledge graph", map[string]interface{}{
		"type":       "object",
		"properties": map[string]interface{}{},
	}, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
		if err := memory.ensureGraphLoaded(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
		}

		memory.mu.RLock()
		defer memory.mu.RUnlock()

		graphData, err := json.MarshalIndent(memory.graph, "", "  ")
		if err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to serialize graph: %v", err)), nil
		}

		return mcp.NewToolResult(mcp.NewTextContent(string(graphData))), nil
	}))

	// Add search_nodes tool
	builder.AddTool(mcp.NewTool("search_nodes", "Search for nodes in the knowledge graph based on a query", map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"query": map[string]interface{}{
				"type":        "string",
				"description": "The search query to match against entity names, types, and observation content",
			},
		},
		"required": []string{"query"},
	}, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
		if err := memory.ensureGraphLoaded(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
		}

		memory.mu.RLock()
		defer memory.mu.RUnlock()

		query, ok := req.Arguments["query"].(string)
		if !ok {
			return mcp.NewToolError("query parameter is required"), nil
		}

		query = strings.ToLower(query)
		var matchedEntities []*Entity

		for _, entity := range memory.graph.Entities {
			if strings.Contains(strings.ToLower(entity.Name), query) ||
				strings.Contains(strings.ToLower(entity.EntityType), query) {
				matchedEntities = append(matchedEntities, entity)
				continue
			}

			for _, observation := range entity.Observations {
				if strings.Contains(strings.ToLower(observation), query) {
					matchedEntities = append(matchedEntities, entity)
					break
				}
			}
		}

		resultData, err := json.MarshalIndent(matchedEntities, "", "  ")
		if err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to serialize results: %v", err)), nil
		}

		return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Found %d matching entities:\n%s", len(matchedEntities), string(resultData)))), nil
	}))

	// Add open_nodes tool
	builder.AddTool(mcp.NewTool("open_nodes", "Open specific nodes in the knowledge graph by their names", map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"names": map[string]interface{}{
				"type": "array",
				"items": map[string]interface{}{
					"type": "string",
				},
				"description": "An array of entity names to retrieve",
			},
		},
		"required": []string{"names"},
	}, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
		if err := memory.ensureGraphLoaded(); err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to load graph: %v", err)), nil
		}

		memory.mu.RLock()
		defer memory.mu.RUnlock()

		namesArg, ok := req.Arguments["names"]
		if !ok {
			return mcp.NewToolError("names parameter is required"), nil
		}

		namesSlice, ok := namesArg.([]interface{})
		if !ok {
			return mcp.NewToolError("names must be an array"), nil
		}

		var foundEntities []*Entity

		for _, nameArg := range namesSlice {
			name, ok := nameArg.(string)
			if !ok {
				continue
			}

			if entity, exists := memory.graph.Entities[name]; exists {
				foundEntities = append(foundEntities, entity)
			}
		}

		resultData, err := json.MarshalIndent(foundEntities, "", "  ")
		if err != nil {
			return mcp.NewToolError(fmt.Sprintf("Failed to serialize results: %v", err)), nil
		}

		return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Found %d entities:\n%s", len(foundEntities), string(resultData)))), nil
	}))

	// Add knowledge-query prompt
	builder.AddPrompt(mcp.NewPrompt("knowledge-query", "Prompt for querying and exploring the knowledge graph", []mcp.PromptArgument{
		{
			Name:        "query",
			Description: "What you want to search for or ask about in the knowledge graph",
			Required:    true,
		},
		{
			Name:        "context",
			Description: "Additional context about your question (optional)",
			Required:    false,
		},
	}, func(req mcp.GetPromptRequest) (mcp.GetPromptResult, error) {
		query, hasQuery := req.Arguments["query"].(string)
		context, hasContext := req.Arguments["context"].(string)

		if !hasQuery || query == "" {
			return mcp.GetPromptResult{}, fmt.Errorf("query argument is required")
		}

		var messages []mcp.PromptMessage

		userContent := fmt.Sprintf(`I want to search the knowledge graph for: %s`, query)
		if hasContext && context != "" {
			userContent += fmt.Sprintf("\n\nAdditional context: %s", context)
		}

		messages = append(messages, mcp.PromptMessage{
			Role:    "user",
			Content: mcp.NewTextContent(userContent),
		})

		assistantContent := fmt.Sprintf(`I'll help you search the knowledge graph for "%s". Here are some strategies you can use:

**Search Commands:**
- Use "search_nodes" tool with query: "%s"
- Use "read_graph" tool to see the entire knowledge structure
- Use "open_nodes" tool if you know specific entity names

**What to look for:**
- Entities with names containing "%s"
- Entity types that match your query
- Observations that mention "%s"
- Related entities through relationships

**Next steps:**
1. Start with a broad search using search_nodes
2. Examine the results to find relevant entities
3. Use open_nodes to get detailed information about specific entities
4. Look at relationships to find connected information

Would you like me to search for this information in the knowledge graph?`, query, query, query, query)

		messages = append(messages, mcp.PromptMessage{
			Role:    "assistant",
			Content: mcp.NewTextContent(assistantContent),
		})

		description := fmt.Sprintf("Knowledge graph search guidance for: %s", query)

		return mcp.GetPromptResult{
			Description: description,
			Messages:    messages,
		}, nil
	}))

	// Add memory:// pattern resource
	builder.AddResource(mcp.NewResource(
		"memory://graph",
		"Knowledge Graph",
		"application/json",
		func(req mcp.ReadResourceRequest) (mcp.ReadResourceResult, error) {
			if err := memory.ensureGraphLoaded(); err != nil {
				return mcp.ReadResourceResult{}, fmt.Errorf("failed to load graph: %v", err)
			}

			memory.mu.RLock()
			defer memory.mu.RUnlock()

			graphData, err := json.MarshalIndent(memory.graph, "", "  ")
			if err != nil {
				return mcp.ReadResourceResult{}, fmt.Errorf("failed to serialize graph: %v", err)
			}

			return mcp.ReadResourceResult{
				Contents: []mcp.Content{
					mcp.NewTextContent(string(graphData)),
				},
			}, nil
		},
	))

	// Add knowledge graph root
	rootName := fmt.Sprintf("Knowledge Graph (%d entities, %d relations)", len(memory.graph.Entities), len(memory.graph.Relations))
	builder.AddRoot(mcp.NewRoot("memory://graph", rootName))

	return builder.Build()
}

// Helper methods for MemoryOperations

func (memory *MemoryOperations) ensureGraphLoaded() error {
	memory.mu.Lock()
	defer memory.mu.Unlock()

	if !memory.loaded {
		memory.loaded = true
		return memory.loadGraphInternal()
	}
	return nil
}

func (memory *MemoryOperations) loadGraphInternal() error {
	if memory.memoryFile == "" {
		return nil
	}

	data, err := os.ReadFile(memory.memoryFile)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("failed to read memory file: %v", err)
	}

	if len(data) == 0 {
		return nil
	}

	var loadedGraph KnowledgeGraph
	if err := json.Unmarshal(data, &loadedGraph); err != nil {
		return fmt.Errorf("failed to parse memory file: %v", err)
	}

	if loadedGraph.Entities != nil {
		memory.graph.Entities = loadedGraph.Entities
	}
	if loadedGraph.Relations != nil {
		memory.graph.Relations = loadedGraph.Relations
	}

	return nil
}

func (memory *MemoryOperations) saveGraph() error {
	if memory.memoryFile == "" {
		return nil
	}

	data, err := json.MarshalIndent(memory.graph, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to serialize graph: %v", err)
	}

	if err := os.WriteFile(memory.memoryFile, data, 0644); err != nil {
		return fmt.Errorf("failed to write memory file: %v", err)
	}

	return nil
}