benchmax 0.1.2.dev31__py3-none-any.whl → 0.1.2.dev33__py3-none-any.whl

benchmax/envs/postgres_search/search_env.py

@@ -285,14 +285,8 @@ tags. Cite your sources inline using [Source: <source_id>] next to each claim.
             if not text.strip():
                 return zeros
 
-            # No final <answer> block → no answer to score. Return all-zero
-            # rewards so conciseness / citations / efficiency can't accrue
-            # from reasoning or tool-call text alone.
-            answer = extract_answer_block(text)
-            if not answer:
-                return zeros
-
             t = task or {}
+            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/envs/reward_helpers.py

@@ -82,16 +82,9 @@ def extract_completion_text(completion: str | list[dict[str, Any]]) -> str:
 
 
 def extract_answer_block(text: str) -> str:
-    """Extract content from ``<answer>`` tags.
-
-    Returns the (stripped) tag contents when an ``<answer>…</answer>`` block
-    is present, otherwise ``""``. A missing answer block is treated as "no
-    final answer" rather than silently falling back to the full completion —
-    consumers can gate rewards on a non-empty result. ``<answer></answer>``
-    likewise yields ``""``.
-    """
+    """Extract content from <answer> tags, or return full text."""
     match = _ANSWER_TAG_RE.search(text or "")
-    return match.group(1).strip() if match else ""
+    return (match.group(1) if match else text).strip()
 
 
 def clip01(value: Any) -> float:
@@ -169,10 +162,8 @@ def citation_score(
                 ref_ids.add(norm_sid)
             break
 
-    if not cited:
+    if not cited or not ref_ids:
         return {"precision": 0.0, "recall": 0.0}
-    if not ref_ids:
-        return {"precision": 1.0, "recall": 0.0}
 
     precision = len(cited & ref_ids) / len(cited)
     recall = len(cited & ref_ids) / len(ref_ids)

benchmax/rag/corpus/chroma/client.py

@@ -16,6 +16,13 @@ from typing import Any
 # Sparse-key name used when setting up BM25 schema
 BM25_KEY = "bm25_embedding"
 
+# Embedding functions that run server-side on Chroma Cloud (embed.trychroma.com)
+# — querying a collection that uses one never downloads a model. Everything else
+# (default all-MiniLM, sentence-transformers / HF / Ollama / ONNX locals,
+# third-party API EFs, or no EF) is treated as unsafe. Add hosted names here as
+# they are verified server-side.
+_SERVER_SIDE_EF_NAMES = frozenset({"chroma-cloud-qwen"})
+
 
 def has_search_api() -> bool:
     """Return True when the chromadb package exposes the Search API."""
@@ -176,6 +183,29 @@ class ChromaClient:
 
         return self._collection
 
+    def dense_embed_is_safe(self) -> bool:
+        """True when a dense (vector) query embeds WITHOUT downloading a model.
+
+        Safe only when we can produce vectors without a client-side model
+        download: either a caller-supplied ``embed_fn``, or a Chroma-hosted
+        server-side embedding function (embeds at embed.trychroma.com). Every
+        other embedder — chromadb's default all-MiniLM, sentence-transformers /
+        HuggingFace / Ollama / ONNX locals, third-party API EFs we lack keys
+        for, or no EF at all — is treated as UNSAFE, so callers refuse the dense
+        path rather than trigger a model download. Conservative by design: an
+        unknown embedder is unsafe.
+        """
+        if self.embed_fn is not None:
+            return True
+        col = self._collection
+        if col is None:
+            return False
+        try:
+            ef = (col._model.configuration_json or {}).get("embedding_function") or {}
+        except Exception:
+            return False
+        return ef.get("name") in _SERVER_SIDE_EF_NAMES
+
     @staticmethod
     def _repair_cloud_embedding_function(collection: Any) -> None:
         """Attach a working EF when chromadb can't rebuild a Cloud hosted one.

benchmax/rag/corpus/chroma/search.py

@@ -10,6 +10,9 @@ from collections.abc import Callable
 from typing import Any
 
 from benchmax.platform.credentials import TokenProvider, as_token_provider, env_token
+from benchmax.rag.corpus.search_schema.search_exceptions import (
+    LocalEmbeddingDownloadDisallowedError,
+)
 
 
 class ChromaSearch:
@@ -113,19 +116,33 @@ class ChromaSearch:
     ) -> list[dict[str, Any]]:
         """Search and return structured results."""
         client = self._get_client()
-
-        if mode == "auto":
-            modes = client.modes
+        # Initialize the collection first so capabilities reflect the real index
+        # (BM25 downgrade) and the embedder config is readable below.
+        client.get_collection()
+        modes = client.modes
+        has_lexical = "lexical" in modes
+
+        # Never download a client-side embedding model at inference/rollout time.
+        # When a dense embed isn't safe — no embed_fn and no Chroma-hosted
+        # server-side embedding function — use the BM25 lexical index if the
+        # collection has one, otherwise refuse rather than fetch all-MiniLM.
+        if not client.dense_embed_is_safe():
+            if not has_lexical:
+                raise LocalEmbeddingDownloadDisallowedError(
+                    "chroma", self._collection_name
+                )
+            mode = "lexical"
+        elif mode == "auto":
             if "hybrid" in modes:
                 mode = "hybrid"
-            elif "lexical" in modes:
+            elif has_lexical:
                 mode = "lexical"
             else:
                 mode = "vector"
-        elif mode not in client.modes:
+        elif mode not in modes:
             raise ValueError(
                 f"ChromaSearch does not support mode '{mode}'. "
-                f"Available modes: {sorted(client.modes)}"
+                f"Available modes: {sorted(modes)}"
             )
 
         if client.search_api and mode in ("lexical", "hybrid"):

benchmax/rag/corpus/chroma/source.py

@@ -17,6 +17,7 @@ from tqdm.auto import tqdm
 from benchmax.rag.chunkers.models import Chunk, ChunkCollection
 from benchmax.rag.corpus.search_schema.search_exceptions import (
     InvalidSearchSpecError,
+    LocalEmbeddingDownloadDisallowedError,
     UnsupportedSearchModeError,
 )
 from benchmax.rag.corpus.search_schema.search_types import (
@@ -642,23 +643,30 @@ class ChromaChunkSource:
         # lack a BM25 index, in which case modes was downgraded to vector-only.
         modes = self._current_modes()
 
-        # Pick mode. "hybrid"/None use the best available strategy and KEEP
-        # lexical enabled as a fallback: hybrid = dense + sparse, and when we
-        # can't produce dense query vectors (no embed_fn, the usual remote case)
-        # the per-query loop below degrades to the sparse/lexical leg — which
-        # needs no embedding. Only an explicit "vector" disables lexical; that's
-        # the dense-only recovery path a caller uses after a lexical/hybrid
-        # failure. (Disabling lexical for "hybrid" silently forced vector search,
-        # which made remote collections dense-embed every query — slow, and on a
-        # default-EF collection it pulls the all-MiniLM model.)
-        if mode == "vector":
-            use_hybrid = use_lexical = False
+        has_lexical = "lexical" in modes
+        has_hybrid = "hybrid" in modes
+
+        # Hard rule: never let chromadb embed a query with a client-side model
+        # (it downloads all-MiniLM and crawls in constrained executors). When a
+        # dense embed isn't safe — no embed_fn and no Chroma-hosted server-side
+        # embedding function — use the BM25 lexical index if the collection has
+        # one, otherwise refuse. This covers every requested mode, including the
+        # linker's "inference" preference for vector.
+        if not self._chroma.dense_embed_is_safe():
+            if not has_lexical:
+                raise LocalEmbeddingDownloadDisallowedError(
+                    "chroma", self._chroma.collection_name
+                )
+            use_hybrid = False
+            use_lexical = True
         elif mode == "lexical":
             use_hybrid = False
-            use_lexical = "lexical" in modes
+            use_lexical = has_lexical
+        elif mode == "vector":
+            use_hybrid = use_lexical = False
         else:  # "hybrid", None, or unrecognized -> best available
-            use_hybrid = "hybrid" in modes
-            use_lexical = "lexical" in modes
+            use_hybrid = has_hybrid
+            use_lexical = has_lexical
 
         # Batch-embed all queries when embed_fn available and vectors needed
         vectors: list[list[float]] | None = None

benchmax/rag/corpus/search_schema/search_exceptions.py

@@ -43,3 +43,21 @@ class UnsupportedSearchModeError(ValueError):
             f"[{backend}] unsupported search mode '{mode}'. "
             f"Supported modes: {sorted(supported_modes)}"
         )
+
+
+class LocalEmbeddingDownloadDisallowedError(RuntimeError):
+    """Raised when serving a search would download a client-side embedding model.
+
+    The collection has no server-side (hosted) embedding function and no BM25
+    index, and the caller supplied no ``embed_fn`` — so embedding a text query
+    would make chromadb download and run a local model (e.g. all-MiniLM). We
+    refuse rather than trigger that download.
+    """
+
+    def __init__(self, backend: str, collection: str):
+        super().__init__(
+            f"[{backend}] collection {collection!r} has no server-side embedding "
+            "function and no BM25 index, so search would download a local "
+            "embedding model. Re-ingest the corpus with a hosted embedder "
+            "(chroma-cloud-qwen) or a BM25 index, or supply an embed_fn."
+        )

{benchmax-0.1.2.dev31.dist-info → benchmax-0.1.2.dev33.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: benchmax
-Version: 0.1.2.dev31
+Version: 0.1.2.dev33
 Summary: Framework-Agnostic RL Environments for LLM Fine-Tuning
 Author: castie@castform.com
 Classifier: Programming Language :: Python :: 3

{benchmax-0.1.2.dev31.dist-info → benchmax-0.1.2.dev33.dist-info}/RECORD

@@ -5,7 +5,7 @@ benchmax/envs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 benchmax/envs/base_env.py,sha256=FoUgWsNGeNpTHeOop0bklRjLfHA90Yi7MW8zCaNh_V0,8976
 benchmax/envs/example_id.py,sha256=WU967Pt2kFvn-W4k5BC6BvKyrTEqioLr7IyWZ3RjGgU,5685
 benchmax/envs/logging.py,sha256=QnXADCp0vWoV_-MK91yX5OFu6GwgIE98dvhaQTPawqQ,5053
-benchmax/envs/reward_helpers.py,sha256=lKbyTvJYU2JoiFItFkUPX6aWwp6JmmgXC76FBaf2rBQ,7740
+benchmax/envs/reward_helpers.py,sha256=-pDqYBazvum8cc8KX7Q_Z0C-Daf3_4TVZuWt-ywhqyY,7364
 benchmax/envs/types.py,sha256=sGKKibQJZQj9RYkFpB3vaUY75tdoHet8yUmdzpZ0SVk,4389
 benchmax/envs/crm/crm_env.py,sha256=ltUtpA45YB_A_hYEpjFTp0nZKwkUvvLSLOAVkaUNz9E,4707
 benchmax/envs/crm/workdir/reward_fn.py,sha256=RY_iy347j79xX4gyCGI7WS0qPmut8Th2rqOiErVbDro,5439
@@ -32,7 +32,7 @@ benchmax/envs/mcp/provisioners/skypilot_provisioner.py,sha256=ACHnzNZE7GfL1WIWf7
 benchmax/envs/mcp/provisioners/utils.py,sha256=ORWJKtPzeS-IdD35p8aZyLMG2RxiB9BAFmU-0pVqiWw,3467
 benchmax/envs/postgres_search/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 benchmax/envs/postgres_search/linker_env.py,sha256=B3cn0TpiqgrYL5NvOQYW3Yxy5DdxPw1kmIgqDs-8Buo,8535
-benchmax/envs/postgres_search/search_env.py,sha256=ldfSNCQonbQZimO7rCO1Jc6im7ff2d0-TLvwryXUOXQ,20181
+benchmax/envs/postgres_search/search_env.py,sha256=IWpqbFr4hjaN_DzdeRchvwvf9qVj5Ut5D-eOsGPyWKQ,19917
 benchmax/envs/telestich/example.py,sha256=cqHIBjD8g7H4-nmspWSKRB2rxeKPOIwkLn136Y04KfQ,28680
 benchmax/envs/telestich/telestich_env.py,sha256=6p6GeyV-9ZIXrAX8zssMFjJgevkV5PfDLMZlslqO8js,61966
 benchmax/envs/wikipedia/utils.py,sha256=YDlxpMfwiVpfMpiZet4kWoeKqNbgTBxeWVEYg5QY3Qs,2879
@@ -64,11 +64,11 @@ benchmax/rag/corpus/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKV
 benchmax/rag/corpus/search_client.py,sha256=171IqQriU6kuQqvSCDgNwOT8SR5pxUPMfCifarrgrFg,1859
 benchmax/rag/corpus/source.py,sha256=dnmReLC8mccHDkg8ZytfXa4AFXrRMCg9v8E2UuVxt8E,4183
 benchmax/rag/corpus/chroma/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-benchmax/rag/corpus/chroma/client.py,sha256=TGf_YEgVBH8p-PLF7QcuSIHkkifrChXPKJB2ENR9OiM,19361
+benchmax/rag/corpus/chroma/client.py,sha256=cYZZKQG09u_VfyjsP6UdCBh-RRNGKa9XisBN4OEejQk,20839
 benchmax/rag/corpus/chroma/files.py,sha256=hSP-J2osPNBAvMZHOWipMVXaWN4tila_tsQaTEPNzgc,5567
 benchmax/rag/corpus/chroma/filter_mapper.py,sha256=Y1FzDwDDg15LZ0-Uh1jzOVcSORiVUy5f1qiaVky3pJI,5074
-benchmax/rag/corpus/chroma/search.py,sha256=4kU1WMwsWQrN03ctVIPdXZoHyZa3jso2fKzcsc5uYr0,6824
-benchmax/rag/corpus/chroma/source.py,sha256=ZOLj_VfixBcB_VE8YLF6X3sKY6XHYupprHNHVxy1hH8,30295
+benchmax/rag/corpus/chroma/search.py,sha256=iO8fBPk50vG3NmkCmAJ2tKnjP_wKnymV3fbfLjkIAJ8,7688
+benchmax/rag/corpus/chroma/source.py,sha256=0azMLUvZS9g4jvxv_KxsPa3-ArQW5WHCq77CQh-qmqY,30440
 benchmax/rag/corpus/pinecone/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 benchmax/rag/corpus/pinecone/files.py,sha256=lhas7-mQ622Ku36QvOavXguBweJyYl78wXIeb_LNqig,5728
 benchmax/rag/corpus/pinecone/filter_mapper.py,sha256=exJ3G34QKeQo1rQ8Pu-iGL0XDXVxCW5dc3q0QoYfCo0,6454
@@ -85,7 +85,7 @@ benchmax/rag/corpus/postgres/source.py,sha256=6ptGHatOscYih42MZ9Wt8MQOrcIEQiJ1X5
 benchmax/rag/corpus/search_schema/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 benchmax/rag/corpus/search_schema/builders.py,sha256=qAMiEOGOLR7xSXWFf12KqzYlrwBZchU_78vkRcOKa8k,1764
 benchmax/rag/corpus/search_schema/dsl_parser.py,sha256=vMijm_nRKztIrsVQP-0OySuCKnrBsbUzet_pwwlU1T8,1586
-benchmax/rag/corpus/search_schema/search_exceptions.py,sha256=vxhJQa7UFHduXDt225onA_R_UWcvGlj3NiS5TFR3M7w,1578
+benchmax/rag/corpus/search_schema/search_exceptions.py,sha256=1ccbLnDAuSMxUnjtyBt-5iXwoKjI3xaZvk9xplCyNFw,2413
 benchmax/rag/corpus/search_schema/search_types.py,sha256=UTkteugSx5OigDRZ8Xqe6itxLUXj2sVeIVxtYbnXGSg,5831
 benchmax/rag/corpus/turbopuffer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 benchmax/rag/corpus/turbopuffer/files.py,sha256=DP80-3NmdyOD34fyQxlzovpLRs_UU1ezQ7PItpY2Nlk,5807
@@ -160,9 +160,9 @@ benchmax/traces/braintrust/adapter.py,sha256=KTeN9qKLwZJJ8TY-KtSudd4J3_nySz1bRts
 benchmax/traces/braintrust/message_extraction.py,sha256=seh3eM_qd9FUPmGOEMChUq_UAMtaIQHYSYDttMgY1go,8409
 benchmax/utils/__init__.py,sha256=FWJVm6jt0m57HS-84bgrb2M-c_EFhf60rWayioUGges,402
 benchmax/utils/checkpoint.py,sha256=htIw9iYjUUHpJqLLZ0y6K4_UYYAkZIx3vdQVY7juKDw,3148
-benchmax-0.1.2.dev31.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-benchmax-0.1.2.dev31.dist-info/METADATA,sha256=A4ysjDpcSRx1jMeFb5xw0r9CinUiGPoz4Vm4x9FXqUQ,2775
-benchmax-0.1.2.dev31.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
-benchmax-0.1.2.dev31.dist-info/entry_points.txt,sha256=qtjqAQsHIwRIaLzwAhGTiRvI91CynwcUO5G95uQuDR4,47
-benchmax-0.1.2.dev31.dist-info/top_level.txt,sha256=ryj4zoahvAKL3BnxOpfJNfyIzhvlED9KJ3Q3k4bb9jc,9
-benchmax-0.1.2.dev31.dist-info/RECORD,,
+benchmax-0.1.2.dev33.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+benchmax-0.1.2.dev33.dist-info/METADATA,sha256=X5P1IBK9INVKaO8xzBqoW8CQYQ2VIVD9IkaQV4tVjFQ,2775
+benchmax-0.1.2.dev33.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+benchmax-0.1.2.dev33.dist-info/entry_points.txt,sha256=qtjqAQsHIwRIaLzwAhGTiRvI91CynwcUO5G95uQuDR4,47
+benchmax-0.1.2.dev33.dist-info/top_level.txt,sha256=ryj4zoahvAKL3BnxOpfJNfyIzhvlED9KJ3Q3k4bb9jc,9
+benchmax-0.1.2.dev33.dist-info/RECORD,,