diff --git a/src/vidxp/application.py b/src/vidxp/application.py index afcf735b..b8462ac5 100644 --- a/src/vidxp/application.py +++ b/src/vidxp/application.py @@ -63,7 +63,10 @@ from vidxp.execution import ExecutionContext, execution_context from vidxp.ports import IndexBackend, ModelRuntimePort, QueryModelPort from vidxp.query_service import GroundedQueryService -from vidxp.search_fusion import fuse_search_results +from vidxp.search_fusion import ( + fuse_search_results, + resolve_candidate_limit, +) from vidxp.model_contracts import ( ModelArtifactDownloadError, ModelArtifactUnavailableError, @@ -529,6 +532,9 @@ def search( config, include_actor=False, ) + candidate_limit = resolve_candidate_limit( + command.top_k, self.settings.search_candidate_depth + ) with self._capability_dependencies(selected): with self.index_backend.open_store(config) as storage: context = CapabilityContext( @@ -542,7 +548,7 @@ def search( modality, query=command.query, media_id=command.media_id, - top_k=command.top_k, + top_k=candidate_limit, context=context, ) for modality in selected @@ -667,6 +673,9 @@ def query_video( results: list[SearchResult] = [] actors: tuple[ActorClusterSummary, ...] = () dependencies = search_modalities + (("actor",) if actor_overview else ()) + candidate_limit = resolve_candidate_limit( + command.top_k, self.settings.search_candidate_depth + ) with self._capability_dependencies(dependencies): with self.index_backend.open_store(config) as storage: context = CapabilityContext( @@ -683,7 +692,7 @@ def query_video( step.modality, query=step.query, media_id=command.media_id, - top_k=command.top_k, + top_k=candidate_limit, context=context, ) ) diff --git a/src/vidxp/search_fusion.py b/src/vidxp/search_fusion.py index b7e41849..db6946a1 100644 --- a/src/vidxp/search_fusion.py +++ b/src/vidxp/search_fusion.py @@ -13,6 +13,26 @@ RRF_RANK_CONSTANT = 60 +DEFAULT_CANDIDATE_DEPTH = 50 +MAX_CANDIDATE_DEPTH = 500 + + +def resolve_candidate_limit( + top_k: int, + candidate_depth: int = DEFAULT_CANDIDATE_DEPTH, + *, + max_candidates: int = MAX_CANDIDATE_DEPTH, +) -> int: + """Determine the per-channel retrieval limit for candidate pools before fusion. + + Uses an independent candidate depth budget, ensuring candidate depth is at + least `top_k` and bounded by `max_candidates`. + """ + if top_k <= 0: + raise ValueError("top_k must be greater than zero.") + if candidate_depth <= 0: + raise ValueError("candidate_depth must be greater than zero.") + return min(max(top_k, candidate_depth), max_candidates) def _query_id( diff --git a/src/vidxp/settings.py b/src/vidxp/settings.py index c5893183..0c18f6a8 100644 --- a/src/vidxp/settings.py +++ b/src/vidxp/settings.py @@ -65,14 +65,10 @@ def _tusd_cors_origins(pattern: str) -> tuple[str, ...]: ) port_text = match.group("port") if port_text is not None and not 1 <= int(port_text) <= 65535: - raise ValueError( - "The upload CORS origin regex contains an invalid port." - ) + raise ValueError("The upload CORS origin regex contains an invalid port.") origins.append(value.replace(r"\.", ".").lower()) if len(set(origins)) != len(origins): - raise ValueError( - "The upload CORS origin regex contains a duplicate origin." - ) + raise ValueError("The upload CORS origin regex contains a duplicate origin.") return tuple(origins) @@ -134,13 +130,9 @@ class VidXPSettings(BaseSettings): mode: ApplicationMode = ApplicationMode.local data_dir: Path = Field(default_factory=default_data_directory) - repository_root: Path = Field( - default_factory=default_repository_directory - ) + repository_root: Path = Field(default_factory=default_repository_directory) runtime_backend: str = "auto" - model_cache: Path = Field( - default_factory=default_model_directory - ) + model_cache: Path = Field(default_factory=default_model_directory) allow_model_downloads: bool = True max_loaded_models: int = Field(default=3, gt=0, le=16) max_concurrent_indexing: int = Field(default=1, gt=0, le=16) @@ -166,6 +158,8 @@ class VidXPSettings(BaseSettings): ) evidence_board_tiles_per_page: int = Field(default=24, gt=0, le=48) evidence_board_pages_per_job: int = Field(default=4, gt=0, le=16) + search_candidate_depth: int = Field(default=50, gt=0, le=500) + http_bind_host: str = Field(default="127.0.0.1", min_length=1) http_port: int = Field(default=DEFAULT_HTTP_PORT, gt=0, le=65535) http_auth_mode: HttpAuthMode = HttpAuthMode.none @@ -382,8 +376,7 @@ def _validate_runtime_backend(cls, value: str) -> str: backend = value.strip().lower() if not re.fullmatch(r"(auto|cpu|mps|cuda(?::[0-9]+)?)", backend): raise ValueError( - "runtime_backend must be auto, cpu, mps, cuda, " - "or cuda:." + "runtime_backend must be auto, cpu, mps, cuda, or cuda:." ) return backend @@ -401,8 +394,7 @@ def _clean_allowlist(cls, values: tuple[str, ...]) -> tuple[str, ...]: ] if invalid: raise ValueError( - "capability_allowlist entries must use " - "DISTRIBUTION:ENTRY_POINT." + "capability_allowlist entries must use DISTRIBUTION:ENTRY_POINT." ) return cleaned @@ -415,9 +407,7 @@ def _clean_allowlist(cls, values: tuple[str, ...]) -> tuple[str, ...]: ) @classmethod def _clean_http_lists(cls, values: tuple[str, ...]) -> tuple[str, ...]: - return tuple( - dict.fromkeys(value.strip() for value in values if value.strip()) - ) + return tuple(dict.fromkeys(value.strip() for value in values if value.strip())) @field_validator("http_trusted_hosts") @classmethod @@ -427,17 +417,10 @@ def _validate_trusted_hosts( ) -> tuple[str, ...]: normalized = tuple(value.lower() for value in values) for value in normalized: - if ( - "*" in value[1:] - or ( - value.startswith("*") - and value != "*" - and not value.startswith("*.") - ) + if "*" in value[1:] or ( + value.startswith("*") and value != "*" and not value.startswith("*.") ): - raise ValueError( - "Trusted-host wildcards must use *.example.com." - ) + raise ValueError("Trusted-host wildcards must use *.example.com.") return normalized @field_validator("mcp_allowed_hosts") @@ -449,13 +432,9 @@ def _validate_mcp_hosts( normalized = tuple(value.lower() for value in values) for value in normalized: if "*" in value and not value.endswith(":*"): - raise ValueError( - "MCP host wildcards are supported only as host:*." - ) + raise ValueError("MCP host wildcards are supported only as host:*.") if "/" in value or "://" in value: - raise ValueError( - "MCP allowed hosts must be Host header values." - ) + raise ValueError("MCP allowed hosts must be Host header values.") return normalized @field_validator("mcp_allowed_origins") @@ -477,9 +456,7 @@ def _validate_mcp_origins( or parsed.query or parsed.fragment ): - raise ValueError( - "MCP allowed origins must be serialized HTTP origins." - ) + raise ValueError("MCP allowed origins must be serialized HTTP origins.") try: parsed.port except ValueError as exc: @@ -487,9 +464,7 @@ def _validate_mcp_origins( "An MCP allowed origin contains an invalid port." ) from exc if "*" in value and not value.endswith(":*"): - raise ValueError( - "MCP origin wildcards are supported only as origin:*." - ) + raise ValueError("MCP origin wildcards are supported only as origin:*.") return values @field_validator("http_oidc_algorithms") @@ -527,23 +502,17 @@ def _validate_oidc_url( if value != value.strip(): raise ValueError(f"{info.field_name} must not contain whitespace.") if "\\" in value or any( - character.isspace() - or ord(character) < 32 - or ord(character) == 127 + character.isspace() or ord(character) < 32 or ord(character) == 127 for character in value ): - raise ValueError( - f"{info.field_name} contains an unsafe URL character." - ) + raise ValueError(f"{info.field_name} contains an unsafe URL character.") parsed = urlsplit(value) if parsed.scheme not in {"http", "https"} or parsed.hostname is None: raise ValueError(f"{info.field_name} must be an HTTP URL.") try: parsed.port except ValueError as exc: - raise ValueError( - f"{info.field_name} contains an invalid port." - ) from exc + raise ValueError(f"{info.field_name} contains an invalid port.") from exc if parsed.username is not None or parsed.password is not None: raise ValueError(f"{info.field_name} must not contain credentials.") if "#" in value: @@ -553,9 +522,7 @@ def _validate_oidc_url( "127.0.0.1", "::1", }: - raise ValueError( - f"{info.field_name} must use HTTPS outside loopback." - ) + raise ValueError(f"{info.field_name} must use HTTPS outside loopback.") if info.field_name == "http_oidc_issuer" and "?" in value: raise ValueError("http_oidc_issuer must not contain a query.") return value @@ -588,32 +555,24 @@ def _validate_service_url( try: parsed.port except ValueError as exc: - raise ValueError( - f"{info.field_name} contains an invalid port." - ) from exc + raise ValueError(f"{info.field_name} contains an invalid port.") from exc if ( info.field_name == "upload_public_endpoint" and parsed.scheme != "https" - and parsed.hostname.lower() not in { + and parsed.hostname.lower() + not in { "localhost", "127.0.0.1", "::1", } ): - raise ValueError( - "upload_public_endpoint must use HTTPS outside loopback." - ) + raise ValueError("upload_public_endpoint must use HTTPS outside loopback.") if info.field_name == "slm_base_url": if parsed.path.rstrip("/") != "/v1": raise ValueError("slm_base_url must end with /v1.") if parsed.hostname.lower() in {"ollama.com", "www.ollama.com"}: - raise ValueError( - "slm_base_url must use a self-hosted Ollama service." - ) - if ( - info.field_name != "slm_base_url" - and not value.endswith("/") - ): + raise ValueError("slm_base_url must use a self-hosted Ollama service.") + if info.field_name != "slm_base_url" and not value.endswith("/"): raise ValueError(f"{info.field_name} must end with a slash.") return value @@ -634,17 +593,13 @@ def _validate_mcp_public_url(cls, value: str | None) -> str | None: or parsed.fragment or parsed.path.rstrip("/") != "/mcp" ): - raise ValueError( - "mcp_public_url must be a plain HTTP URL ending in /mcp." - ) - if ( - parsed.scheme != "https" - and parsed.hostname.lower() - not in {"localhost", "127.0.0.1", "::1"} - ): - raise ValueError( - "mcp_public_url must use HTTPS outside loopback." - ) + raise ValueError("mcp_public_url must be a plain HTTP URL ending in /mcp.") + if parsed.scheme != "https" and parsed.hostname.lower() not in { + "localhost", + "127.0.0.1", + "::1", + }: + raise ValueError("mcp_public_url must use HTTPS outside loopback.") return value.rstrip("/") @field_validator("upload_handoff_public_url") @@ -677,9 +632,7 @@ def _clean_import_roots( cls, values: tuple[Path, ...], ) -> tuple[Path, ...]: - return tuple( - dict.fromkeys(path.expanduser() for path in values) - ) + return tuple(dict.fromkeys(path.expanduser() for path in values)) @model_validator(mode="after") def _require_explicit_server_backend(self) -> "VidXPSettings": @@ -689,17 +642,17 @@ def _require_explicit_server_backend(self) -> "VidXPSettings": "Connect agents to the remote MCP endpoint or use the HTTP " "API directly." ) - if ( - self.mode == ApplicationMode.server - and not re.fullmatch(r"(cpu|cuda(?::[0-9]+)?)", self.runtime_backend) + if self.mode == ApplicationMode.server and not re.fullmatch( + r"(cpu|cuda(?::[0-9]+)?)", self.runtime_backend ): raise ValueError( "Server mode requires an explicit cpu or cuda runtime backend." ) if self.http_auth_mode == HttpAuthMode.static: - if self.http_static_bearer_token is None or len( - self.http_static_bearer_token.get_secret_value() - ) < 32: + if ( + self.http_static_bearer_token is None + or len(self.http_static_bearer_token.get_secret_value()) < 32 + ): raise ValueError( "Static HTTP authentication requires a bearer token of " "at least 32 characters." @@ -805,9 +758,7 @@ def _require_explicit_server_backend(self) -> "VidXPSettings": "per-file upload limit." ) if self.upload_cors_origin_regex is not None: - allowed_origins = _tusd_cors_origins( - self.upload_cors_origin_regex - ) + allowed_origins = _tusd_cors_origins(self.upload_cors_origin_regex) if self.upload_handoff_public_url is not None: parsed_handoff = urlsplit(self.upload_handoff_public_url) handoff_origin = ( @@ -818,9 +769,7 @@ def _require_explicit_server_backend(self) -> "VidXPSettings": "The upload CORS origin regex must allow the handoff origin." ) if (self.slm_base_url is None) != (self.slm_model is None): - raise ValueError( - "slm_base_url and slm_model must be configured together." - ) + raise ValueError("slm_base_url and slm_model must be configured together.") if self.slm_model is not None and self.slm_model.endswith("-cloud"): raise ValueError("slm_model must be a self-hosted Ollama model.") return self @@ -831,13 +780,13 @@ def validate_http_server(self) -> None: and self.http_auth_mode == HttpAuthMode.none ): raise ValueError( - "Server-mode HTTP requires static bearer or OIDC " - "authentication." + "Server-mode HTTP requires static bearer or OIDC authentication." ) - if ( - self.http_auth_mode == HttpAuthMode.none - and self.http_bind_host not in {"127.0.0.1", "::1", "localhost"} - ): + if self.http_auth_mode == HttpAuthMode.none and self.http_bind_host not in { + "127.0.0.1", + "::1", + "localhost", + }: raise ValueError( "Unauthenticated HTTP may bind only to a loopback address." ) @@ -845,13 +794,8 @@ def validate_http_server(self) -> None: raise ValueError("At least one trusted HTTP host is required.") if not self.mcp_allowed_hosts: raise ValueError("At least one allowed MCP host is required.") - if ( - self.http_auth_mode == HttpAuthMode.oidc - and self.mcp_public_url is None - ): - raise ValueError( - "OIDC MCP authentication requires mcp_public_url." - ) + if self.http_auth_mode == HttpAuthMode.oidc and self.mcp_public_url is None: + raise ValueError("OIDC MCP authentication requires mcp_public_url.") @property def layout(self) -> RepositoryLayout: diff --git a/tests/test_application.py b/tests/test_application.py index 57a38cdc..c18a5cc9 100644 --- a/tests/test_application.py +++ b/tests/test_application.py @@ -423,9 +423,7 @@ def detections_handler(context, request): ) self.assertEqual(len({id(context) for context in contexts}), 1) application.artifacts.create_actor_overlay.assert_called_once() - artifact_call = ( - application.artifacts.create_actor_overlay.call_args.kwargs - ) + artifact_call = application.artifacts.create_actor_overlay.call_args.kwargs self.assertEqual(artifact_call["media_id"], MEDIA_ID) self.assertEqual(artifact_call["generation_id"], GENERATION_ID) @@ -436,9 +434,7 @@ def test_create_index_builds_one_central_config(self): original_filename="video.mp4", sha256="1" * 64, ) - application.media.content.return_value = Mock( - path=Path("managed.mp4") - ) + application.media.content.return_value = Mock(path=Path("managed.mp4")) backend.create.return_value = { "media_id": MEDIA_ID, "generation_id": GENERATION_ID, @@ -537,7 +533,7 @@ def handler(context, request): self.assertIsInstance(result, FusedSearchResult) self.assertEqual(result.modalities, ("indexed",)) self.assertEqual(calls[0][1].query, "yellow taxi") - self.assertEqual(calls[0][1].top_k, 7) + self.assertEqual(calls[0][1].top_k, 50) self.assertIs( calls[0][0].storage, manager.__enter__.return_value, @@ -685,6 +681,242 @@ def handler(_context, request): self.assertEqual(searched, ["scene", "speech"]) self.assertEqual(result.modalities, ("scene", "speech")) + def test_search_uses_candidate_limit_for_channel_searches(self): + searched_top_ks = [] + + def search_plugin(name: str) -> CapabilityPlugin: + definition = CapabilityDefinition( + name=name, + description="Search capability.", + extra=name, + collection_name=name, + index_stage=name, + execution_group=name, + operations={ + "search": OperationDefinition( + input_model=SearchInput, + output_model=SearchResult, + ) + }, + ) + + def handler(_context, request): + searched_top_ks.append((name, request.top_k)) + return SearchResult( + query_id=f"{name}:query", + query=request.query, + modality=name, + ) + + return CapabilityPlugin( + definition=definition, + executor_factory=lambda: CapabilityExecutor( + indexer=Mock(), + operations={"search": handler}, + ), + ) + + registry = CapabilityRegistry((search_plugin("scene"),)) + manager = MagicMock() + manager.__enter__.return_value = Mock(spec=IndexStore) + application, backend = self.application("repository", registry=registry) + backend.active_config.return_value = IndexConfig.local( + enabled_modalities=("scene",), + collection_names={"scene": "scene"}, + ) + backend.open_store.return_value = manager + + # Request top_k=3 to caller + application.search(SearchCommand(query="taxi", top_k=3)) + + # Capability search operation should have received expanded candidate limit (50) + self.assertEqual(searched_top_ks, [("scene", 50)]) + + def test_query_video_uses_candidate_limit_for_channel_searches(self): + requests = [] + + def handler(_context, request): + requests.append(request) + return SearchResult( + query_id="indexed:1", + query=request.query, + modality="indexed", + ) + + manager = MagicMock() + manager.__enter__.return_value = Mock(spec=IndexStore) + application = self.indexed_application(handler, manager) + pinned = IndexConfig.local( + enabled_modalities=("indexed",), + collection_names={"indexed": "indexed"}, + snapshot_id=SNAPSHOT_ID, + snapshot_sha256=SNAPSHOT_SHA256, + ) + application.index_backend.config_for_snapshot.return_value = pinned + + # Request top_k=3 for video query + result = application.query_video( + QueryVideoCommand( + question="What happens?", + media_id=MEDIA_ID, + modalities=("indexed",), + top_k=3, + ), + snapshot=IndexSnapshotReference( + snapshot_id=SNAPSHOT_ID, + snapshot_sha256=SNAPSHOT_SHA256, + ), + ) + + # Underlying search operation should receive expanded candidate limit (50) + self.assertEqual(requests[0].top_k, 50) + # Fused moments / evidence should respect requested top_k=3 + self.assertLessEqual(len(result.moments), 3) + + def test_application_search_recovers_moment_ranked_past_public_top_k(self): + def search_plugin(name: str, hits_fn) -> CapabilityPlugin: + definition = CapabilityDefinition( + name=name, + description="Search capability.", + extra=name, + collection_name=name, + index_stage=name, + execution_group=name, + operations={ + "search": OperationDefinition( + input_model=SearchInput, + output_model=SearchResult, + ) + }, + ) + + def handler(_context, request): + hits = hits_fn(request.top_k) + return SearchResult( + query_id=f"{name}:query", + query=request.query, + modality=name, + hits=tuple(hits), + ) + + return CapabilityPlugin( + definition=definition, + executor_factory=lambda: CapabilityExecutor( + indexer=Mock(), + operations={"search": handler}, + ), + ) + + def scene_hits(limit: int): + all_hits = [ + SearchHit( + rank=1, + media_id=MEDIA_ID, + video_id=MEDIA_ID, + generation_id=GENERATION_ID, + start=10, + end=20, + score=-1.0, + raw_distance=1.0, + modality="scene", + source_id="s1", + ), + SearchHit( + rank=2, + media_id=MEDIA_ID, + video_id=MEDIA_ID, + generation_id=GENERATION_ID, + start=30, + end=40, + score=-2.0, + raw_distance=2.0, + modality="scene", + source_id="s2", + ), + SearchHit( + rank=3, + media_id=MEDIA_ID, + video_id=MEDIA_ID, + generation_id=GENERATION_ID, + start=50, + end=60, + score=-3.0, + raw_distance=3.0, + modality="scene", + source_id="s3_shared", + ), + ] + return all_hits[:limit] + + def speech_hits(limit: int): + all_hits = [ + SearchHit( + rank=1, + media_id=MEDIA_ID, + video_id=MEDIA_ID, + generation_id=GENERATION_ID, + start=70, + end=80, + score=-1.0, + raw_distance=1.0, + modality="speech", + source_id="p1", + ), + SearchHit( + rank=2, + media_id=MEDIA_ID, + video_id=MEDIA_ID, + generation_id=GENERATION_ID, + start=90, + end=100, + score=-2.0, + raw_distance=2.0, + modality="speech", + source_id="p2", + ), + SearchHit( + rank=3, + media_id=MEDIA_ID, + video_id=MEDIA_ID, + generation_id=GENERATION_ID, + start=50, + end=60, + score=-3.0, + raw_distance=3.0, + modality="speech", + source_id="s3_shared_speech", + ), + ] + return all_hits[:limit] + + registry = CapabilityRegistry( + ( + search_plugin("scene", scene_hits), + search_plugin("speech", speech_hits), + ) + ) + manager = MagicMock() + manager.__enter__.return_value = Mock(spec=IndexStore) + application, backend = self.application("repository", registry=registry) + backend.active_config.return_value = IndexConfig.local( + enabled_modalities=("scene", "speech"), + collection_names={"scene": "scene", "speech": "speech"}, + ) + backend.open_store.return_value = manager + + # Public top_k is 2. Moment at [50, 60] is rank 3 in scene and rank 3 in speech. + fused = application.search( + SearchCommand(query="test", modalities=("scene", "speech"), top_k=2) + ) + + # Fused output is trimmed to public top_k=2 + self.assertEqual(len(fused.moments), 2) + # Shared rank-3 moment retrieved into candidate pool becomes #1 fused result + top_moment = fused.moments[0] + self.assertEqual(top_moment.start, 50) + self.assertEqual(top_moment.end, 60) + self.assertEqual(top_moment.modalities, ("scene", "speech")) + def test_application_boundary_returns_stable_validation_error(self): application, _ = self.application("unused") @@ -749,9 +981,7 @@ def test_validation_error_does_not_expose_input(self): def test_missing_media_error_does_not_expose_path(self): application, _ = self.application("unused") secret_path = Path("private/customer/video.mp4") - application.media.require_record.side_effect = ( - MediaUnavailableError("secret") - ) + application.media.require_record.side_effect = MediaUnavailableError("secret") with self.assertRaises(ApplicationError) as raised: application.create_index( @@ -836,9 +1066,7 @@ def test_prepare_dependency_failure_is_stable(self): "modality": "noop", } }, - prepare=Mock( - side_effect=ModuleNotFoundError("provider.internal") - ), + prepare=Mock(side_effect=ModuleNotFoundError("provider.internal")), ), ) registry = CapabilityRegistry((plugin,)) @@ -861,9 +1089,7 @@ def test_missing_model_is_not_misclassified_as_a_package_dependency(self): original_filename="video.mp4", sha256="1" * 64, ) - application.media.content.return_value = Mock( - path=Path("video.mp4") - ) + application.media.content.return_value = Mock(path=Path("video.mp4")) backend.create.side_effect = ModelArtifactUnavailableError("scene") with self.assertRaises(ModelUnavailableError) as raised: @@ -1170,9 +1396,7 @@ def test_remove_delegates_to_index_backend(self): backend.remove.return_value = True self.assertTrue( - application.remove_from_index( - RemoveIndexCommand(media_id=MEDIA_ID) - ) + application.remove_from_index(RemoveIndexCommand(media_id=MEDIA_ID)) ) config, media_id = backend.remove.call_args.args @@ -1284,9 +1508,7 @@ def test_local_backend_injects_storage_for_generation_build(self): index_video.call_args.kwargs["manifest_store"].runtime, backend.runtime, ) - self.assertIsNotNone( - index_video.call_args.kwargs["config"].generation_id - ) + self.assertIsNotNone(index_video.call_args.kwargs["config"].generation_id) cleanup_storage.__exit__.assert_called_once() storage.__exit__.assert_called_once() diff --git a/tests/test_search_fusion.py b/tests/test_search_fusion.py index 622e0210..9c2486ac 100644 --- a/tests/test_search_fusion.py +++ b/tests/test_search_fusion.py @@ -1,7 +1,11 @@ import unittest from vidxp.application_models import SearchHit, SearchResult -from vidxp.search_fusion import RRF_RANK_CONSTANT, fuse_search_results +from vidxp.search_fusion import ( + RRF_RANK_CONSTANT, + fuse_search_results, + resolve_candidate_limit, +) MEDIA_ID = "123456781234423481234567890abcde" @@ -96,9 +100,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 +111,61 @@ def test_rewritten_atomic_query_identity_changes_fused_identity(self): self.assertNotEqual(first.query_id, second.query_id) + def test_resolve_candidate_limit_bounds(self): + self.assertEqual(resolve_candidate_limit(3, 50), 50) + self.assertEqual(resolve_candidate_limit(10, 50), 50) + self.assertEqual(resolve_candidate_limit(100, 50), 100) + self.assertEqual(resolve_candidate_limit(600, 50), 500) + with self.assertRaises(ValueError): + resolve_candidate_limit(0) + with self.assertRaises(ValueError): + resolve_candidate_limit(-1) + with self.assertRaises(ValueError): + resolve_candidate_limit(5, candidate_depth=0) + + def test_fusion_recovers_top_result_when_candidates_extend_beyond_final_top_k(self): + # Top-k requested is 2. + # Moment 1 is rank 1 in scene, rank 10 in speech. + # Moment 2 is rank 1 in speech, rank 10 in scene. + # Moment 3 is rank 3 in scene AND rank 3 in speech. + # If each search only returned top_k=2 candidates, Moment 3 would be missed completely. + # When candidate depth extends past rank 2, Moment 3 is included in both channels + # and becomes the top fused result due to strong combined rank (1/63 + 1/63 > 1/61). + scene = SearchResult( + query_id="scene:q", + query="test", + modality="scene", + hits=( + hit("scene", 1, 10, 20, "s1"), + hit("scene", 2, 30, 40, "s2"), + hit("scene", 3, 50, 60, "s3"), + ), + ) + speech = SearchResult( + query_id="speech:q", + query="test", + modality="speech", + hits=( + hit("speech", 1, 70, 80, "p1"), + hit("speech", 2, 90, 100, "p2"), + hit("speech", 3, 50, 60, "s3_speech"), + ), + ) + + fused = fuse_search_results( + query="test", + requested_modalities=("scene", "speech"), + results=(scene, speech), + top_k=2, + ) + + self.assertEqual(len(fused.moments), 2) + top_moment = fused.moments[0] + self.assertEqual(top_moment.start, 50) + self.assertEqual(top_moment.end, 60) + self.assertEqual(top_moment.modalities, ("scene", "speech")) + self.assertAlmostEqual(top_moment.score, 2 / (RRF_RANK_CONSTANT + 3)) + if __name__ == "__main__": unittest.main()