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.
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/PKG-INFO +3 -1
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/README.md +2 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/pyproject.toml +1 -1
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/auth.py +49 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/bundle.py +3 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/__init__.py +2 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/environment.py +6 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/env.py +109 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/shared_types.py +11 -0
- benchmax-0.2.3/src/benchmax/rag/__init__.py +13 -0
- benchmax-0.2.3/src/benchmax/rag/embed.py +79 -0
- benchmax-0.2.3/src/benchmax/rag/env.py +699 -0
- benchmax-0.2.3/src/benchmax/rag/search.py +64 -0
- benchmax-0.2.3/tests/unit/rag/corpus/test_embed.py +126 -0
- benchmax-0.2.3/tests/unit/rag/test_rag_env.py +790 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/.gitignore +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/pytest.ini +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/README.md +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/base/README.md +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/base/__init__.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/base/dataset.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/base/env.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/base/openai_types.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/dataset.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/README.md +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/__init__.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/bundled_agent.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/credentials.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/dataset.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/dep_check.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/harbor/types.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/identity.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/envs/logging.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/README.md +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/__init__.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/adaptive.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/deterministic.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/diversity.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/judge.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/prompts.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/rubric.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/src/benchmax/rewards/scoring.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/conftest.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/bundle/test_artifact.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/bundle/test_source_capture.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/envs/test_base_dataset.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/envs/test_base_env_group.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/envs/test_contract_types.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/envs/test_environment_group.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/envs/test_example_id.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/fakes/model_server.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/harbor/test_bundled_agent.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/harbor/test_harbor_dataset.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/harbor/test_harbor_env.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/conftest.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/test_adaptive.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/test_deterministic.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/test_diversity.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/test_diversity_env.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/test_judge.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/test_rubric.py +0 -0
- {benchmax-0.2.2.dev0 → benchmax-0.2.3}/tests/unit/rewards/test_rubric_rewards.py +0 -0
- {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.
|
|
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)
|
|
@@ -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()
|