retrive more documents - needs more check

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)}")

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]

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")