benchmax 0.2.2.dev0__tar.gz → 0.2.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (63) hide show
  1. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/PKG-INFO +3 -1
  2. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/README.md +2 -0
  3. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/pyproject.toml +1 -1
  4. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/auth.py +49 -0
  5. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/bundle.py +3 -0
  6. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/__init__.py +2 -0
  7. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/environment.py +6 -0
  8. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/env.py +109 -0
  9. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/shared_types.py +11 -0
  10. benchmax-0.2.3/src/benchmax/rag/__init__.py +13 -0
  11. benchmax-0.2.3/src/benchmax/rag/embed.py +79 -0
  12. benchmax-0.2.3/src/benchmax/rag/env.py +699 -0
  13. benchmax-0.2.3/src/benchmax/rag/search.py +64 -0
  14. benchmax-0.2.3/tests/unit/rag/corpus/test_embed.py +126 -0
  15. benchmax-0.2.3/tests/unit/rag/test_rag_env.py +790 -0
  16. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/.gitignore +0 -0
  17. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/pytest.ini +0 -0
  18. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/README.md +0 -0
  19. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/base/README.md +0 -0
  20. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/base/__init__.py +0 -0
  21. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/base/dataset.py +0 -0
  22. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/base/env.py +0 -0
  23. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/base/openai_types.py +0 -0
  24. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/dataset.py +0 -0
  25. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/README.md +0 -0
  26. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/__init__.py +0 -0
  27. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/bundled_agent.py +0 -0
  28. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/credentials.py +0 -0
  29. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/dataset.py +0 -0
  30. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/dep_check.py +0 -0
  31. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/types.py +0 -0
  32. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/identity.py +0 -0
  33. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/logging.py +0 -0
  34. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/README.md +0 -0
  35. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/__init__.py +0 -0
  36. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/adaptive.py +0 -0
  37. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/deterministic.py +0 -0
  38. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/diversity.py +0 -0
  39. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/judge.py +0 -0
  40. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/prompts.py +0 -0
  41. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/rubric.py +0 -0
  42. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/scoring.py +0 -0
  43. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/conftest.py +0 -0
  44. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/bundle/test_artifact.py +0 -0
  45. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/bundle/test_source_capture.py +0 -0
  46. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/envs/test_base_dataset.py +0 -0
  47. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/envs/test_base_env_group.py +0 -0
  48. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/envs/test_contract_types.py +0 -0
  49. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/envs/test_environment_group.py +0 -0
  50. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/envs/test_example_id.py +0 -0
  51. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/fakes/model_server.py +0 -0
  52. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/harbor/test_bundled_agent.py +0 -0
  53. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/harbor/test_harbor_dataset.py +0 -0
  54. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/harbor/test_harbor_env.py +0 -0
  55. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/conftest.py +0 -0
  56. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/test_adaptive.py +0 -0
  57. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/test_deterministic.py +0 -0
  58. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/test_diversity.py +0 -0
  59. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/test_diversity_env.py +0 -0
  60. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/test_judge.py +0 -0
  61. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/test_rubric.py +0 -0
  62. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/test_rubric_rewards.py +0 -0
  63. {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/test_auth.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: benchmax
3
- Version: 0.2.2.dev0
3
+ Version: 0.2.3
4
4
  Summary: Platform-independent runtime for grouped LLM environments
5
5
  Author: benchmax Authors
6
6
  Classifier: Operating System :: OS Independent
@@ -98,6 +98,8 @@ environment class + constructor arguments
98
98
 
99
99
  benchmax creates the portable artifact so the same environment can be loaded outside the author's checkout. `castform` handles uploading the bundle, validating it remotely, and using it for training.
100
100
 
101
+ Remote runtimes install `pip_dependencies` without enabling prereleases globally. If a dependency graph needs a prerelease, list that package explicitly even when it would normally be transitive. For example, use `pip_dependencies=["parent-package==1.0.0", "transitive-package==2.0.0rc1"]`.
102
+
101
103
  ## further reading
102
104
 
103
105
  - [base environment guide](src/benchmax/envs/base/README.md)
@@ -83,6 +83,8 @@ environment class + constructor arguments
83
83
 
84
84
  benchmax creates the portable artifact so the same environment can be loaded outside the author's checkout. `castform` handles uploading the bundle, validating it remotely, and using it for training.
85
85
 
86
+ Remote runtimes install `pip_dependencies` without enabling prereleases globally. If a dependency graph needs a prerelease, list that package explicitly even when it would normally be transitive. For example, use `pip_dependencies=["parent-package==1.0.0", "transitive-package==2.0.0rc1"]`.
87
+
86
88
  ## further reading
87
89
 
88
90
  - [base environment guide](src/benchmax/envs/base/README.md)
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "benchmax"
3
- version = "0.2.2.dev0"
3
+ version = "0.2.3"
4
4
  description = "Platform-independent runtime for grouped LLM environments"
5
5
  readme = "README.md"
6
6
  authors = [{ name = "benchmax Authors" }]
@@ -16,6 +16,7 @@ from dataclasses import dataclass, field
16
16
  from typing import Protocol, runtime_checkable
17
17
 
18
18
  import httpx
19
+ from openai import AsyncOpenAI, OpenAI
19
20
 
20
21
  __all__ = [
21
22
  "InjectedAuth",
@@ -24,6 +25,8 @@ __all__ = [
24
25
  "RequestModelAuth",
25
26
  "StaticBearerAuth",
26
27
  "bind_model_auth",
28
+ "create_async_openai_client",
29
+ "create_openai_client",
27
30
  ]
28
31
 
29
32
 
@@ -99,6 +102,52 @@ class RequestModelAuth(httpx.Auth):
99
102
  yield request
100
103
 
101
104
 
105
+ def create_openai_client(
106
+ *,
107
+ model: str,
108
+ base_url: str,
109
+ auth: ModelAuth,
110
+ request_id: str,
111
+ max_retries: int = 2,
112
+ ) -> OpenAI:
113
+ """Create a synchronous OpenAI-compatible client with explicit auth."""
114
+
115
+ context = ModelRequestContext(
116
+ base_url=base_url,
117
+ model=model,
118
+ rollout_id=request_id,
119
+ )
120
+ return OpenAI(
121
+ base_url=base_url,
122
+ api_key="benchmax-explicit-auth",
123
+ http_client=httpx.Client(auth=RequestModelAuth(auth, context)),
124
+ max_retries=max_retries,
125
+ )
126
+
127
+
128
+ def create_async_openai_client(
129
+ *,
130
+ model: str,
131
+ base_url: str,
132
+ auth: ModelAuth,
133
+ request_id: str,
134
+ max_retries: int = 2,
135
+ ) -> AsyncOpenAI:
136
+ """Create an asynchronous OpenAI-compatible client with explicit auth."""
137
+
138
+ context = ModelRequestContext(
139
+ base_url=base_url,
140
+ model=model,
141
+ rollout_id=request_id,
142
+ )
143
+ return AsyncOpenAI(
144
+ base_url=base_url,
145
+ api_key="benchmax-explicit-auth",
146
+ http_client=httpx.AsyncClient(auth=RequestModelAuth(auth, context)),
147
+ max_retries=max_retries,
148
+ )
149
+
150
+
102
151
  def _resolve_headers_sync(
103
152
  auth: ModelAuth,
104
153
  context: ModelRequestContext,
@@ -203,6 +203,9 @@ def dump_bundle(
203
203
  env_class: A concrete Environment implementation.
204
204
  constructor_args: kwargs for ``env_class(**...)`` on load.
205
205
  pip_dependencies: Recorded in metadata. NOT installed by this call.
206
+ List every required prerelease explicitly, including a prerelease
207
+ that would otherwise only be a transitive dependency. Remote
208
+ runtimes do not enable prereleases globally.
206
209
  local_modules: Additional modules to pickle by value. Use this for
207
210
  local source outside the environment's own Python project. Such
208
211
  source otherwise fails loudly unless its installed distribution is
@@ -19,6 +19,7 @@ from benchmax.envs.shared_types import (
19
19
  RolloutFailure,
20
20
  RolloutOutcome,
21
21
  RolloutRequest,
22
+ ValidationDiagnostic,
22
23
  )
23
24
 
24
25
  __all__ = [
@@ -39,6 +40,7 @@ __all__ = [
39
40
  "RolloutFailure",
40
41
  "RolloutOutcome",
41
42
  "RolloutRequest",
43
+ "ValidationDiagnostic",
42
44
  "Tool",
43
45
  "StaticBearerAuth",
44
46
  "bind_model_auth",
@@ -15,6 +15,7 @@ from benchmax.envs.shared_types import (
15
15
  RolloutFailure,
16
16
  RolloutOutcome,
17
17
  RolloutRequest,
18
+ ValidationDiagnostic,
18
19
  )
19
20
 
20
21
  __all__ = ["Environment"]
@@ -42,6 +43,11 @@ class Environment[Payload, Attempt: RolloutAttempt](ABC):
42
43
 
43
44
  return False
44
45
 
46
+ def validation_diagnostics(self) -> Sequence[ValidationDiagnostic]:
47
+ """Return static warnings and errors for environment configuration."""
48
+
49
+ return ()
50
+
45
51
  @abstractmethod
46
52
  async def create_dataset(
47
53
  self,
@@ -26,6 +26,7 @@ from benchmax.envs.shared_types import (
26
26
  RolloutAttempt,
27
27
  RolloutOutcome,
28
28
  RolloutRequest,
29
+ ValidationDiagnostic,
29
30
  )
30
31
 
31
32
  if TYPE_CHECKING:
@@ -59,6 +60,72 @@ _TERMINATION_REASON_BY_EXCEPTION = {
59
60
  _HARNESS_REPORTED_TERMINATION_REASONS = frozenset(
60
61
  {"context_exceeded", "output_exceeded", "max_turns_exceeded", "tool_budget_exceeded"}
61
62
  )
63
+ _TRAINER_OWNED_MODEL_FIELDS = frozenset(
64
+ {
65
+ "temperature",
66
+ "top_p",
67
+ "top_k",
68
+ "presence_penalty",
69
+ "frequency_penalty",
70
+ "seed",
71
+ "stop",
72
+ }
73
+ )
74
+ _OUTPUT_CAP_FIELDS = frozenset({"max_tokens", "max_completion_tokens"})
75
+ _UNSUPPORTED_MODEL_FIELDS = frozenset(
76
+ {
77
+ "best_of",
78
+ "do_sample",
79
+ "function_call",
80
+ "functions",
81
+ "grammar",
82
+ "max_new_tokens",
83
+ "min_p",
84
+ "min_tokens",
85
+ "num_beams",
86
+ "prediction",
87
+ "reasoning_effort",
88
+ "repetition_penalty",
89
+ "typical_p",
90
+ "verbosity",
91
+ "web_search_options",
92
+ }
93
+ )
94
+
95
+
96
+ def _walk_model_controls(
97
+ value: Mapping[str, object],
98
+ prefix: str = "agent.kwargs",
99
+ ) -> Sequence[tuple[str, str, object]]:
100
+ controls: list[tuple[str, str, object]] = []
101
+ for field, child in value.items():
102
+ path = f"{prefix}.{field}"
103
+ controls.append((path, field, child))
104
+ if isinstance(child, Mapping):
105
+ controls.extend(_walk_model_controls(child, path))
106
+ return controls
107
+
108
+
109
+ def _unsupported_model_control(field: str, value: object) -> bool:
110
+ if field == "return_routed_experts":
111
+ return True
112
+ if field in _UNSUPPORTED_MODEL_FIELDS:
113
+ return value is not None
114
+ if field == "n":
115
+ return value not in (None, 1)
116
+ if field == "tool_choice":
117
+ return value not in (None, "auto")
118
+ if field == "logprobs":
119
+ return value not in (None, False)
120
+ if field == "top_logprobs":
121
+ return value is not None
122
+ if field == "parallel_tool_calls":
123
+ return value not in (None, True)
124
+ if field == "response_format":
125
+ return value not in (None, {}, {"type": "text"})
126
+ if field == "logit_bias":
127
+ return value not in (None, {})
128
+ return False
62
129
 
63
130
 
64
131
  class HarborEnv(Environment["TaskConfig", RolloutAttempt]):
@@ -74,6 +141,48 @@ class HarborEnv(Environment["TaskConfig", RolloutAttempt]):
74
141
 
75
142
  return self._requires_public_model_endpoint
76
143
 
144
+ def validation_diagnostics(self) -> Sequence[ValidationDiagnostic]:
145
+ """Flag Harbor harness controls that conflict with tracked training."""
146
+
147
+ agent = self._trial.agent
148
+ config = agent.config if isinstance(agent, BundledHarborAgent) else agent
149
+ kwargs = getattr(config, "kwargs", None)
150
+ if not isinstance(kwargs, Mapping):
151
+ return ()
152
+ diagnostics: list[ValidationDiagnostic] = []
153
+ for path, field, value in _walk_model_controls(kwargs):
154
+ if field in _OUTPUT_CAP_FIELDS:
155
+ diagnostics.append(
156
+ ValidationDiagnostic(
157
+ severity="warning",
158
+ code="harness_output_cap",
159
+ location=path,
160
+ message=(
161
+ f"{path} is a harness-requested output cap; Castform may clamp it "
162
+ "to the remaining trainer context budget"
163
+ ),
164
+ )
165
+ )
166
+ elif field in _TRAINER_OWNED_MODEL_FIELDS:
167
+ diagnostics.append(
168
+ ValidationDiagnostic(
169
+ severity="error",
170
+ code="trainer_owned_model_control",
171
+ location=path,
172
+ message=(f"{path} is trainer-owned and cannot be set by a Harbor harness"),
173
+ )
174
+ )
175
+ elif _unsupported_model_control(field, value):
176
+ diagnostics.append(
177
+ ValidationDiagnostic(
178
+ severity="error",
179
+ code="unsupported_model_control",
180
+ location=path,
181
+ message=f"{path}={value!r} is unsupported by tracked training sessions",
182
+ )
183
+ )
184
+ return tuple(diagnostics)
185
+
77
186
  def __init__(
78
187
  self,
79
188
  *,
@@ -13,12 +13,23 @@ __all__ = [
13
13
  "RolloutFailure",
14
14
  "RolloutOutcome",
15
15
  "RolloutRequest",
16
+ "ValidationDiagnostic",
16
17
  ]
17
18
 
18
19
  DatasetSplit = Literal["train", "eval"]
19
20
  type RewardMap = Mapping[str, float]
20
21
 
21
22
 
23
+ @dataclass(frozen=True, slots=True)
24
+ class ValidationDiagnostic:
25
+ """A configuration issue discoverable before executing a rollout."""
26
+
27
+ severity: Literal["warning", "error"]
28
+ code: str
29
+ message: str
30
+ location: str | None = None
31
+
32
+
22
33
  class RolloutFailure(RuntimeError): # noqa: N818 — public exported name
23
34
  """Operational rollout failure that should become a terminal outcome.
24
35
 
@@ -0,0 +1,13 @@
1
+ """Runtime building blocks for retrieval-augmented Benchmax environments."""
2
+
3
+ from benchmax.rag.embed import DEFAULT_EMBED_MODEL, OpenAIEmbedder
4
+ from benchmax.rag.env import RagEnv
5
+ from benchmax.rag.search import SearchClient, SearchResult
6
+
7
+ __all__ = [
8
+ "DEFAULT_EMBED_MODEL",
9
+ "OpenAIEmbedder",
10
+ "SearchClient",
11
+ "RagEnv",
12
+ "SearchResult",
13
+ ]
@@ -0,0 +1,79 @@
1
+ """OpenAI-compatible query embedding for RAG search clients.
2
+
3
+ The rollout-facing provider search clients (Turbopuffer / Pinecone / Chroma)
4
+ accept an optional async ``embed_fn``. Wiring one in makes vector / hybrid
5
+ retrieval work regardless of how the user's index was built — turbopuffer
6
+ vector/hybrid (no server-side embed), a pinecone index NOT on the hosted model,
7
+ or a non-cloud chroma collection.
8
+
9
+ ``qa-gen`` does NOT need this (it reads chunks directly); it's only for retrieval.
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ from dataclasses import dataclass
15
+
16
+ from benchmax.auth import ModelAuth, create_async_openai_client
17
+
18
+ # The served, non-hidden embeddings model in the llm-proxy catalog.
19
+ DEFAULT_EMBED_MODEL = "text-embedding-3-large"
20
+
21
+
22
+ @dataclass(frozen=True, slots=True)
23
+ class OpenAIEmbedder:
24
+ """Pickle-safe embedding callable with explicit call-time model auth.
25
+
26
+ ``auth`` follows the same contract as :class:`benchmax.rewards.Judge`.
27
+ Managed environments use ``InjectedAuth("embedding")``; customer-owned
28
+ endpoints use ``StaticBearerAuth``. No platform or SDK environment
29
+ credential is inferred.
30
+ """
31
+
32
+ model: str
33
+ base_url: str
34
+ auth: ModelAuth
35
+ timeout: float | None = 60.0
36
+ max_retries: int = 2
37
+
38
+ def __post_init__(self) -> None:
39
+ if not isinstance(self.auth, ModelAuth):
40
+ raise TypeError("embedding auth must implement ModelAuth")
41
+ if not isinstance(self.model, str) or not self.model.strip():
42
+ raise ValueError("embedding model must be non-empty")
43
+ if not isinstance(self.base_url, str) or not self.base_url.strip():
44
+ raise ValueError("embedding base_url must be non-empty")
45
+ if self.timeout is not None and (
46
+ isinstance(self.timeout, bool)
47
+ or not isinstance(self.timeout, (int, float))
48
+ or self.timeout <= 0
49
+ ):
50
+ raise ValueError("embedding timeout must be positive or None")
51
+ if (
52
+ isinstance(self.max_retries, bool)
53
+ or not isinstance(self.max_retries, int)
54
+ or self.max_retries < 0
55
+ ):
56
+ raise ValueError("embedding max_retries must be non-negative")
57
+
58
+ async def __call__(self, texts: list[str]) -> list[list[float]]:
59
+ if not isinstance(texts, list) or any(not isinstance(text, str) for text in texts):
60
+ raise TypeError("embedding input must be a list of strings")
61
+ if not texts:
62
+ return []
63
+
64
+ client = create_async_openai_client(
65
+ model=self.model,
66
+ base_url=self.base_url,
67
+ auth=self.auth,
68
+ request_id="rag-embedding",
69
+ max_retries=self.max_retries,
70
+ )
71
+ try:
72
+ response = await client.embeddings.create(
73
+ model=self.model,
74
+ input=texts,
75
+ timeout=self.timeout,
76
+ )
77
+ return [item.embedding for item in response.data]
78
+ finally:
79
+ await client.close()