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.
This commit is contained in:
Daniel Chalef 2025-10-30 22:40:29 -07:00
parent dbd22fa527
commit de43038f5e
3 changed files with 44 additions and 10 deletions

View file

@ -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:

View file

@ -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):

View file

@ -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]: