Memory - add annotations, format return data, improve error handling (#113)

* move Neo4jMemory class to separate file, add ToolResult and structured output

* update changelog, add try-catch for neo4j errors

* Update CHANGELOG.md

* add tool annotations, update tool arg types

* replace List with list type hints

Author: a-s-g93

servers/mcp-neo4j-memory/CHANGELOG.md

@@ -3,8 +3,16 @@
 ### Fixed
 
 ### Changed
+* Update tool return type hints for structured output
+* Move `Neo4jMemory` class and related classes to separate file
+* Change tool responses to return the `ToolResponse` object
+* Updated tool argument types with Pydantic models
 
 ### Added
+* Add structured output to tool responses
+* Add error handling to catch Neo4j specific errors and improve error responses
+* Implement `ToolError` class from FastMCP
+* Add tool annotations
 
 ## v0.2.0
 

servers/mcp-neo4j-memory/src/mcp_neo4j_memory/neo4j_memory.py

@@ -0,0 +1,245 @@
+import logging
+from typing import Any, Dict, List
+
+from neo4j import AsyncDriver, RoutingControl
+from pydantic import BaseModel
+
+
+# Set up logging
+logger = logging.getLogger('mcp_neo4j_memory')
+logger.setLevel(logging.INFO)
+
+# Models for our knowledge graph
+class Entity(BaseModel):
+    name: str
+    type: str
+    observations: List[str]
+
+class Relation(BaseModel):
+    source: str
+    target: str
+    relationType: str
+
+class KnowledgeGraph(BaseModel):
+    entities: List[Entity]
+    relations: List[Relation]
+
+class ObservationAddition(BaseModel):
+    entityName: str
+    observations: List[str]
+
+class ObservationDeletion(BaseModel):
+    entityName: str
+    observations: List[str]
+
+class Neo4jMemory:
+    def __init__(self, neo4j_driver: AsyncDriver):
+        self.driver = neo4j_driver
+
+    async def create_fulltext_index(self):
+        """Create a fulltext search index for entities if it doesn't exist."""
+        try:
+            query = "CREATE FULLTEXT INDEX search IF NOT EXISTS FOR (m:Memory) ON EACH [m.name, m.type, m.observations];"
+            await self.driver.execute_query(query, routing_control=RoutingControl.WRITE)
+            logger.info("Created fulltext search index")
+        except Exception as e:
+            # Index might already exist, which is fine
+            logger.debug(f"Fulltext index creation: {e}")
+
+    async def load_graph(self, filter_query: str = "*"):
+        """Load the entire knowledge graph from Neo4j."""
+        logger.info("Loading knowledge graph from Neo4j")
+        query = """
+            CALL db.index.fulltext.queryNodes('search', $filter) yield node as entity, score
+            OPTIONAL MATCH (entity)-[r]-(other)
+            RETURN collect(distinct {
+                name: entity.name, 
+                type: entity.type, 
+                observations: entity.observations
+            }) as nodes,
+            collect(distinct {
+                source: startNode(r).name, 
+                target: endNode(r).name, 
+                relationType: type(r)
+            }) as relations
+        """
+        
+        result = await self.driver.execute_query(query, {"filter": filter_query}, routing_control=RoutingControl.READ)
+        
+        if not result.records:
+            return KnowledgeGraph(entities=[], relations=[])
+        
+        record = result.records[0]
+        nodes = record.get('nodes', list())
+        rels = record.get('relations', list())
+        
+        entities = [
+            Entity(
+                name=node['name'],
+                type=node['type'],
+                observations=node.get('observations', list())
+            )
+            for node in nodes if node.get('name')
+        ]
+        
+        relations = [
+            Relation(
+                source=rel['source'],
+                target=rel['target'],
+                relationType=rel['relationType']
+            )
+            for rel in rels if rel.get('relationType')
+        ]
+        
+        logger.debug(f"Loaded entities: {entities}")
+        logger.debug(f"Loaded relations: {relations}")
+        
+        return KnowledgeGraph(entities=entities, relations=relations)
+
+    async def create_entities(self, entities: List[Entity]) -> List[Entity]:
+        """Create multiple new entities in the knowledge graph."""
+        logger.info(f"Creating {len(entities)} entities")
+        for entity in entities:
+            query = f"""
+            WITH $entity as entity
+            MERGE (e:Memory {{ name: entity.name }})
+            SET e += entity {{ .type, .observations }}
+            SET e:{entity.type}
+            """
+            await self.driver.execute_query(query, {"entity": entity.model_dump()}, routing_control=RoutingControl.WRITE)
+
+        return entities
+
+    async def create_relations(self, relations: List[Relation]) -> List[Relation]:
+        """Create multiple new relations between entities."""
+        logger.info(f"Creating {len(relations)} relations")
+        for relation in relations:
+            query = f"""
+            WITH $relation as relation
+            MATCH (from:Memory),(to:Memory)
+            WHERE from.name = relation.source
+            AND  to.name = relation.target
+            MERGE (from)-[r:{relation.relationType}]->(to)
+            """
+            
+            await self.driver.execute_query(
+                query, 
+                {"relation": relation.model_dump()},
+                routing_control=RoutingControl.WRITE
+            )
+
+        return relations
+
+    async def add_observations(self, observations: List[ObservationAddition]) -> List[Dict[str, Any]]:
+        """Add new observations to existing entities."""
+        logger.info(f"Adding observations to {len(observations)} entities")
+        query = """
+        UNWIND $observations as obs  
+        MATCH (e:Memory { name: obs.entityName })
+        WITH e, [o in obs.observations WHERE NOT o IN e.observations] as new
+        SET e.observations = coalesce(e.observations,[]) + new
+        RETURN e.name as name, new
+        """
+            
+        result = await self.driver.execute_query(
+            query, 
+            {"observations": [obs.model_dump() for obs in observations]},
+            routing_control=RoutingControl.WRITE
+        )
+
+        results = [{"entityName": record.get("name"), "addedObservations": record.get("new")} for record in result.records]
+        return results
+
+    async def delete_entities(self, entity_names: List[str]) -> None:
+        """Delete multiple entities and their associated relations."""
+        logger.info(f"Deleting {len(entity_names)} entities")
+        query = """
+        UNWIND $entities as name
+        MATCH (e:Memory { name: name })
+        DETACH DELETE e
+        """
+        
+        await self.driver.execute_query(query, {"entities": entity_names}, routing_control=RoutingControl.WRITE)
+        logger.info(f"Successfully deleted {len(entity_names)} entities")
+
+    async def delete_observations(self, deletions: List[ObservationDeletion]) -> None:
+        """Delete specific observations from entities."""
+        logger.info(f"Deleting observations from {len(deletions)} entities")
+        query = """
+        UNWIND $deletions as d  
+        MATCH (e:Memory { name: d.entityName })
+        SET e.observations = [o in coalesce(e.observations,[]) WHERE NOT o IN d.observations]
+        """
+        await self.driver.execute_query(
+            query, 
+            {"deletions": [deletion.model_dump() for deletion in deletions]},
+            routing_control=RoutingControl.WRITE
+        )
+        logger.info(f"Successfully deleted observations from {len(deletions)} entities")
+
+    async def delete_relations(self, relations: List[Relation]) -> None:
+        """Delete multiple relations from the graph."""
+        logger.info(f"Deleting {len(relations)} relations")
+        for relation in relations:
+            query = f"""
+            WITH $relation as relation
+            MATCH (source:Memory)-[r:{relation.relationType}]->(target:Memory)
+            WHERE source.name = relation.source
+            AND target.name = relation.target
+            DELETE r
+            """
+            await self.driver.execute_query(
+                query, 
+                {"relation": relation.model_dump()},
+                routing_control=RoutingControl.WRITE
+            )
+        logger.info(f"Successfully deleted {len(relations)} relations")
+
+    async def read_graph(self) -> KnowledgeGraph:
+        """Read the entire knowledge graph."""
+        return await self.load_graph()
+
+    async def search_memories(self, query: str) -> KnowledgeGraph:
+        """Search for memories based on a query with Fulltext Search."""
+        logger.info(f"Searching for memories with query: '{query}'")
+        return await self.load_graph(query)
+
+    async def find_memories_by_name(self, names: List[str]) -> KnowledgeGraph:
+        """Find specific memories by their names. This does not use fulltext search."""
+        logger.info(f"Finding {len(names)} memories by name")
+        query = """
+        MATCH (e:Memory)
+        WHERE e.name IN $names
+        RETURN  e.name as name, 
+                e.type as type, 
+                e.observations as observations
+        """
+        result_nodes = await self.driver.execute_query(query, {"names": names}, routing_control=RoutingControl.READ)
+        entities: list[Entity] = list()
+        for record in result_nodes.records:
+            entities.append(Entity(
+                name=record['name'],
+                type=record['type'],
+                observations=record.get('observations', list())
+            ))
+        
+        # Get relations for found entities
+        relations: list[Relation] = list()
+        if entities:
+            query = """
+            MATCH (source:Memory)-[r]->(target:Memory)
+            WHERE source.name IN $names OR target.name IN $names
+            RETURN  source.name as source, 
+                    target.name as target, 
+                    type(r) as relationType
+            """
+            result_relations = await self.driver.execute_query(query, {"names": names}, routing_control=RoutingControl.READ)
+            for record in result_relations.records:
+                relations.append(Relation(
+                    source=record["source"],
+                    target=record["target"],
+                    relationType=record["relationType"]
+                ))
+        
+        logger.info(f"Found {len(entities)} entities and {len(relations)} relations")
+        return KnowledgeGraph(entities=entities, relations=relations)