Kg batch clustering

backend/alembic/versions/cec7ec36c505_kgentity_parent.py

@@ -0,0 +1,82 @@
+"""kgentity_parent
+
+Revision ID: cec7ec36c505
+Revises: 495cb26ce93e
+Create Date: 2025-06-07 20:07:46.400770
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy import text
+
+
+# revision identifiers, used by Alembic.
+revision = "cec7ec36c505"
+down_revision = "495cb26ce93e"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "kg_entity",
+        sa.Column("parent_key", sa.String(), nullable=True, index=True),
+    )
+
+    # Set parent key to the immediate parent entity (if any)
+
+    # Each entity e could have at most D=KG_MAX_PARENT_RECURSION_DEPTH ancestors (default=2)
+    # The ancestors (i.e., parent candidates) can be found through the relationship pi__has_subcomponent__e
+    # If p is an immediate parent of e, then there is no p__has_subcomponent_pi for all pi
+    # If p is not an immediate parent of e, then there exists a pi such that p__has_subcomponent__pi
+    # We can thus find the immediate parent for N many entities in O(N*D^2) time and O(D) space
+
+    # Note: this won't work if the tenant increased KG_MAX_PARENT_RECURSION_DEPTH after extraction
+    # In this case, resetting the kg data before migration and re-extracting later is recommended
+    op.execute(
+        text(
+            """
+            UPDATE kg_entity
+            SET parent_key = (
+                -- 1. find all ancestors (parent candidates) pi of kg_entity
+                SELECT pi.entity_key
+                FROM kg_relationship r
+                JOIN kg_entity pi ON r.source_node = pi.id_name
+                WHERE
+                    r.target_node = kg_entity.id_name
+                    AND r.type = 'has_subcomponent'
+                    -- 2. exclude any parent candidate that is a parent of another candidate
+                    AND NOT EXISTS (
+                        SELECT 1
+                        FROM kg_relationship r2
+                        WHERE
+                            r2.source_node = r.source_node
+                            AND r2.type = 'has_subcomponent'
+                            AND r2.target_node IN (
+                                -- 3. get the other parent candidates
+                                SELECT r3.source_node
+                                FROM kg_relationship r3
+                                WHERE
+                                    r3.target_node = kg_entity.id_name
+                                    AND r3.type = 'has_subcomponent'
+                                    AND r3.source_node != r.source_node
+                            )
+                    )
+                LIMIT 1
+            )
+            -- only bother setting parent_key for entities that have a parent
+            WHERE EXISTS (
+                SELECT 1
+                FROM kg_relationship
+                WHERE
+                    target_node = kg_entity.id_name
+                    AND type = 'has_subcomponent'
+            );
+            """
+        )
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("kg_entity", "parent_key")

backend/onyx/db/entities.py

@@ -130,6 +130,7 @@ def transfer_entity(
             entity_class=entity.entity_class,
             entity_subtype=entity.entity_subtype,
             entity_key=entity.entity_key,
+            parent_key=entity.parent_key,
             alternative_names=entity.alternative_names or [],
             entity_type_id_name=entity.entity_type_id_name,
             document_id=entity.document_id,

backend/onyx/db/models.py

@@ -849,6 +849,9 @@ class KGEntity(Base):
     entity_key: Mapped[str] = mapped_column(
         NullFilteredString, nullable=True, index=True
     )
+    parent_key: Mapped[str | None] = mapped_column(
+        NullFilteredString, nullable=True, index=True
+    )
     entity_subtype: Mapped[str] = mapped_column(
         NullFilteredString, nullable=True, index=True
     )
@@ -1003,7 +1006,7 @@ class KGEntityExtractionStaging(Base):
     )
 
     # Basic entity information
-    parent_key: Mapped[str] = mapped_column(
+    parent_key: Mapped[str | None] = mapped_column(
         NullFilteredString, nullable=True, index=True
     )
 

backend/onyx/db/relationships.py

@@ -27,7 +27,7 @@
 def upsert_staging_relationship(
     db_session: Session,
     relationship_id_name: str,
-    source_document_id: str,
+    source_document_id: str | None,
     occurrences: int = 1,
 ) -> KGRelationshipExtractionStaging:
     """
@@ -99,6 +99,72 @@ def upsert_staging_relationship(
     return result
 
 
+def upsert_relationship(
+    db_session: Session,
+    relationship_id_name: str,
+    source_document_id: str | None,
+    occurrences: int = 1,
+) -> KGRelationship:
+    """
+    Upsert a new relationship directly to the database.
+
+    Args:
+        db_session: SQLAlchemy database session
+        relationship_id_name: The ID name of the relationship in format "source__relationship__target"
+        source_document_id: ID of the source document
+        occurrences: Number of times this relationship has been found
+    Returns:
+        The created or updated KGRelationship object
+
+    Raises:
+        sqlalchemy.exc.IntegrityError: If there's an error with the database operation
+    """
+    # Generate a unique ID for the relationship
+    relationship_id_name = format_relationship_id(relationship_id_name)
+    (
+        source_entity_id_name,
+        relationship_string,
+        target_entity_id_name,
+    ) = split_relationship_id(relationship_id_name)
+
+    source_entity_type = get_entity_type(source_entity_id_name)
+    target_entity_type = get_entity_type(target_entity_id_name)
+    relationship_type = extract_relationship_type_id(relationship_id_name)
+
+    # Insert the new relationship
+    stmt = (
+        postgresql.insert(KGRelationship)
+        .values(
+            {
+                "id_name": relationship_id_name,
+                "source_node": source_entity_id_name,
+                "target_node": target_entity_id_name,
+                "source_node_type": source_entity_type,
+                "target_node_type": target_entity_type,
+                "type": relationship_string.lower(),
+                "relationship_type_id_name": relationship_type,
+                "source_document": source_document_id,
+                "occurrences": occurrences,
+            }
+        )
+        .on_conflict_do_update(
+            index_elements=["id_name", "source_document"],
+            set_=dict(
+                occurrences=KGRelationship.occurrences + occurrences,
+            ),
+        )
+        .returning(KGRelationship)
+    )
+
+    new_relationship = db_session.execute(stmt).scalar()
+    if new_relationship is None:
+        raise RuntimeError(
+            f"Failed to upsert relationship with id_name: {relationship_id_name}"
+        )
+    db_session.flush()
+    return new_relationship
+
+
 def transfer_relationship(
     db_session: Session,
     relationship: KGRelationshipExtractionStaging,
@@ -218,6 +284,65 @@ def upsert_staging_relationship_type(
     return result
 
 
+def upsert_relationship_type(
+    db_session: Session,
+    source_entity_type: str,
+    relationship_type: str,
+    target_entity_type: str,
+    definition: bool = False,
+    extraction_count: int = 1,
+) -> KGRelationshipType:
+    """
+    Upsert a new relationship type directly to the database.
+
+    Args:
+        db_session: SQLAlchemy session
+        source_entity_type: Type of the source entity
+        relationship_type: Type of relationship
+        target_entity_type: Type of the target entity
+        definition: Whether this relationship type represents a definition (default False)
+
+    Returns:
+        The created KGRelationshipType object
+    """
+
+    id_name = make_relationship_type_id(
+        source_entity_type, relationship_type, target_entity_type
+    )
+
+    # Create new relationship type
+    stmt = (
+        postgresql.insert(KGRelationshipType)
+        .values(
+            {
+                "id_name": id_name,
+                "name": relationship_type,
+                "source_entity_type_id_name": source_entity_type.upper(),
+                "target_entity_type_id_name": target_entity_type.upper(),
+                "definition": definition,
+                "occurrences": extraction_count,
+                "type": relationship_type,  # Using the relationship_type as the type
+                "active": True,  # Setting as active by default
+            }
+        )
+        .on_conflict_do_update(
+            index_elements=["id_name"],
+            set_=dict(
+                occurrences=KGRelationshipType.occurrences + extraction_count,
+            ),
+        )
+        .returning(KGRelationshipType)
+    )
+
+    new_relationship_type = db_session.execute(stmt).scalar()
+    if new_relationship_type is None:
+        raise RuntimeError(
+            f"Failed to upsert relationship type with id_name: {id_name}"
+        )
+    db_session.flush()
+    return new_relationship_type
+
+
 def transfer_relationship_type(
     db_session: Session,
     relationship_type: KGRelationshipTypeExtractionStaging,
@@ -262,112 +387,6 @@ def transfer_relationship_type(
     return new_relationship_type
 
 
-def get_parent_child_relationships_and_types(
-    db_session: Session,
-    depth: int,
-) -> tuple[
-    list[KGRelationshipExtractionStaging], list[KGRelationshipTypeExtractionStaging]
-]:
-    """
-    Create parent-child relationships and relationship types from staging entities with
-    a parent key, if the parent exists in the normalized entities table. Will create
-    relationships up to depth levels. E.g., if depth is 2, a relationship will be created
-    between the entity and its parent, and the entity and its grandparents (if any).
-    A relationship will not be created if the parent does not exist.
-    """
-    relationship_types: dict[str, KGRelationshipTypeExtractionStaging] = {}
-    relationships: dict[tuple[str, str | None], KGRelationshipExtractionStaging] = {}
-
-    parented_entities = (
-        db_session.query(KGEntityExtractionStaging)
-        .filter(KGEntityExtractionStaging.parent_key.isnot(None))
-        .all()
-    )
-
-    # create has_subcomponent relationships and relationship types
-    for entity in parented_entities:
-        child = entity
-        if entity.transferred_id_name is None:
-            logger.warning(f"Entity {entity.id_name} has not yet been transferred")
-            continue
-
-        for i in range(depth):
-            if not child.parent_key:
-                break
-
-            # find the transferred parent entity
-            parent = (
-                db_session.query(KGEntity)
-                .filter(
-                    KGEntity.entity_class == child.entity_class,
-                    KGEntity.entity_key == child.parent_key,
-                )
-                .first()
-            )
-            if parent is None:
-                logger.warning(f"Parent entity not found for {entity.id_name}")
-                break
-
-            # create the relationship type
-            relationship_type = upsert_staging_relationship_type(
-                db_session=db_session,
-                source_entity_type=parent.entity_type_id_name,
-                relationship_type="has_subcomponent",
-                target_entity_type=entity.entity_type_id_name,
-                definition=False,
-                extraction_count=1,
-            )
-            relationship_types[relationship_type.id_name] = relationship_type
-
-            # create the relationship
-            # (don't add it to the table as we're using the transferred id, which breaks fk constraints)
-            relationship_id_name = make_relationship_id(
-                parent.id_name, "has_subcomponent", entity.transferred_id_name
-            )
-            if (parent.id_name, entity.document_id) not in relationships:
-                (
-                    source_entity_id_name,
-                    relationship_string,
-                    target_entity_id_name,
-                ) = split_relationship_id(relationship_id_name)
-
-                source_entity_type = get_entity_type(source_entity_id_name)
-                target_entity_type = get_entity_type(target_entity_id_name)
-                relationship_type_id_name = extract_relationship_type_id(
-                    relationship_id_name
-                )
-                relationships[(relationship_id_name, entity.document_id)] = (
-                    KGRelationshipExtractionStaging(
-                        id_name=relationship_id_name,
-                        source_node=source_entity_id_name,
-                        target_node=target_entity_id_name,
-                        source_node_type=source_entity_type,
-                        target_node_type=target_entity_type,
-                        type=relationship_string,
-                        relationship_type_id_name=relationship_type_id_name,
-                        source_document=entity.document_id,
-                        occurrences=1,
-                    )
-                )
-            else:
-                relationships[(parent.id_name, entity.document_id)].occurrences += 1
-
-            # set parent as the next child (unless we're at the max depth)
-            if i < depth - 1:
-                parent_staging = (
-                    db_session.query(KGEntityExtractionStaging)
-                    .filter(
-                        KGEntityExtractionStaging.transferred_id_name == parent.id_name
-                    )
-                    .first()
-                )
-                if parent_staging is None:
-                    break
-                child = parent_staging
-
-    return list(relationships.values()), list(relationship_types.values())
-
-
 def delete_relationships_by_id_names(
     db_session: Session, id_names: list[str], kg_stage: KGStage
 ) -> int: