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=[])
|
return NodeSearchResponse(message='No relevant nodes found', nodes=[])
|
||||||
|
|
||||||
# Format the results
|
# Format the results
|
||||||
node_results = [
|
node_results = []
|
||||||
NodeResult(
|
for node in nodes:
|
||||||
uuid=node.uuid,
|
# Get attributes and ensure no embeddings are included
|
||||||
name=node.name,
|
attrs = node.attributes if hasattr(node, 'attributes') else {}
|
||||||
type=node.labels[0] if node.labels else 'Unknown',
|
# Remove any embedding keys that might be in attributes
|
||||||
created_at=node.created_at.isoformat() if node.created_at else None,
|
attrs = {k: v for k, v in attrs.items() if 'embedding' not in k.lower()}
|
||||||
summary=node.summary,
|
|
||||||
|
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)
|
return NodeSearchResponse(message='Nodes retrieved successfully', nodes=node_results)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|
|
||||||
|
|
@ -14,9 +14,11 @@ class SuccessResponse(TypedDict):
|
||||||
class NodeResult(TypedDict):
|
class NodeResult(TypedDict):
|
||||||
uuid: str
|
uuid: str
|
||||||
name: str
|
name: str
|
||||||
type: str
|
labels: list[str]
|
||||||
created_at: str | None
|
created_at: str | None
|
||||||
summary: str | None
|
summary: str | None
|
||||||
|
group_id: str
|
||||||
|
attributes: dict[str, Any]
|
||||||
|
|
||||||
|
|
||||||
class NodeSearchResponse(TypedDict):
|
class NodeSearchResponse(TypedDict):
|
||||||
|
|
|
||||||
|
|
@ -3,6 +3,30 @@
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from graphiti_core.edges import EntityEdge
|
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]:
|
def format_fact_result(edge: EntityEdge) -> dict[str, Any]:
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue