diff --git a/backend/onyx/db/entities.py b/backend/onyx/db/entities.py index 9696ce9267..ade11bfd51 100644 --- a/backend/onyx/db/entities.py +++ b/backend/onyx/db/entities.py @@ -12,6 +12,7 @@ from sqlalchemy.orm import Session import onyx.db.document as dbdocument +from onyx.db.entity_type import UNGROUNDED_SOURCE_NAME from onyx.db.models import Document from onyx.db.models import KGEntity from onyx.db.models import KGEntityExtractionStaging @@ -328,7 +329,13 @@ def get_entity_stats_by_grounded_source_name( .group_by(KGEntityType.grounded_source_name) .all() ) + + # `row.grounded_source_name` is NULLABLE in the database schema. + # Thus, for all "ungrounded" entity-types, we use a default name. return { - row.grounded_source_name: (row.last_updated, row.entities_count) + (row.grounded_source_name or UNGROUNDED_SOURCE_NAME): ( + row.last_updated, + row.entities_count, + ) for row in results } diff --git a/backend/onyx/db/entity_type.py b/backend/onyx/db/entity_type.py index 54b7bfaff3..a88ac9c2ac 100644 --- a/backend/onyx/db/entity_type.py +++ b/backend/onyx/db/entity_type.py @@ -11,7 +11,7 @@ from onyx.server.kg.models import EntityType -_UNGROUNDED_SOURCE_NAME = "Ungrounded" +UNGROUNDED_SOURCE_NAME = "Ungrounded" def get_entity_types_with_grounded_source_name( @@ -87,7 +87,7 @@ def get_configured_entity_types(db_session: Session) -> dict[str, list[KGEntityT et_map = defaultdict(list) for et in ets: - key = et.grounded_source_name or _UNGROUNDED_SOURCE_NAME + key = et.grounded_source_name or UNGROUNDED_SOURCE_NAME et_map[key].append(et) return et_map diff --git a/backend/onyx/server/kg/api.py b/backend/onyx/server/kg/api.py index 5608199742..8d15e2c24d 100644 --- a/backend/onyx/server/kg/api.py +++ b/backend/onyx/server/kg/api.py @@ -179,15 +179,19 @@ def get_kg_entity_types( ) -> SourceAndEntityTypeView: # when using for the first time, populate with default entity types entity_types = { - key: [EntityType.from_model(et) for et in ets] - for key, ets in get_configured_entity_types(db_session=db_session).items() + source_name: [EntityType.from_model(et) for et in ets] + for source_name, ets in get_configured_entity_types( + db_session=db_session + ).items() } source_statistics = { - key: SourceStatistics( - source_name=key, last_updated=last_updated, entities_count=entities_count + source_name: SourceStatistics( + source_name=source_name, + last_updated=last_updated, + entities_count=entities_count, ) - for key, ( + for source_name, ( last_updated, entities_count, ) in get_entity_stats_by_grounded_source_name(db_session=db_session).items()