test and bugfix

Author: Orbital-Web

backend/onyx/db/document.py

@@ -937,9 +937,14 @@ def get_unprocessed_kg_document_batch_for_connector(
         .where(
             and_(
                 DocumentByConnectorCredentialPair.connector_id == connector_id,
-                DbDocument.doc_updated_at >= kg_coverage_start,
-                DbDocument.doc_updated_at
-                >= datetime.now() - timedelta(days=kg_max_coverage_days),
+                or_(
+                    DbDocument.doc_updated_at.is_(None),
+                    DbDocument.doc_updated_at
+                    >= max(
+                        kg_coverage_start,
+                        datetime.now() - timedelta(days=kg_max_coverage_days),
+                    ),
+                ),
                 or_(
                     DbDocument.kg_stage.is_(None),
                     DbDocument.kg_stage == KGStage.NOT_STARTED,
@@ -1185,14 +1190,17 @@ def check_for_documents_needing_kg_processing(
         .where(
             and_(
                 Connector.kg_processing_enabled.is_(True),
-                DbDocument.doc_updated_at >= kg_coverage_start,
-                DbDocument.doc_updated_at
-                >= datetime.now() - timedelta(days=kg_max_coverage_days),
                 or_(
-                    or_(
-                        DbDocument.kg_stage.is_(None),
-                        DbDocument.kg_stage == KGStage.NOT_STARTED,
+                    DbDocument.doc_updated_at.is_(None),
+                    DbDocument.doc_updated_at
+                    >= max(
+                        kg_coverage_start,
+                        datetime.now() - timedelta(days=kg_max_coverage_days),
                     ),
+                ),
+                or_(
+                    DbDocument.kg_stage.is_(None),
+                    DbDocument.kg_stage == KGStage.NOT_STARTED,
                     DbDocument.doc_updated_at > DbDocument.kg_processing_time,
                 ),
             )

backend/onyx/db/entities.py

@@ -130,6 +130,8 @@ def transfer_entity(
             set_=dict(
                 occurrences=KGEntity.occurrences + entity.occurrences,
                 attributes=entity.attributes,  # attribute can get updated after re-indexing
+                entity_key=entity.entity_key,
+                parent_key=entity.parent_key,
                 event_time=entity.event_time,
                 time_updated=datetime.now(),
             ),
@@ -196,6 +198,9 @@ def merge_entities(
             document_id=document_id,
             alternative_names=list(alternative_names),
             occurrences=parent.occurrences + child.occurrences,
+            attributes=parent.attributes | child.attributes,
+            entity_key=child.entity_key,
+            parent_key=child.parent_key,
         )
         .returning(KGEntity)
     )

backend/onyx/kg/utils/extraction_utils.py

@@ -451,16 +451,16 @@ def track_metadata(
                 )
                 self.entity_attr_info[entity_type][attribute] = trackinfo
 
-            # if we see to many different values, we stop tracking
-            if (
-                trackinfo.values is None
-                or len(trackinfo.values) > KG_METADATA_TRACKING_THRESHOLD
-            ):
-                trackinfo.values = None
+            # None means marked as don't track
+            if trackinfo.values is None:
                 continue
 
             # track the value
             if isinstance(value, str):
                 trackinfo.values.add(value)
             else:
                 trackinfo.values.update(value)
+
+            # if we see to many different values, we stop tracking
+            if len(trackinfo.values) > KG_METADATA_TRACKING_THRESHOLD:
+                trackinfo.values = None

backend/tests/integration/common_utils/managers/document.py

@@ -66,8 +66,14 @@ def _generate_dummy_document(
     document_id: str,
     cc_pair_id: int,
     content: str | None = None,
+    extra_metadata: dict | None = None,
 ) -> dict:
     text = content if content else f"This is test document {document_id}"
+
+    metadata: dict = {"document_id": document_id}
+    if extra_metadata:
+        metadata.update(extra_metadata)
+
     return {
         "document": {
             "id": document_id,
@@ -78,8 +84,7 @@ def _generate_dummy_document(
                 }
             ],
             "source": DocumentSource.NOT_APPLICABLE,
-            # just for testing metadata
-            "metadata": {"document_id": document_id},
+            "metadata": metadata,
             "semantic_identifier": f"Test Document {document_id}",
             "from_ingestion_api": True,
         },
@@ -128,12 +133,18 @@ def seed_doc_with_content(
         content: str,
         document_id: str | None = None,
         api_key: DATestAPIKey | None = None,
+        metadata: dict | None = None,
     ) -> SimpleTestDocument:
         # Use provided document_ids if available, otherwise generate random UUIDs
         if document_id is None:
             document_id = f"test-doc-{uuid4()}"
         # Create and ingest some documents
-        document: dict = _generate_dummy_document(document_id, cc_pair.id, content)
+        document: dict = _generate_dummy_document(
+            document_id,
+            cc_pair.id,
+            content,
+            extra_metadata=metadata,
+        )
         response = requests.post(
             f"{API_SERVER_URL}/onyx-api/ingestion",
             json=document,