Skip to content

Commit 29bdded

Browse files
feat: Include LLM headers in ModelConfig
1 parent 02ec43b commit 29bdded

File tree

10 files changed

+363
-5
lines changed

10 files changed

+363
-5
lines changed

.stats.yml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
configured_endpoints: 8
2-
openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/browserbase%2Fstagehand-a4a5ea048bb50d06460d81d6828b53b12b19e9224121ee6338dcd1f0781e22a1.yml
3-
openapi_spec_hash: 9b81c0ae04576318d13d7a80d4ab7b5a
2+
openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/browserbase%2Fstagehand-bc309fd00fe0507f4cbe3bc77fa27d0fbffeaa6e71998778da34de42608a67e8.yml
3+
openapi_spec_hash: 1db1af5c1b068bba1d652102f4454668
44
config_hash: d6c6f623d03971bdba921650e5eb7e5f

stagehand-java-core/src/main/kotlin/com/browserbase/api/models/sessions/ModelConfig.kt

Lines changed: 138 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import com.browserbase.api.core.JsonField
88
import com.browserbase.api.core.JsonMissing
99
import com.browserbase.api.core.JsonValue
1010
import com.browserbase.api.core.checkRequired
11+
import com.browserbase.api.core.toImmutable
1112
import com.browserbase.api.errors.StagehandInvalidDataException
1213
import com.fasterxml.jackson.annotation.JsonAnyGetter
1314
import com.fasterxml.jackson.annotation.JsonAnySetter
@@ -24,6 +25,7 @@ private constructor(
2425
private val modelName: JsonField<String>,
2526
private val apiKey: JsonField<String>,
2627
private val baseUrl: JsonField<String>,
28+
private val headers: JsonField<Headers>,
2729
private val provider: JsonField<Provider>,
2830
private val additionalProperties: MutableMap<String, JsonValue>,
2931
) {
@@ -33,8 +35,9 @@ private constructor(
3335
@JsonProperty("modelName") @ExcludeMissing modelName: JsonField<String> = JsonMissing.of(),
3436
@JsonProperty("apiKey") @ExcludeMissing apiKey: JsonField<String> = JsonMissing.of(),
3537
@JsonProperty("baseURL") @ExcludeMissing baseUrl: JsonField<String> = JsonMissing.of(),
38+
@JsonProperty("headers") @ExcludeMissing headers: JsonField<Headers> = JsonMissing.of(),
3639
@JsonProperty("provider") @ExcludeMissing provider: JsonField<Provider> = JsonMissing.of(),
37-
) : this(modelName, apiKey, baseUrl, provider, mutableMapOf())
40+
) : this(modelName, apiKey, baseUrl, headers, provider, mutableMapOf())
3841

3942
/**
4043
* Model name string with provider prefix (e.g., 'openai/gpt-5-nano')
@@ -60,6 +63,14 @@ private constructor(
6063
*/
6164
fun baseUrl(): Optional<String> = baseUrl.getOptional("baseURL")
6265

66+
/**
67+
* Custom headers sent with every request to the model provider
68+
*
69+
* @throws StagehandInvalidDataException if the JSON field has an unexpected type (e.g. if the
70+
* server responded with an unexpected value).
71+
*/
72+
fun headers(): Optional<Headers> = headers.getOptional("headers")
73+
6374
/**
6475
* AI provider for the model (or provide a baseURL endpoint instead)
6576
*
@@ -89,6 +100,13 @@ private constructor(
89100
*/
90101
@JsonProperty("baseURL") @ExcludeMissing fun _baseUrl(): JsonField<String> = baseUrl
91102

103+
/**
104+
* Returns the raw JSON value of [headers].
105+
*
106+
* Unlike [headers], this method doesn't throw if the JSON field has an unexpected type.
107+
*/
108+
@JsonProperty("headers") @ExcludeMissing fun _headers(): JsonField<Headers> = headers
109+
92110
/**
93111
* Returns the raw JSON value of [provider].
94112
*
@@ -127,6 +145,7 @@ private constructor(
127145
private var modelName: JsonField<String>? = null
128146
private var apiKey: JsonField<String> = JsonMissing.of()
129147
private var baseUrl: JsonField<String> = JsonMissing.of()
148+
private var headers: JsonField<Headers> = JsonMissing.of()
130149
private var provider: JsonField<Provider> = JsonMissing.of()
131150
private var additionalProperties: MutableMap<String, JsonValue> = mutableMapOf()
132151

@@ -135,6 +154,7 @@ private constructor(
135154
modelName = modelConfig.modelName
136155
apiKey = modelConfig.apiKey
137156
baseUrl = modelConfig.baseUrl
157+
headers = modelConfig.headers
138158
provider = modelConfig.provider
139159
additionalProperties = modelConfig.additionalProperties.toMutableMap()
140160
}
@@ -173,6 +193,17 @@ private constructor(
173193
*/
174194
fun baseUrl(baseUrl: JsonField<String>) = apply { this.baseUrl = baseUrl }
175195

196+
/** Custom headers sent with every request to the model provider */
197+
fun headers(headers: Headers) = headers(JsonField.of(headers))
198+
199+
/**
200+
* Sets [Builder.headers] to an arbitrary JSON value.
201+
*
202+
* You should usually call [Builder.headers] with a well-typed [Headers] value instead. This
203+
* method is primarily for setting the field to an undocumented or not yet supported value.
204+
*/
205+
fun headers(headers: JsonField<Headers>) = apply { this.headers = headers }
206+
176207
/** AI provider for the model (or provide a baseURL endpoint instead) */
177208
fun provider(provider: Provider) = provider(JsonField.of(provider))
178209

@@ -221,6 +252,7 @@ private constructor(
221252
checkRequired("modelName", modelName),
222253
apiKey,
223254
baseUrl,
255+
headers,
224256
provider,
225257
additionalProperties.toMutableMap(),
226258
)
@@ -236,6 +268,7 @@ private constructor(
236268
modelName()
237269
apiKey()
238270
baseUrl()
271+
headers().ifPresent { it.validate() }
239272
provider().ifPresent { it.validate() }
240273
validated = true
241274
}
@@ -258,8 +291,109 @@ private constructor(
258291
(if (modelName.asKnown().isPresent) 1 else 0) +
259292
(if (apiKey.asKnown().isPresent) 1 else 0) +
260293
(if (baseUrl.asKnown().isPresent) 1 else 0) +
294+
(headers.asKnown().getOrNull()?.validity() ?: 0) +
261295
(provider.asKnown().getOrNull()?.validity() ?: 0)
262296

297+
/** Custom headers sent with every request to the model provider */
298+
class Headers
299+
@JsonCreator
300+
private constructor(
301+
@com.fasterxml.jackson.annotation.JsonValue
302+
private val additionalProperties: Map<String, JsonValue>
303+
) {
304+
305+
@JsonAnyGetter
306+
@ExcludeMissing
307+
fun _additionalProperties(): Map<String, JsonValue> = additionalProperties
308+
309+
fun toBuilder() = Builder().from(this)
310+
311+
companion object {
312+
313+
/** Returns a mutable builder for constructing an instance of [Headers]. */
314+
@JvmStatic fun builder() = Builder()
315+
}
316+
317+
/** A builder for [Headers]. */
318+
class Builder internal constructor() {
319+
320+
private var additionalProperties: MutableMap<String, JsonValue> = mutableMapOf()
321+
322+
@JvmSynthetic
323+
internal fun from(headers: Headers) = apply {
324+
additionalProperties = headers.additionalProperties.toMutableMap()
325+
}
326+
327+
fun additionalProperties(additionalProperties: Map<String, JsonValue>) = apply {
328+
this.additionalProperties.clear()
329+
putAllAdditionalProperties(additionalProperties)
330+
}
331+
332+
fun putAdditionalProperty(key: String, value: JsonValue) = apply {
333+
additionalProperties.put(key, value)
334+
}
335+
336+
fun putAllAdditionalProperties(additionalProperties: Map<String, JsonValue>) = apply {
337+
this.additionalProperties.putAll(additionalProperties)
338+
}
339+
340+
fun removeAdditionalProperty(key: String) = apply { additionalProperties.remove(key) }
341+
342+
fun removeAllAdditionalProperties(keys: Set<String>) = apply {
343+
keys.forEach(::removeAdditionalProperty)
344+
}
345+
346+
/**
347+
* Returns an immutable instance of [Headers].
348+
*
349+
* Further updates to this [Builder] will not mutate the returned instance.
350+
*/
351+
fun build(): Headers = Headers(additionalProperties.toImmutable())
352+
}
353+
354+
private var validated: Boolean = false
355+
356+
fun validate(): Headers = apply {
357+
if (validated) {
358+
return@apply
359+
}
360+
361+
validated = true
362+
}
363+
364+
fun isValid(): Boolean =
365+
try {
366+
validate()
367+
true
368+
} catch (e: StagehandInvalidDataException) {
369+
false
370+
}
371+
372+
/**
373+
* Returns a score indicating how many valid values are contained in this object
374+
* recursively.
375+
*
376+
* Used for best match union deserialization.
377+
*/
378+
@JvmSynthetic
379+
internal fun validity(): Int =
380+
additionalProperties.count { (_, value) -> !value.isNull() && !value.isMissing() }
381+
382+
override fun equals(other: Any?): Boolean {
383+
if (this === other) {
384+
return true
385+
}
386+
387+
return other is Headers && additionalProperties == other.additionalProperties
388+
}
389+
390+
private val hashCode: Int by lazy { Objects.hash(additionalProperties) }
391+
392+
override fun hashCode(): Int = hashCode
393+
394+
override fun toString() = "Headers{additionalProperties=$additionalProperties}"
395+
}
396+
263397
/** AI provider for the model (or provide a baseURL endpoint instead) */
264398
class Provider @JsonCreator private constructor(private val value: JsonField<String>) : Enum {
265399

@@ -415,16 +549,17 @@ private constructor(
415549
modelName == other.modelName &&
416550
apiKey == other.apiKey &&
417551
baseUrl == other.baseUrl &&
552+
headers == other.headers &&
418553
provider == other.provider &&
419554
additionalProperties == other.additionalProperties
420555
}
421556

422557
private val hashCode: Int by lazy {
423-
Objects.hash(modelName, apiKey, baseUrl, provider, additionalProperties)
558+
Objects.hash(modelName, apiKey, baseUrl, headers, provider, additionalProperties)
424559
}
425560

426561
override fun hashCode(): Int = hashCode
427562

428563
override fun toString() =
429-
"ModelConfig{modelName=$modelName, apiKey=$apiKey, baseUrl=$baseUrl, provider=$provider, additionalProperties=$additionalProperties}"
564+
"ModelConfig{modelName=$modelName, apiKey=$apiKey, baseUrl=$baseUrl, headers=$headers, provider=$provider, additionalProperties=$additionalProperties}"
430565
}

stagehand-java-core/src/test/kotlin/com/browserbase/api/models/sessions/ModelConfigTest.kt

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
package com.browserbase.api.models.sessions
44

5+
import com.browserbase.api.core.JsonValue
56
import com.browserbase.api.core.jsonMapper
67
import com.fasterxml.jackson.module.kotlin.jacksonTypeRef
78
import org.assertj.core.api.Assertions.assertThat
@@ -16,12 +17,23 @@ internal class ModelConfigTest {
1617
.modelName("openai/gpt-5-nano")
1718
.apiKey("sk-some-openai-api-key")
1819
.baseUrl("https://api.openai.com/v1")
20+
.headers(
21+
ModelConfig.Headers.builder()
22+
.putAdditionalProperty("foo", JsonValue.from("string"))
23+
.build()
24+
)
1925
.provider(ModelConfig.Provider.OPENAI)
2026
.build()
2127

2228
assertThat(modelConfig.modelName()).isEqualTo("openai/gpt-5-nano")
2329
assertThat(modelConfig.apiKey()).contains("sk-some-openai-api-key")
2430
assertThat(modelConfig.baseUrl()).contains("https://api.openai.com/v1")
31+
assertThat(modelConfig.headers())
32+
.contains(
33+
ModelConfig.Headers.builder()
34+
.putAdditionalProperty("foo", JsonValue.from("string"))
35+
.build()
36+
)
2537
assertThat(modelConfig.provider()).contains(ModelConfig.Provider.OPENAI)
2638
}
2739

@@ -33,6 +45,11 @@ internal class ModelConfigTest {
3345
.modelName("openai/gpt-5-nano")
3446
.apiKey("sk-some-openai-api-key")
3547
.baseUrl("https://api.openai.com/v1")
48+
.headers(
49+
ModelConfig.Headers.builder()
50+
.putAdditionalProperty("foo", JsonValue.from("string"))
51+
.build()
52+
)
3653
.provider(ModelConfig.Provider.OPENAI)
3754
.build()
3855

stagehand-java-core/src/test/kotlin/com/browserbase/api/models/sessions/SessionActParamsTest.kt

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,11 @@ internal class SessionActParamsTest {
2323
.modelName("openai/gpt-5-nano")
2424
.apiKey("sk-some-openai-api-key")
2525
.baseUrl("https://api.openai.com/v1")
26+
.headers(
27+
ModelConfig.Headers.builder()
28+
.putAdditionalProperty("foo", JsonValue.from("string"))
29+
.build()
30+
)
2631
.provider(ModelConfig.Provider.OPENAI)
2732
.build()
2833
)
@@ -74,6 +79,11 @@ internal class SessionActParamsTest {
7479
.modelName("openai/gpt-5-nano")
7580
.apiKey("sk-some-openai-api-key")
7681
.baseUrl("https://api.openai.com/v1")
82+
.headers(
83+
ModelConfig.Headers.builder()
84+
.putAdditionalProperty("foo", JsonValue.from("string"))
85+
.build()
86+
)
7787
.provider(ModelConfig.Provider.OPENAI)
7888
.build()
7989
)
@@ -129,6 +139,11 @@ internal class SessionActParamsTest {
129139
.modelName("openai/gpt-5-nano")
130140
.apiKey("sk-some-openai-api-key")
131141
.baseUrl("https://api.openai.com/v1")
142+
.headers(
143+
ModelConfig.Headers.builder()
144+
.putAdditionalProperty("foo", JsonValue.from("string"))
145+
.build()
146+
)
132147
.provider(ModelConfig.Provider.OPENAI)
133148
.build()
134149
)
@@ -164,6 +179,11 @@ internal class SessionActParamsTest {
164179
.modelName("openai/gpt-5-nano")
165180
.apiKey("sk-some-openai-api-key")
166181
.baseUrl("https://api.openai.com/v1")
182+
.headers(
183+
ModelConfig.Headers.builder()
184+
.putAdditionalProperty("foo", JsonValue.from("string"))
185+
.build()
186+
)
167187
.provider(ModelConfig.Provider.OPENAI)
168188
.build()
169189
)

0 commit comments

Comments
 (0)