From de43038f5e9339613f5dd2a5594f22bfcb85a67e Mon Sep 17 00:00:00 2001
From: Daniel Chalef <131175+danielchalef@users.noreply.github.com>
Date: Thu, 30 Oct 2025 22:40:29 -0700
Subject: [PATCH] feat: Return complete node properties and exclude all
 embeddings

Enhanced node search results to include all relevant properties:
- Added `attributes` dict for custom entity properties
- Changed from single `type` to full `labels` array
- Added `group_id` for partition information
- Added safety filter to strip any keys containing "embedding" from attributes

Added format_node_result() helper function for consistent node formatting
that excludes name_embedding vectors, matching the pattern used for edges.

Embeddings are now explicitly excluded in all data returns:
- EntityNode: name_embedding excluded + attributes filtered
- EntityEdge: fact_embedding excluded (existing)
- EpisodicNode: No embeddings to exclude

This ensures clients receive complete metadata while keeping payload
sizes manageable and avoiding exposure of internal vector representations.
---
 mcp_server/src/graphiti_mcp_server.py   | 26 ++++++++++++++++---------
 mcp_server/src/models/response_types.py |  4 +++-
 mcp_server/src/utils/formatting.py      | 24 +++++++++++++++++++++++
 3 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/mcp_server/src/graphiti_mcp_server.py b/mcp_server/src/graphiti_mcp_server.py
index 1fcb746f..0c9a568a 100644
--- a/mcp_server/src/graphiti_mcp_server.py
+++ b/mcp_server/src/graphiti_mcp_server.py
@@ -458,16 +458,24 @@ async def search_nodes(
             return NodeSearchResponse(message='No relevant nodes found', nodes=[])
 
         # Format the results
-        node_results = [
-            NodeResult(
-                uuid=node.uuid,
-                name=node.name,
-                type=node.labels[0] if node.labels else 'Unknown',
-                created_at=node.created_at.isoformat() if node.created_at else None,
-                summary=node.summary,
+        node_results = []
+        for node in nodes:
+            # Get attributes and ensure no embeddings are included
+            attrs = node.attributes if hasattr(node, 'attributes') else {}
+            # Remove any embedding keys that might be in attributes
+            attrs = {k: v for k, v in attrs.items() if 'embedding' not in k.lower()}
+
+            node_results.append(
+                NodeResult(
+                    uuid=node.uuid,
+                    name=node.name,
+                    labels=node.labels if node.labels else [],
+                    created_at=node.created_at.isoformat() if node.created_at else None,
+                    summary=node.summary,
+                    group_id=node.group_id,
+                    attributes=attrs,
+                )
             )
-            for node in nodes
-        ]
 
         return NodeSearchResponse(message='Nodes retrieved successfully', nodes=node_results)
     except Exception as e:
diff --git a/mcp_server/src/models/response_types.py b/mcp_server/src/models/response_types.py
index 81032afc..eca20324 100644
--- a/mcp_server/src/models/response_types.py
+++ b/mcp_server/src/models/response_types.py
@@ -14,9 +14,11 @@ class SuccessResponse(TypedDict):
 class NodeResult(TypedDict):
     uuid: str
     name: str
-    type: str
+    labels: list[str]
     created_at: str | None
     summary: str | None
+    group_id: str
+    attributes: dict[str, Any]
 
 
 class NodeSearchResponse(TypedDict):
diff --git a/mcp_server/src/utils/formatting.py b/mcp_server/src/utils/formatting.py
index 53ff937e..595a6116 100644
--- a/mcp_server/src/utils/formatting.py
+++ b/mcp_server/src/utils/formatting.py
@@ -3,6 +3,30 @@
 from typing import Any
 
 from graphiti_core.edges import EntityEdge
+from graphiti_core.nodes import EntityNode
+
+
+def format_node_result(node: EntityNode) -> dict[str, Any]:
+    """Format an entity node into a readable result.
+
+    Since EntityNode is a Pydantic BaseModel, we can use its built-in serialization capabilities.
+    Excludes embedding vectors to reduce payload size and avoid exposing internal representations.
+
+    Args:
+        node: The EntityNode to format
+
+    Returns:
+        A dictionary representation of the node with serialized dates and excluded embeddings
+    """
+    result = node.model_dump(
+        mode='json',
+        exclude={
+            'name_embedding',
+        },
+    )
+    # Remove any embedding that might be in attributes
+    result.get('attributes', {}).pop('name_embedding', None)
+    return result
 
 
 def format_fact_result(edge: EntityEdge) -> dict[str, Any]: