From e757996a7c7cb5354681ae0e352449ba51ffbcab Mon Sep 17 00:00:00 2001 From: Paco Gomez Date: Tue, 9 Jun 2026 13:59:54 +0200 Subject: [PATCH 1/2] Fix malformed entity lookup retry loop. Skip malformed Knowledge Graph entity lookup results during search post-processing so searches do not retry the same batch indefinitely. --- CHANGELOG.md | 6 +++ pyproject.toml | 2 +- .../search/search_utils.py | 18 +++++++-- tests/test_search/test_search_utils.py | 38 +++++++++++++++++++ uv.lock | 4 +- 5 files changed, 62 insertions(+), 6 deletions(-) create mode 100644 tests/test_search/test_search_utils.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d82bf59..d78e282 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,12 @@ file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [1.1.2] - 2026-06-09 + +## Fixed +- Skip malformed Knowledge Graph entity lookup results during search post-processing instead + of retrying the same batch indefinitely. + ## [1.1.1] - 2025-12-17 ## Fixed diff --git a/pyproject.toml b/pyproject.toml index 11c878a..0cbd64a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "bigdata-research-tools" -version = "1.1.1" +version = "1.1.2" description = "Bigdata.com API High-Efficiency Tools at Scale" readme = "README.md" authors = [{ name = "Bigdata.com", email = "support@ravenpack.com" }] diff --git a/src/bigdata_research_tools/search/search_utils.py b/src/bigdata_research_tools/search/search_utils.py index a538aa7..4e28ee5 100644 --- a/src/bigdata_research_tools/search/search_utils.py +++ b/src/bigdata_research_tools/search/search_utils.py @@ -53,6 +53,20 @@ def _look_up_entities_binary_search( entities = [] non_entities = [] + def convert_sdk_entities(batch_lookup: list[object]) -> list[BigdataEntity]: + converted_entities = [] + for sdk_entity in batch_lookup: + if sdk_entity is None: + continue + try: + converted_entities.append(BigdataEntity.from_sdk(sdk_entity)) + except AssertionError as e: + logger.warning( + "Skipping malformed entity returned by Knowledge Graph lookup: %s", + e, + ) + return converted_entities + def depth_first_search(batch: list[str]) -> None: """ Recursively lookup entities in a depth-first search manner. @@ -68,9 +82,7 @@ def depth_first_search(batch: list[str]) -> None: try: batch_lookup = bigdata.knowledge_graph.get_entities(batch) - entities.extend( - [BigdataEntity.from_sdk(ent) for ent in batch_lookup if ent is not None] - ) + entities.extend(convert_sdk_entities(batch_lookup)) except ValidationError as e: non_entities_found = findall(non_entity_key_pattern, str(e)) non_entities.extend(non_entities_found) diff --git a/tests/test_search/test_search_utils.py b/tests/test_search/test_search_utils.py new file mode 100644 index 0000000..8d5bfe0 --- /dev/null +++ b/tests/test_search/test_search_utils.py @@ -0,0 +1,38 @@ +from dataclasses import dataclass + +from bigdata_research_tools.search import search_utils + + +@dataclass +class FakeSdkEntity: + id: str + name: str + entity_type: str = "COMP" + + +@dataclass +class MalformedSdkEntity: + name: str + entity_type: str = "COMP" + + +class FakeKnowledgeGraph: + def get_entities(self, ids: list[str]) -> list[FakeSdkEntity | MalformedSdkEntity]: + return [ + FakeSdkEntity(id=ids[0], name="Valid Company"), + MalformedSdkEntity(name="Malformed Company"), + ] + + +class FakeBigdata: + knowledge_graph = FakeKnowledgeGraph() + + +def test_lookup_entities_skips_malformed_sdk_entities(monkeypatch) -> None: + monkeypatch.setattr(search_utils, "bigdata_connection", lambda: FakeBigdata()) + + entities = search_utils._look_up_entities_binary_search(["ABC123", "DEF456"]) + + assert len(entities) == 1 + assert entities[0].id == "ABC123" + assert entities[0].name == "Valid Company" diff --git a/uv.lock b/uv.lock index 85bfaa2..f702b38 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10, <4.0" resolution-markers = [ "python_full_version >= '3.12'", @@ -216,7 +216,7 @@ wheels = [ [[package]] name = "bigdata-research-tools" -version = "1.1.1" +version = "1.1.2" source = { editable = "." } dependencies = [ { name = "bigdata-client" }, From cd538390ff1619f3f6286a528b9cc7a61d4906e5 Mon Sep 17 00:00:00 2001 From: Paco Gomez Date: Tue, 9 Jun 2026 14:06:13 +0200 Subject: [PATCH 2/2] linting error --- src/bigdata_research_tools/search/search_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/bigdata_research_tools/search/search_utils.py b/src/bigdata_research_tools/search/search_utils.py index 4e28ee5..dd7a325 100644 --- a/src/bigdata_research_tools/search/search_utils.py +++ b/src/bigdata_research_tools/search/search_utils.py @@ -7,6 +7,7 @@ from bigdata_client.connection import RequestMaxLimitExceeds from bigdata_client.document import Document from bigdata_client.models.document import DocumentChunk +from bigdata_client.models.entities import QueryComponentMixin from bigdata_client.query_type import QueryType from pydantic import ValidationError @@ -53,7 +54,9 @@ def _look_up_entities_binary_search( entities = [] non_entities = [] - def convert_sdk_entities(batch_lookup: list[object]) -> list[BigdataEntity]: + def convert_sdk_entities( + batch_lookup: list[QueryComponentMixin | None], + ) -> list[BigdataEntity]: converted_entities = [] for sdk_entity in batch_lookup: if sdk_entity is None: