Fix entity/relation creation endpoints to properly update vector stores

- Changed create_entity to use rag.acreate_entity() instead of direct graph manipulation
  - Changed create_relation to use rag.acreate_relation() instead of direct graph manipulation
  - This ensures vector embeddings are created and entities/relations are searchable
  - Adds proper concurrency locks and metadata population
This commit is contained in:
NeelM0906 2025-10-09 17:02:17 -04:00
parent f6d1fb98ac
commit b7c77396a0

View file

@ -264,24 +264,20 @@ def create_graph_routes(rag, api_key: Optional[str] = None):
} }
""" """
try: try:
# Check if entity already exists # Use the proper acreate_entity method which handles:
exists = await rag.chunk_entity_relation_graph.has_node(request.entity_name) # - Graph lock for concurrency
if exists: # - Vector embedding creation in entities_vdb
raise ValueError(f"Entity '{request.entity_name}' already exists") # - Metadata population and defaults
# - Index consistency via _edit_entity_done
# Prepare entity data result = await rag.acreate_entity(
entity_data = request.entity_data.copy() entity_name=request.entity_name,
entity_data["entity_id"] = request.entity_name entity_data=request.entity_data,
# Create the entity
await rag.chunk_entity_relation_graph.upsert_node(
request.entity_name, entity_data
) )
return { return {
"status": "success", "status": "success",
"message": f"Entity '{request.entity_name}' created successfully", "message": f"Entity '{request.entity_name}' created successfully",
"data": entity_data, "data": result,
} }
except ValueError as ve: except ValueError as ve:
logger.error( logger.error(
@ -321,36 +317,22 @@ def create_graph_routes(rag, api_key: Optional[str] = None):
} }
""" """
try: try:
# Check if both entities exist # Use the proper acreate_relation method which handles:
source_exists = await rag.chunk_entity_relation_graph.has_node( # - Graph lock for concurrency
request.source_entity # - Entity existence validation
) # - Duplicate relation checks
target_exists = await rag.chunk_entity_relation_graph.has_node( # - Vector embedding creation in relationships_vdb
request.target_entity # - Index consistency via _edit_relation_done
) result = await rag.acreate_relation(
source_entity=request.source_entity,
if not source_exists: target_entity=request.target_entity,
raise ValueError( relation_data=request.relation_data,
f"Source entity '{request.source_entity}' does not exist"
)
if not target_exists:
raise ValueError(
f"Target entity '{request.target_entity}' does not exist"
)
# Create the relationship
await rag.chunk_entity_relation_graph.upsert_edge(
request.source_entity, request.target_entity, request.relation_data
) )
return { return {
"status": "success", "status": "success",
"message": f"Relation created successfully between '{request.source_entity}' and '{request.target_entity}'", "message": f"Relation created successfully between '{request.source_entity}' and '{request.target_entity}'",
"data": { "data": result,
"source": request.source_entity,
"target": request.target_entity,
**request.relation_data,
},
} }
except ValueError as ve: except ValueError as ve:
logger.error( logger.error(