diff --git a/learn2rag/pipeline/config.py b/learn2rag/pipeline/config.py index b0677fb..9225631 100644 --- a/learn2rag/pipeline/config.py +++ b/learn2rag/pipeline/config.py @@ -1,13 +1,20 @@ import json import os import logging +from pathlib import Path -with open(os.environ.get("PIPELINE_USER_CONFIG", "learn2rag/pipeline/user_config.json"), "r") as file: +BASE_DIR = Path(__file__).resolve().parent + +default_user_config = BASE_DIR / "user_config.json" +default_importer_config = BASE_DIR.parent / "importer" / "config" / "config.json" +default_opt_config = BASE_DIR / "opt_config.json" + +with open(os.environ.get("PIPELINE_USER_CONFIG", default_user_config), "r") as file: user_config = json.load(file) -with open(os.environ.get("IMPORTER_CONFIG", "learn2rag/importer/config/config.json"), "r") as file: +with open(os.environ.get("IMPORTER_CONFIG", default_importer_config), "r") as file: importer_config = json.load(file) -with open(os.environ.get("PIPELINE_OPT_CONFIG", "learn2rag/pipeline/opt_config.json"), "r") as file: +with open(os.environ.get("PIPELINE_OPT_CONFIG", default_opt_config), "r") as file: opt_config = json.load(file) logging.info(f"Loaded opt_config:\n{json.dumps(opt_config, indent=4)}") diff --git a/learn2rag/pipeline/search.py b/learn2rag/pipeline/search.py index b35c49c..734979a 100644 --- a/learn2rag/pipeline/search.py +++ b/learn2rag/pipeline/search.py @@ -470,8 +470,42 @@ def search_multi(multi_query: dict[str, str], user_config: dict[str, Any], opt_c async def search_authorized(question: str, user_auths: Mapping[str, Any], *, request_id: str | None = None, user_config: dict[str, Any] = user_config, opt_config: dict[str, Any] = opt_config) -> List[ScoredPoint]: - points = _collect_query_points(question, user_config, opt_config, request_id=request_id) - query_response = QueryResponse(points=points) - authorized_points = await filter_authorized(user_auths, query_response) - # keep deterministic order after auth filter - return _sort_and_deduplicate(list(authorized_points)) + max_retries = opt_config.get("max_auth_retries", 3) + target_k = opt_config.get("top_k", 10) + current_multiplier = opt_config.get("auth_oversample_start", 2) + step_multiplier = opt_config.get("auth_oversample_step", 2) + deduplicated: list[ScoredPoint] = [] + for attempt in range(max_retries): + local_opt = copy.deepcopy(opt_config) + local_opt["top_k"] = target_k * current_multiplier + if "top_k_reranker" in local_opt: + local_opt["top_k_reranker"] = local_opt["top_k_reranker"] * current_multiplier + if "top_k_subqueries" in local_opt: + local_opt["top_k_subqueries"] = local_opt["top_k_subqueries"] * current_multiplier + if "top_k_keywords" in local_opt: + local_opt["top_k_keywords"] = local_opt["top_k_keywords"] * current_multiplier + + for key in ["prefetch_limit_sparse", "prefetch_limit_dense", "prefetch_limit_colbert"]: + if key in local_opt: + local_opt[key] = local_opt[key] * current_multiplier + + profilingLogger.info( + "authorized_search_attempt attempt=%d/%d multiplier=%d target_k=%d", + attempt + 1, + max_retries, + current_multiplier, + target_k, + extra={'activity': 'search_authorized', 'request_id': request_id}, + ) + + points = _collect_query_points(question, user_config, local_opt, request_id=request_id) + query_response = QueryResponse(points=points) + + authorized_points = await filter_authorized(user_auths, query_response) + deduplicated = _sort_and_deduplicate(list(authorized_points)) + + if len(deduplicated) >= target_k: + return deduplicated[:target_k] + current_multiplier += step_multiplier + # If all retries are exhausted, return whatever we managed to authorize + return deduplicated[:target_k] diff --git a/learn2rag/pipeline/tests/test_search_authorized.py b/learn2rag/pipeline/tests/test_search_authorized.py new file mode 100644 index 0000000..e84fcb2 --- /dev/null +++ b/learn2rag/pipeline/tests/test_search_authorized.py @@ -0,0 +1,158 @@ +import sys +import inspect +import unittest +from typing import Any, cast + +from qdrant_client.http.models import ScoredPoint + +# Import the target function +from ..search import search_authorized + + +class SearchAuthorizedTestCase(unittest.IsolatedAsyncioTestCase): + def setUp(self) -> None: + # 1. Verify the search.py file was actually updated! + source = inspect.getsource(search_authorized) + if "max_auth_retries" not in source: + self.fail("CRITICAL ERROR: pipeline/search.py does not contain the updated retry loop.") + + self.user_auths = {"roles": ["user", "admin"]} + self.user_config = {"collection_name": "test_collection"} + + self.opt_config: dict[str, Any] = { + "top_k": 3, + "max_auth_retries": 3, + "auth_oversample_start": 2, + "auth_oversample_step": 2, + "top_k_reranker": 3, + "prefetch_limit_dense": 10, + } + + # 2. BRUTE FORCE MOCKING: Find ALL instances of the search module in memory + self.search_modules = [ + mod for name, mod in sys.modules.items() + if name.endswith('search') and hasattr(mod, '_collect_query_points') + ] + + # Store original functions to restore them cleanly after tests + self.originals = { + mod: (getattr(mod, '_collect_query_points'), getattr(mod, 'filter_authorized', None)) + for mod in self.search_modules + } + + self.collect_calls: list[dict[str, Any]] = [] + self.filter_call_count = 0 + + def tearDown(self) -> None: + # Restore all original functions to memory + for mod, (orig_collect, orig_filter) in self.originals.items(): + setattr(mod, '_collect_query_points', orig_collect) + if orig_filter: + setattr(mod, 'filter_authorized', orig_filter) + + search_authorized.__globals__['_collect_query_points'] = self.originals[self.search_modules[0]][0] + search_authorized.__globals__['filter_authorized'] = self.originals[self.search_modules[0]][1] + + def _apply_patches(self, fake_collect: Any, fake_filter: Any) -> None: + """Inject our fakes into every possible memory space where the code might execute""" + for mod in self.search_modules: + setattr(mod, '_collect_query_points', fake_collect) + setattr(mod, 'filter_authorized', fake_filter) + + # Also patch the direct function globals as a fallback + search_authorized.__globals__['_collect_query_points'] = fake_collect + search_authorized.__globals__['filter_authorized'] = fake_filter + + def _make_mock_points(self, n: int) -> list[ScoredPoint]: + return [ + ScoredPoint( + id=i + 1, + score=1.0 - (i * 0.01), + version=1, + payload={"content": f"mock doc {i}"} + ) for i in range(n) + ] + + def _extract_opt(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]: + if len(args) >= 3: + return cast(dict[str, Any], args[2]) + opt = kwargs.get('opt_config', kwargs.get('local_opt', {})) + return cast(dict[str, Any], opt) + + async def test_success_on_first_attempt(self) -> None: + def fake_collect(*args: Any, **kwargs: Any) -> list[ScoredPoint]: + self.collect_calls.append(self._extract_opt(args, kwargs)) + return self._make_mock_points(6) + + async def fake_filter(*args: Any, **kwargs: Any) -> list[ScoredPoint]: + self.filter_call_count += 1 + return self._make_mock_points(4) + + self._apply_patches(fake_collect, fake_filter) + + results = await search_authorized( + "test query", self.user_auths, user_config=self.user_config, opt_config=self.opt_config + ) + + self.assertEqual(len(self.collect_calls), 1, "Collect was not called exactly once.") + self.assertEqual(len(results), 3) + self.assertEqual(self.collect_calls[0]["top_k"], 6) + self.assertEqual(self.collect_calls[0]["top_k_reranker"], 6) + + async def test_success_on_second_attempt_after_scaling(self) -> None: + def fake_collect(*args: Any, **kwargs: Any) -> list[ScoredPoint]: + self.collect_calls.append(self._extract_opt(args, kwargs)) + return self._make_mock_points(12) + + async def fake_filter(*args: Any, **kwargs: Any) -> list[ScoredPoint]: + self.filter_call_count += 1 + if self.filter_call_count == 1: + return self._make_mock_points(1) # Fail first try + return self._make_mock_points(5) # Succeed second try + + self._apply_patches(fake_collect, fake_filter) + + results = await search_authorized( + "test query", self.user_auths, user_config=self.user_config, opt_config=self.opt_config + ) + + self.assertEqual(len(self.collect_calls), 2, "It should have retried twice") + self.assertEqual(len(results), 3) + self.assertEqual(self.collect_calls[1]["top_k"], 12) + + async def test_max_retries_exhausted(self) -> None: + def fake_collect(*args: Any, **kwargs: Any) -> list[ScoredPoint]: + self.collect_calls.append(self._extract_opt(args, kwargs)) + return self._make_mock_points(12) + + async def fake_filter(*args: Any, **kwargs: Any) -> list[ScoredPoint]: + return self._make_mock_points(1) # Always drop all but 1 + + self._apply_patches(fake_collect, fake_filter) + self.opt_config["max_auth_retries"] = 2 + + results = await search_authorized( + "test query", self.user_auths, user_config=self.user_config, opt_config=self.opt_config + ) + + self.assertEqual(len(self.collect_calls), 2, "Should exhaust exactly 2 max retries") + self.assertEqual(len(results), 1, "Should return what it managed to find") + + async def test_maintains_original_opt_config_immutability(self) -> None: + def fake_collect(*args: Any, **kwargs: Any) -> list[ScoredPoint]: + self.collect_calls.append(self._extract_opt(args, kwargs)) + return self._make_mock_points(12) + + async def fake_filter(*args: Any, **kwargs: Any) -> list[ScoredPoint]: + self.filter_call_count += 1 + if self.filter_call_count == 1: + return self._make_mock_points(1) + return self._make_mock_points(4) + + self._apply_patches(fake_collect, fake_filter) + + await search_authorized( + "test query", self.user_auths, user_config=self.user_config, opt_config=self.opt_config + ) + + self.assertEqual(self.opt_config["top_k"], 3, "Original dictionary should not have been mutated") \ No newline at end of file