diff --git a/src/vidxp/application_models.py b/src/vidxp/application_models.py index 395bde4..f029bfa 100644 --- a/src/vidxp/application_models.py +++ b/src/vidxp/application_models.py @@ -962,7 +962,7 @@ def to_prediction(self) -> dict[str, list[dict[str, Any]]]: class FusionProvenance(ApplicationModel): profile: Literal[FusionProfile.reciprocal_rank] = FusionProfile.reciprocal_rank rank_constant: int = Field(default=60, gt=0) - overlap_rule: Literal["connected_intervals"] = "connected_intervals" + overlap_rule: Literal["connected_intervals", "shared_overlap"] = "connected_intervals" requested_modalities: tuple[Identifier, ...] = () searched_modalities: tuple[Identifier, ...] = () diff --git a/src/vidxp/search_fusion.py b/src/vidxp/search_fusion.py index b7e4184..415eeec 100644 --- a/src/vidxp/search_fusion.py +++ b/src/vidxp/search_fusion.py @@ -33,7 +33,7 @@ def _query_id( return "fused:" + hashlib.sha256(identity.encode("utf-8")).hexdigest() -def _connected_components( +def _shared_overlap_components( hits: tuple[SearchHit, ...], ) -> list[list[SearchHit]]: ordered = sorted( @@ -50,17 +50,21 @@ def _connected_components( components: list[list[SearchHit]] = [] current: list[SearchHit] = [] current_media: str | None = None - current_end = 0.0 + current_overlap_end = 0.0 for hit in ordered: - if not current or hit.media_id != current_media or hit.start > current_end: + if ( + not current + or hit.media_id != current_media + or hit.start >= current_overlap_end + ): if current: components.append(current) current = [hit] current_media = hit.media_id - current_end = hit.end + current_overlap_end = hit.end else: current.append(hit) - current_end = max(current_end, hit.end) + current_overlap_end = min(current_overlap_end, hit.end) if current: components.append(current) return components @@ -137,7 +141,7 @@ def fuse_search_results( ordered_results = tuple(by_modality[modality] for modality in searched_modalities) flattened = tuple(hit for result in ordered_results for hit in result.hits) candidates = [] - for hits in _connected_components(flattened): + for hits in _shared_overlap_components(flattened): ordered_hits = tuple( sorted( hits, @@ -152,8 +156,8 @@ def fuse_search_results( { "score": _score(hits), "media_id": hits[0].media_id, - "start": min(hit.start for hit in hits), - "end": max(hit.end for hit in hits), + "start": max(hit.start for hit in hits), + "end": min(hit.end for hit in hits), "modalities": tuple(sorted({hit.modality for hit in hits})), "hits": ordered_hits, } @@ -192,6 +196,7 @@ def fuse_search_results( modalities=searched_modalities, moments=moments, fusion=FusionProvenance( + overlap_rule="shared_overlap", requested_modalities=requested_modalities, searched_modalities=searched_modalities, ), diff --git a/tests/test_search_fusion.py b/tests/test_search_fusion.py index 622e021..08a2c1b 100644 --- a/tests/test_search_fusion.py +++ b/tests/test_search_fusion.py @@ -1,6 +1,6 @@ import unittest -from vidxp.application_models import SearchHit, SearchResult +from vidxp.application_models import FusedSearchResult, SearchHit, SearchResult from vidxp.search_fusion import RRF_RANK_CONSTANT, fuse_search_results @@ -57,8 +57,8 @@ def test_rrf_counts_only_the_best_rank_per_modality_in_a_moment(self): moment = result.moments[0] self.assertAlmostEqual(moment.score, 2 / (RRF_RANK_CONSTANT + 1)) self.assertEqual(len(moment.hits), 3) - self.assertEqual(moment.start, 1) - self.assertEqual(moment.end, 4) + self.assertEqual(moment.start, 2.5) + self.assertEqual(moment.end, 3.0) def test_result_order_does_not_change_fusion_identity_or_output(self): scene = SearchResult( @@ -96,9 +96,7 @@ def test_rewritten_atomic_query_identity_changes_fused_identity(self): modality="scene", hits=(hit("scene", 1, 1, 2, "scene:1"),), ) - rewritten = original.model_copy( - update={"query_id": "scene:rewritten"} - ) + rewritten = original.model_copy(update={"query_id": "scene:rewritten"}) arguments = { "query": "Where is the taxi?", "requested_modalities": ("scene",), @@ -109,6 +107,103 @@ def test_rewritten_atomic_query_identity_changes_fused_identity(self): self.assertNotEqual(first.query_id, second.query_id) + def test_bridging_hit_does_not_merge_separate_moments(self): + hit_a = hit("scene", 1, 10.0, 12.0, "scene:a") + hit_b = hit("speech", 1, 11.0, 25.0, "speech:b") + hit_c = hit("scene", 2, 24.0, 26.0, "scene:c") + + scene = SearchResult( + query_id="scene:q", + query="car", + modality="scene", + hits=(hit_a, hit_c), + ) + speech = SearchResult( + query_id="speech:q", + query="car", + modality="speech", + hits=(hit_b,), + ) + + result = fuse_search_results( + query="car", + requested_modalities=("scene", "speech"), + results=(scene, speech), + ) + + self.assertEqual(len(result.moments), 2) + moment_1, moment_2 = result.moments + self.assertEqual(moment_1.start, 11.0) + self.assertEqual(moment_1.end, 12.0) + self.assertIn("scene:a", [h.source_id for h in moment_1.hits]) + self.assertEqual(moment_2.start, 24.0) + self.assertEqual(moment_2.end, 26.0) + self.assertIn("scene:c", [h.source_id for h in moment_2.hits]) + self.assertNotIn("scene:c", [h.source_id for h in moment_1.hits]) + + def test_touching_hits_remain_separate_and_provenance_round_trips(self): + scene = SearchResult( + query_id="scene:q", + query="car", + modality="scene", + hits=( + hit("scene", 1, 10.0, 11.0, "scene:a"), + hit("scene", 2, 11.0, 12.0, "scene:b"), + ), + ) + result = fuse_search_results( + query="car", requested_modalities=("scene",), results=(scene,) + ) + + self.assertEqual( + [(moment.start, moment.end) for moment in result.moments], + [(10.0, 11.0), (11.0, 12.0)], + ) + self.assertEqual(result.fusion.overlap_rule, "shared_overlap") + self.assertEqual( + FusedSearchResult.model_validate_json(result.model_dump_json()), result + ) + legacy = result.model_dump(mode="json") + legacy["fusion"]["overlap_rule"] = "connected_intervals" + self.assertEqual( + FusedSearchResult.model_validate(legacy).fusion.overlap_rule, + "connected_intervals", + ) + del legacy["fusion"]["overlap_rule"] + self.assertEqual( + FusedSearchResult.model_validate(legacy).fusion.overlap_rule, + "connected_intervals", + ) + + def test_nearby_duplicate_hits_combine_into_one_moment(self): + hit_1 = hit("scene", 1, 1.0, 3.0, "scene:1") + hit_2 = hit("scene", 2, 2.0, 3.5, "scene:2") + hit_3 = hit("speech", 1, 2.2, 2.8, "speech:1") + + scene = SearchResult( + query_id="scene:q", + query="dog", + modality="scene", + hits=(hit_1, hit_2), + ) + speech = SearchResult( + query_id="speech:q", + query="dog", + modality="speech", + hits=(hit_3,), + ) + + result = fuse_search_results( + query="dog", + requested_modalities=("scene", "speech"), + results=(scene, speech), + ) + + self.assertEqual(len(result.moments), 1) + self.assertEqual(len(result.moments[0].hits), 3) + self.assertEqual(result.moments[0].start, 2.2) + self.assertEqual(result.moments[0].end, 2.8) + if __name__ == "__main__": unittest.main()