benchmax 0.1.2.dev34__tar.gz → 0.1.2.dev35__tar.gz

{benchmax-0.1.2.dev34/src/benchmax.egg-info → benchmax-0.1.2.dev35}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: benchmax
-Version: 0.1.2.dev34
+Version: 0.1.2.dev35
 Summary: Framework-Agnostic RL Environments for LLM Fine-Tuning
 Author: castie@castform.com
 Classifier: Programming Language :: Python :: 3
@@ -12,7 +12,7 @@ Requires-Dist: aiohttp>=3.13.1
 Requires-Dist: asyncio>=4.0.0
 Requires-Dist: cloudpickle>=3.0.0
 Requires-Dist: datasets>=4.0.0
-Requires-Dist: httpx>=0.27.0
+Requires-Dist: httpx[http2]>=0.27.0
 Requires-Dist: json-repair>=0.59.10
 Requires-Dist: openai>=2.15.0
 Requires-Dist: pydantic>=2.0.0

{benchmax-0.1.2.dev34 → benchmax-0.1.2.dev35}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "benchmax"
-version = "0.1.2.dev34"
+version = "0.1.2.dev35"
 description = "Framework-Agnostic RL Environments for LLM Fine-Tuning"
 readme = "README.md"
 authors = [{ name = "castie@castform.com" }]
@@ -10,7 +10,7 @@ dependencies = [
     "asyncio>=4.0.0",
     "cloudpickle>=3.0.0",
     "datasets>=4.0.0",
-    "httpx>=0.27.0",
+    "httpx[http2]>=0.27.0",
     "json-repair>=0.59.10",
     "openai>=2.15.0",
     "pydantic>=2.0.0",

{benchmax-0.1.2.dev34 → benchmax-0.1.2.dev35}/src/benchmax/cli/corpus.py

@@ -110,7 +110,7 @@ def _cmd_corpus_list(args: argparse.Namespace) -> int:
     if not corpora:
         print("No corpora yet. Create one: castform corpus ingest <folder>")
         return 0
-    print(f"{len(corpora)}/5 corpora:")
+    print(f"{len(corpora)}/20 corpora:")
     for c in corpora:
         print(f"  {c.name}  (id: {c.id})")
     return 0
@@ -206,7 +206,7 @@ def register(sub: argparse._SubParsersAction) -> None:
     p_ing.add_argument("--json", action="store_true", help="Emit raw JSON")
     p_ing.set_defaults(func=_cmd_corpus_ingest)
 
-    p_ls = corpus_sub.add_parser("list", help="List your corpora (and the 5-corpus cap)")
+    p_ls = corpus_sub.add_parser("list", help="List your corpora (and the 20-corpus cap)")
     p_ls.add_argument("--json", action="store_true", help="Emit raw JSON")
     p_ls.set_defaults(func=_cmd_corpus_list)
 

{benchmax-0.1.2.dev34 → benchmax-0.1.2.dev35}/src/benchmax/cli/launch.py

@@ -126,7 +126,7 @@ def _cmd_launch(args: argparse.Namespace) -> int:
                 )
                 return 1
             reply = input(
-                f"Launch '{run_name}' (type={args.type}) — incurs GPU cost. Continue? [y/N] "
+                f"Launch '{run_name}' — incurs GPU cost. Continue? [y/N] "
             )
             if reply.strip().lower() not in ("y", "yes"):
                 print("Aborted.")
@@ -177,7 +177,6 @@ def _cmd_launch(args: argparse.Namespace) -> int:
         with warnings.catch_warnings(record=True) as caught:
             warnings.simplefilter("always")
             run_id = client.launch_training_run(
-                training_run_type=args.type,
                 name=run_name,
                 launcher_args=launcher_args or None,
                 **dataclasses.asdict(uploaded),
@@ -211,11 +210,6 @@ def register(sub: argparse._SubParsersAction) -> None:
     )
     p.add_argument("--eval", default="eval_dataset.jsonl", help="Eval dataset (jsonl)")
     p.add_argument("--name", help="Run name (default: the env class name)")
-    p.add_argument(
-        "--type",
-        default="simple",
-        help="Training run type: simple (GPU) or simple-cpu (smoke)",
-    )
     p.add_argument(
         "--env-arg", action="append", metavar="KEY=VALUE", help="Env constructor arg"
     )

{benchmax-0.1.2.dev34 → benchmax-0.1.2.dev35}/src/benchmax/cli/scaffold/skills/launch-run/SKILL.md

@@ -66,9 +66,3 @@ truncated in training — keep `MAX_SEARCH_CALLS` ≤ 8 (see design-environment'
 Server-controlled fields — `save`, `load`, `global_batch_size`, the eval mirrors —
 are **not settable**: the launch handler fills them in and rejects caller input
 that carries them. (`rollout_batch_size` is derived too, not a launch arg.)
-
-### Run types
-
-`--type simple` (default) is the GPU training pool. `--type simple-cpu` is a
-CPU-only smoke pool (cheap) for exercising the launch lifecycle without GPU.
-(`simple-r5` from older docs is not implemented.)

{benchmax-0.1.2.dev34 → benchmax-0.1.2.dev35}/src/benchmax/envs/postgres_search/search_env.py

@@ -23,7 +23,6 @@ from benchmax.envs.example_id import make_example
 from benchmax.envs.reward_helpers import (
     clip01,
     count_search_calls,
-    extract_answer_block,
     extract_completion_text,
     search_within_budget,
 )
@@ -40,6 +39,19 @@ logger = logging.getLogger(__name__)
 
 _CITATION_RE = re.compile(r"\[Source:\s*([^\]]+)\]", re.IGNORECASE)
 
+_ANSWER_TAG_RE = re.compile(r"<answer>(.*?)</answer>", re.DOTALL | re.IGNORECASE)
+
+
+def _extract_answer_block(text: str) -> str:
+    """Extract content from <answer> tags; return "" if no answer tag is present.
+
+    Strict variant of ``reward_helpers.extract_answer_block``: a completion
+    without an explicit ``<answer>`` block scores as no answer rather than
+    falling back to the full text.
+    """
+    match = _ANSWER_TAG_RE.search(text or "")
+    return match.group(1).strip() if match else ""
+
 # Match Python-style `{name}` placeholders with word-char names only —
 # leaves JSON-like literals (e.g. `{"answer": "X"}`) and unknown keys
 # untouched, so a user-edited SYSTEM_PROMPT_TEMPLATE that contains JSON
@@ -70,7 +82,6 @@ _CORRECTNESS_RUBRIC = Rubric(
     type="positive",
     score_map={
         0: "Provided answer is missing or incorrect.",
-        0.5: "Partially correct — captures some facts but missing key details.",
         1: "Fully correct and factually consistent.",
     },
 )
@@ -286,7 +297,7 @@ tags. Cite your sources inline using [Source: <source_id>] next to each claim.
                 return zeros
 
             t = task or {}
-            answer = extract_answer_block(text)
+            answer = _extract_answer_block(text)
             prompt = str(t.get("question") or t.get("prompt") or "")
             gt_str = str(t.get("ground_truth") or "")
             reference_chunks = t.get("reference_chunks", [])

{benchmax-0.1.2.dev34 → benchmax-0.1.2.dev35}/src/benchmax/envs/telestich/example.py

@@ -642,12 +642,11 @@ if __name__ == "__main__":
     ):
         print(f"  {label:<14}: {path}")
 
-    # 4. Launch the training run. training_run_type="simple" + the `model` arg select
-    #    the trainer YAML/pool server-side (Qwen3.5-4B→gpu4, Qwen3.5-35B-A3B→gpu8).
+    # 4. Launch the training run. The model arg selects the trainer YAML/pool
+    #    server-side (Qwen3.5-4B→gpu4, Qwen3.5-35B-A3B→gpu8).
     print(f"\nLaunching training run (model={MODEL}) ...")
     with TrainerClient(api_key="", base_url=BASE_URL) as trainer:
         run_id = trainer.launch_training_run(
-            training_run_type="simple",
             env_cls_path=uploaded.env_cls_path,
             env_metadata_path=uploaded.env_metadata_path,
             train_dataset_path=uploaded.train_dataset_path,

{benchmax-0.1.2.dev34 → benchmax-0.1.2.dev35}/src/benchmax/platform/client.py

@@ -360,7 +360,6 @@ class TrainerClient:
     Example:
         client = TrainerClient(api_key="sk_...", base_url="http://localhost:3000")
         run_id = client.launch_training_run(
-            training_run_type="simple",
             env_cls_path="envs/run-abc/abc123/env-cls.pkl",
             env_metadata_path="envs/run-abc/abc123/env-metadata.json",
             train_dataset_path="datasets/run-abc/def456/train.jsonl",
@@ -413,7 +412,6 @@ class TrainerClient:
 
     def launch_training_run(
         self,
-        training_run_type: str,
         env_cls_path: str,
         env_metadata_path: str,
         train_dataset_path: str,
@@ -421,12 +419,9 @@ class TrainerClient:
         name: str | None = None,
         launcher_args: dict[str, Any] | None = None,
     ) -> str:
-        """Launch a new training run from a job template.
+        """Launch a new training run.
 
         Args:
-            training_run_type: Job template selector. ``"simple"`` (GPU pool —
-                gpu4 for 4B, gpu8 for 35B) or ``"simple-cpu"`` (CPU-only smoke
-                pool, no GPU).
             env_cls_path: Path to the environment class pickle (.pkl file)
             env_metadata_path: Path to the environment metadata JSON file
             train_dataset_path: Path to the training dataset
@@ -453,7 +448,7 @@ class TrainerClient:
         response = self._http_client.post(
             "/v1/train/runs/launch",
             json={
-                "type": training_run_type,
+                "type": "simple",
                 "name": name,
                 "args": args,
             },

{benchmax-0.1.2.dev34 → benchmax-0.1.2.dev35}/src/benchmax/platform/training_run.py

@@ -36,7 +36,6 @@ class UploadedTrainingRun:
 
         uploaded = upload_training_run(...)
         run_id = trainer.launch_training_run(
-            training_run_type="simple",
             **dataclasses.asdict(uploaded),
         )
     """

{benchmax-0.1.2.dev34 → benchmax-0.1.2.dev35}/src/benchmax/rag/corpus/postgres/client.py

@@ -2,7 +2,9 @@
 
 from __future__ import annotations
 
+import asyncio
 import logging
+import threading
 import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass, field
@@ -15,7 +17,6 @@ from benchmax.platform.credentials import TokenProvider, platform_bearer
 
 from .exceptions import (
     AuthenticationError,
-    ChunkLimitError,
     CorpusAPIError,
     CorpusLimitError,
     CorpusNotFoundError,
@@ -52,11 +53,38 @@ class CorpusClient:
     max_retries: int = 5
     retry_backoff_seconds: float = 0.5
     token_provider: TokenProvider = platform_bearer
-    _http_client: httpx.Client = field(init=False, repr=False)
-
-    def __post_init__(self) -> None:
-        """Initialize the persistent HTTP client. Auth is attached per request
-        in ``_request`` — not baked here."""
+    # Enable HTTP/2 multiplexing on the async client. Safe there (one client
+    # bound to one event loop), unlike the shared sync client across threads.
+    async_http2: bool = True
+    # HTTP clients are created lazily, one per thread (see ``_http_client``).
+    # httpx.Client's connection pool is not safe to share across threads at high
+    # parallelism: the QA-gen work queue hits this client from every batch
+    # thread, which raced the shared pool's sockets into
+    # ``ReadError: [Errno 9] Bad file descriptor``. A client-per-thread avoids
+    # the shared-pool race entirely.
+    _local: threading.local = field(
+        init=False, repr=False, default_factory=threading.local
+    )
+    _client_override: httpx.Client | None = field(
+        init=False, repr=False, default=None
+    )
+    _client_registry: list[httpx.Client] = field(
+        init=False, repr=False, default_factory=list
+    )
+    _registry_lock: threading.Lock = field(
+        init=False, repr=False, default_factory=threading.Lock
+    )
+    # Single async client, lazily bound to the running event loop (rebuilt if the
+    # loop changes — asyncio.run() mints a fresh loop per Pipeline.run()).
+    _async_client: httpx.AsyncClient | None = field(
+        init=False, repr=False, default=None
+    )
+    _async_client_loop: Any = field(init=False, repr=False, default=None)
+
+    def _build_http_client(self) -> httpx.Client:
+        """Create a new HTTP client and register it for ``close()``.
+
+        Auth is attached per request in ``_request`` — not baked here."""
         timeout_config = httpx.Timeout(
             timeout=self.timeout,
             connect=self.timeout,
@@ -64,11 +92,37 @@ class CorpusClient:
             write=self.timeout,
             pool=self.timeout,
         )
-        self._http_client = httpx.Client(
+        client = httpx.Client(
             base_url=self.base_url,
             headers={"Content-Type": "application/json"},
             timeout=timeout_config,
         )
+        with self._registry_lock:
+            self._client_registry.append(client)
+        return client
+
+    @property
+    def _http_client(self) -> httpx.Client:
+        """The HTTP client for the calling thread.
+
+        Returns an explicitly installed override if present (e.g. a profiling
+        harness swapping in an HTTP/2 client); otherwise a lazily-created
+        per-thread client so concurrent batch threads never share a pool."""
+        if self._client_override is not None:
+            return self._client_override
+        client = getattr(self._local, "client", None)
+        if client is None:
+            client = self._build_http_client()
+            self._local.client = client
+        return client
+
+    @_http_client.setter
+    def _http_client(self, value: httpx.Client) -> None:
+        """Install a single client shared across all threads. Intended for
+        single-threaded or multiplexed (HTTP/2) setups, not pool-per-thread."""
+        self._client_override = value
+        with self._registry_lock:
+            self._client_registry.append(value)
 
     def _request(self, method: str, path: str, **kwargs: Any) -> httpx.Response:
         """Execute an HTTP request with retry/backoff for transient network failures.
@@ -156,6 +210,97 @@ class CorpusClient:
                 pass
         return self.retry_backoff_seconds * (2 ** (attempt - 1))
 
+    def _get_async_client(self) -> httpx.AsyncClient:
+        """The async client bound to the running event loop.
+
+        Rebuilt when the loop changes (``asyncio.run`` mints a fresh loop per
+        ``Pipeline.run``) or the client was closed. Creation does not ``await``,
+        so on a single event loop the check-then-build is race-free."""
+        loop = asyncio.get_running_loop()
+        client = self._async_client
+        if client is None or client.is_closed or self._async_client_loop is not loop:
+            timeout_config = httpx.Timeout(
+                timeout=self.timeout,
+                connect=self.timeout,
+                read=self.timeout,
+                write=self.timeout,
+                pool=self.timeout,
+            )
+            client = httpx.AsyncClient(
+                base_url=self.base_url,
+                headers={"Content-Type": "application/json"},
+                timeout=timeout_config,
+                http2=self.async_http2,
+            )
+            self._async_client = client
+            self._async_client_loop = loop
+        return client
+
+    async def _arequest(self, method: str, path: str, **kwargs: Any) -> httpx.Response:
+        """Async twin of ``_request`` — same retry/backoff and 429 handling, with
+        ``await asyncio.sleep`` instead of ``time.sleep`` so the loop stays free."""
+        try:
+            bearer = self.token_provider()
+        except RuntimeError as exc:
+            # The seam (platform_bearer) raises when no credential resolves; surface
+            # it as an auth error so callers catch it like any other Corpora failure.
+            raise AuthenticationError(
+                f"No Castform platform credential available for the Corpora API: {exc}"
+            ) from exc
+        headers = {
+            **kwargs.pop("headers", {}),
+            "Authorization": f"Bearer {bearer}",
+        }
+        retries = max(1, int(self.max_retries))
+        client = self._get_async_client()
+        attempt = 1
+        while True:
+            try:
+                response = await client.request(method, path, headers=headers, **kwargs)
+            except (httpx.ConnectTimeout, httpx.ConnectError, httpx.ReadTimeout) as exc:
+                if attempt >= retries:
+                    raise CorpusAPIError(
+                        (
+                            "Corpora API request failed after retries due to a network timeout/error. "
+                            f"method={method} path={path} base_url={self.base_url} "
+                            f"attempts={retries} last_error={exc!s}"
+                        ),
+                        status_code=503,
+                    ) from exc
+                delay = self.retry_backoff_seconds * (2 ** (attempt - 1))
+                logger.warning(
+                    "Corpora API request attempt %s/%s failed (%s). Retrying in %.2fs. "
+                    "method=%s path=%s base_url=%s",
+                    attempt,
+                    retries,
+                    type(exc).__name__,
+                    delay,
+                    method,
+                    path,
+                    self.base_url,
+                )
+                await asyncio.sleep(delay)
+                attempt += 1
+                continue
+
+            if response.status_code == 429 and attempt < retries:
+                delay = self._retry_after_delay(response, attempt)
+                logger.warning(
+                    "Corpora API rate-limited (429) on attempt %s/%s. Retrying in %.2fs. "
+                    "method=%s path=%s base_url=%s",
+                    attempt,
+                    retries,
+                    delay,
+                    method,
+                    path,
+                    self.base_url,
+                )
+                await asyncio.sleep(delay)
+                attempt += 1
+                continue
+
+            return response
+
     def __enter__(self) -> "CorpusClient":
         return self
 
@@ -163,8 +308,28 @@ class CorpusClient:
         self.close()
 
     def close(self) -> None:
-        """Close the HTTP client."""
-        self._http_client.close()
+        """Close every HTTP client this instance created (one per thread, plus
+        any installed override)."""
+        with self._registry_lock:
+            clients = list(self._client_registry)
+            self._client_registry.clear()
+        self._client_override = None
+        for client in clients:
+            try:
+                client.close()
+            except Exception:  # noqa: BLE001 — best-effort cleanup
+                logger.debug("Error closing corpus HTTP client", exc_info=True)
+
+    async def aclose(self) -> None:
+        """Close the async client. Call from within its event loop."""
+        client = self._async_client
+        self._async_client = None
+        self._async_client_loop = None
+        if client is not None and not client.is_closed:
+            try:
+                await client.aclose()
+            except Exception:  # noqa: BLE001 — best-effort cleanup
+                logger.debug("Error closing corpus async client", exc_info=True)
 
     def _handle_response_errors(self, response: httpx.Response) -> None:
         """Convert HTTP errors to appropriate exceptions."""
@@ -181,7 +346,7 @@ class CorpusClient:
             raise AuthenticationError(message)
 
         if response.status_code == 400:
-            if "Maximum of 5 corpora" in message:
+            if "Maximum of 20 corpora" in message:
                 raise CorpusLimitError()
             if "Chunk limit exceeded" in message:
                 raise CorpusAPIError(message, 400)
@@ -204,7 +369,7 @@ class CorpusClient:
             Corpus object with id, name, timestamps
 
         Raises:
-            CorpusLimitError: If max 5 corpora limit reached
+            CorpusLimitError: If max 20 corpora limit reached
             AuthenticationError: If API key is invalid
         """
         response = self._request("POST", "/v1/corpora", json={"name": name})
@@ -306,7 +471,7 @@ class CorpusClient:
             print(f"     ID: {corpus.id}")
             print(f"     Created: {corpus.created_at}")
 
-        print(f"\n  0. Cancel operation")
+        print("\n  0. Cancel operation")
         print()
 
         while True:
@@ -572,3 +737,63 @@ class CorpusClient:
                 matched.append((local_chunk, corpus_chunk.score or 0.0))
 
         return matched
+
+    async def asearch(
+        self,
+        corpus_id: str,
+        query: str,
+        limit: int = 10,
+        offset: int = 0,
+        metadata: dict[str, Any] | None = None,
+        filters: dict[str, Any] | None = None,
+    ) -> SearchResult:
+        """Async twin of ``search``. Same payload + response shape, async I/O."""
+        payload: dict[str, Any] = {"query": query, "limit": limit, "offset": offset}
+        if metadata:
+            payload["metadata"] = metadata
+        if filters:
+            payload["filters"] = filters
+
+        response = await self._arequest(
+            "POST", f"/v1/corpora/{corpus_id}/search", json=payload
+        )
+        self._handle_response_errors(response)
+
+        data = response.json()
+        results = [
+            CorpusChunk(
+                id=r["id"],
+                content=r["content"],
+                metadata=r.get("metadata") or {},
+                score=r.get("score"),
+            )
+            for r in data.get("results", [])
+        ]
+
+        return SearchResult(results=results, total=data.get("total", 0), query=query)
+
+    async def asearch_with_chunks(
+        self,
+        corpus_id: str,
+        query: str,
+        collection: ChunkCollection,
+        limit: int = 10,
+        metadata: dict[str, Any] | None = None,
+        filters: dict[str, Any] | None = None,
+    ) -> list[tuple[Chunk, float]]:
+        """Async twin of ``search_with_chunks``."""
+        result = await self.asearch(
+            corpus_id=corpus_id,
+            query=query,
+            limit=limit,
+            metadata=metadata,
+            filters=filters,
+        )
+
+        matched: list[tuple[Chunk, float]] = []
+        for corpus_chunk in result.results:
+            local_chunk = collection.get_chunk_by_hash(corpus_chunk.id)
+            if local_chunk:
+                matched.append((local_chunk, corpus_chunk.score or 0.0))
+
+        return matched

{benchmax-0.1.2.dev34 → benchmax-0.1.2.dev35}/src/benchmax/rag/corpus/postgres/exceptions.py

@@ -24,10 +24,10 @@ class AuthenticationError(CorpusAPIError):
 
 
 class CorpusLimitError(CorpusAPIError):
-    """Maximum corpus limit (5) reached."""
+    """Maximum corpus limit (20) reached."""
 
     def __init__(self, existing_corpora: list[Corpus] | None = None):
-        super().__init__("Maximum of 5 corpora per user reached", 400)
+        super().__init__("Maximum of 20 corpora per user reached", 400)
         self.existing_corpora = existing_corpora or []
 
 

{benchmax-0.1.2.dev34 → benchmax-0.1.2.dev35}/src/benchmax/rag/corpus/postgres/source.py

@@ -329,36 +329,93 @@ class PostgresChunkSource:
                 collection=self.collection,
                 limit=top_k,
             )
+            self._accumulate_related(related_map, source, query, matched_chunks, top_k)
 
-            for result_chunk, score in matched_chunks[:top_k]:
-                if result_chunk.hash == source.hash:
+        return self._sorted_related(related_map)
+
+    async def asearch_related(
+        self,
+        source: Chunk,
+        queries: list[str],
+        top_k: int = 5,
+        mode: SearchMode | None = None,
+        hybrid: HybridOptions | None = None,
+    ) -> list[dict]:
+        """Async twin of ``search_related`` — identical dedup/neighbor-skip/scoring,
+        async corpus I/O. Queries run sequentially for parity with the sync path;
+        cross-batch search concurrency comes from the async work queue."""
+        if hybrid is not None:
+            warnings.warn(
+                "PostgresChunkSource does not support hybrid search; 'hybrid' parameter is ignored.",
+                stacklevel=2,
+            )
+        if mode is not None and mode != "lexical":
+            warnings.warn(
+                f"PostgresChunkSource only supports 'lexical' mode; '{mode}' will be ignored.",
+                stacklevel=2,
+            )
+        self._assert_ready()
+        related_map: dict[str, dict] = {}
+
+        for query in queries:
+            matched_chunks = await self._client.asearch_with_chunks(
+                corpus_id=self._corpus.id,
+                query=query,
+                collection=self.collection,
+                limit=top_k,
+            )
+            self._accumulate_related(related_map, source, query, matched_chunks, top_k)
+
+        return self._sorted_related(related_map)
+
+    async def aclose(self) -> None:
+        """Close the underlying corpus client's async transport (best-effort).
+        Call from within the event loop that used it."""
+        await self._client.aclose()
+
+    @staticmethod
+    def _accumulate_related(
+        related_map: dict[str, dict],
+        source: Chunk,
+        query: str,
+        matched_chunks: list[tuple[Chunk, float]],
+        top_k: int,
+    ) -> None:
+        """Merge one query's results into ``related_map``: skip the source chunk
+        and its same-file neighbors, dedup by hash, aggregate queries + max score."""
+        for result_chunk, score in matched_chunks[:top_k]:
+            if result_chunk.hash == source.hash:
+                continue
+
+            is_same_file = result_chunk.get_metadata("file") == source.get_metadata(
+                "file"
+            )
+            if is_same_file:
+                index_diff = abs(
+                    result_chunk.get_metadata("index", 0)
+                    - source.get_metadata("index", 0)
+                )
+                if index_diff <= 1:
                     continue
 
-                is_same_file = result_chunk.get_metadata("file") == source.get_metadata(
-                    "file"
+            if result_chunk.hash not in related_map:
+                related_map[result_chunk.hash] = {
+                    "chunk": result_chunk,
+                    "queries": [],
+                    "same_file": is_same_file,
+                    "max_score": score,
+                }
+            else:
+                related_map[result_chunk.hash]["max_score"] = max(
+                    related_map[result_chunk.hash]["max_score"], score
                 )
-                if is_same_file:
-                    index_diff = abs(
-                        result_chunk.get_metadata("index", 0)
-                        - source.get_metadata("index", 0)
-                    )
-                    if index_diff <= 1:
-                        continue
-
-                if result_chunk.hash not in related_map:
-                    related_map[result_chunk.hash] = {
-                        "chunk": result_chunk,
-                        "queries": [],
-                        "same_file": is_same_file,
-                        "max_score": score,
-                    }
-                else:
-                    related_map[result_chunk.hash]["max_score"] = max(
-                        related_map[result_chunk.hash]["max_score"], score
-                    )
 
-                related_map[result_chunk.hash]["queries"].append(query)
+            related_map[result_chunk.hash]["queries"].append(query)
 
+    @staticmethod
+    def _sorted_related(related_map: dict[str, dict]) -> list[dict]:
+        """Sort related chunks: most matching queries first, cross-file before
+        same-file, then max BM25 score — all descending."""
         return sorted(
             related_map.values(),
             key=lambda x: (len(x["queries"]), not x["same_file"], x["max_score"]),