Skip to content

Commit

Permalink
remove chunk_entity
Browse files Browse the repository at this point in the history
  • Loading branch information
NolanTrem committed Nov 30, 2024
1 parent 7620d9f commit 6d395a1
Show file tree
Hide file tree
Showing 2 changed files with 0 additions and 68 deletions.
8 changes: 0 additions & 8 deletions py/core/main/services/kg_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -566,14 +566,6 @@ async def kg_entity_description(

return all_results

@telemetry_event("get_graph_status")
async def get_graph_status(
self,
collection_id: UUID,
**kwargs,
):
raise NotImplementedError("Not implemented")

@telemetry_event("kg_clustering")
async def kg_clustering(
self,
Expand Down
60 changes: 0 additions & 60 deletions py/core/providers/database/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -2199,7 +2199,6 @@ async def delete_node_via_document_id(

# Execute separate DELETE queries
delete_queries = [
f"DELETE FROM {self._get_table_name('chunk_entity')} WHERE document_id = $1",
f"DELETE FROM {self._get_table_name('relationship')} WHERE document_id = $1",
f"DELETE FROM {self._get_table_name('entity')} WHERE document_id = $1",
]
Expand Down Expand Up @@ -2716,7 +2715,6 @@ async def delete_graph_for_collection(
# TODO: make these queries more efficient. Pass the document_ids as params.
if cascade:
DELETE_QUERIES += [
f"DELETE FROM {self._get_table_name('chunk_entity')} WHERE document_id = ANY($1::uuid[]);",
f"DELETE FROM {self._get_table_name('relationship')} WHERE document_id = ANY($1::uuid[]);",
f"DELETE FROM {self._get_table_name('entity')} WHERE document_id = ANY($1::uuid[]);",
f"DELETE FROM {self._get_table_name('graph_entity')} WHERE collection_id = $1;",
Expand Down Expand Up @@ -2905,64 +2903,6 @@ async def get_entity_map(

return entity_map

async def get_graph_status(self, collection_id: UUID) -> dict:
# check document_info table for the documents in the collection and return the status of each document
kg_extraction_statuses = await self.connection_manager.fetch_query(
f"SELECT document_id, extraction_status FROM {self._get_table_name('document_info')} WHERE collection_id = $1",
[collection_id],
)

document_ids = [
doc_id["document_id"] for doc_id in kg_extraction_statuses
]

graph_cluster_statuses = await self.connection_manager.fetch_query(
f"SELECT enrichment_status FROM {self._get_table_name(PostgresCollectionHandler.TABLE_NAME)} WHERE id = $1",
[collection_id],
)

# entity and relationship counts
chunk_entity_count = await self.connection_manager.fetch_query(
f"SELECT COUNT(*) FROM {self._get_table_name('chunk_entity')} WHERE document_id = ANY($1)",
[document_ids],
)

relationship_count = await self.connection_manager.fetch_query(
f"SELECT COUNT(*) FROM {self._get_table_name('relationship')} WHERE document_id = ANY($1)",
[document_ids],
)

entity_count = await self.connection_manager.fetch_query(
f"SELECT COUNT(*) FROM {self._get_table_name('entity')} WHERE document_id = ANY($1)",
[document_ids],
)

graph_entity_count = await self.connection_manager.fetch_query(
f"SELECT COUNT(*) FROM {self._get_table_name('graph_entity')} WHERE collection_id = $1",
[collection_id],
)

community_count = await self.connection_manager.fetch_query(
f"SELECT COUNT(*) FROM {self._get_table_name('community')} WHERE collection_id = $1",
[collection_id],
)

return {
"kg_extraction_statuses": kg_extraction_statuses,
"graph_cluster_status": graph_cluster_statuses[0][
"enrichment_status"
],
"chunk_entity_count": chunk_entity_count[0]["count"],
"relationship_count": relationship_count[0]["count"],
"entity_count": entity_count[0]["count"],
"graph_entity_count": graph_entity_count[0]["count"],
"community_count": community_count[0]["count"],
}

####################### ESTIMATION METHODS #######################

####################### GRAPH SEARCH METHODS #######################

def _build_filters(
self, filters: dict, parameters: list[Union[str, int, bytes]]
) -> str:
Expand Down

0 comments on commit 6d395a1

Please sign in to comment.