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:
parent
dbd22fa527
commit
de43038f5e
3 changed files with 44 additions and 10 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue