diff --git a/anthropic-java-vertex/src/main/kotlin/com/anthropic/vertex/backends/VertexBackend.kt b/anthropic-java-vertex/src/main/kotlin/com/anthropic/vertex/backends/VertexBackend.kt index 6967781da..bb6ab0614 100644 --- a/anthropic-java-vertex/src/main/kotlin/com/anthropic/vertex/backends/VertexBackend.kt +++ b/anthropic-java-vertex/src/main/kotlin/com/anthropic/vertex/backends/VertexBackend.kt @@ -29,6 +29,7 @@ private constructor( @get:JvmName("googleCredentials") val googleCredentials: GoogleCredentials, @get:JvmName("region") val region: String, @get:JvmName("project") val project: String, + private val baseUrlOverride: String?, ) : Backend { private val jsonMapper = jsonMapper() @@ -45,12 +46,13 @@ private constructor( } override fun baseUrl(): String = - when (region) { - "global" -> "https://aiplatform.googleapis.com" - "us" -> "https://aiplatform.us.rep.googleapis.com" - "eu" -> "https://aiplatform.eu.rep.googleapis.com" - else -> "https://$region-aiplatform.googleapis.com" - } + baseUrlOverride + ?: when (region) { + "global" -> "https://aiplatform.googleapis.com" + "us" -> "https://aiplatform.us.rep.googleapis.com" + "eu" -> "https://aiplatform.eu.rep.googleapis.com" + else -> "https://$region-aiplatform.googleapis.com" + } override fun prepareRequest(request: HttpRequest): HttpRequest { val pathSegments = request.pathSegments @@ -150,8 +152,8 @@ private constructor( * * The credentials can be extracted from the environment and set on the builder by calling * [fromEnv] before calling [build] to create the [VertexBackend]. Alternatively, set the Google - * credentials, region and project explicitly via [googleCredentials], [region] and [project] - * before calling [build]. + * credentials, region and project explicitly via [googleCredentials], [region] and [project]. + * If requests should use a non-standard endpoint, set [baseUrl] before calling [build]. */ class Builder internal constructor() { private var googleCredentials: GoogleCredentials? = null @@ -160,6 +162,8 @@ private constructor( private var project: String? = null + private var baseUrlOverride: String? = null + fun fromEnv() = apply { try { googleCredentials = GoogleCredentials.getApplicationDefault() @@ -188,11 +192,15 @@ private constructor( fun project(project: String) = apply { this.project = project } + /** Overrides the base URL for requests. */ + fun baseUrl(baseUrl: String) = apply { this.baseUrlOverride = baseUrl } + fun build(): VertexBackend = VertexBackend( checkRequired("googleCredentials", googleCredentials), checkRequired("region", region), checkRequired("project", project), + baseUrlOverride, ) } } diff --git a/anthropic-java-vertex/src/test/kotlin/com/anthropic/vertex/backends/VertexBackendTest.kt b/anthropic-java-vertex/src/test/kotlin/com/anthropic/vertex/backends/VertexBackendTest.kt index 9eebe63a0..e1b3ba1f0 100644 --- a/anthropic-java-vertex/src/test/kotlin/com/anthropic/vertex/backends/VertexBackendTest.kt +++ b/anthropic-java-vertex/src/test/kotlin/com/anthropic/vertex/backends/VertexBackendTest.kt @@ -24,6 +24,7 @@ internal class VertexBackendTest { private const val GC_PROJECT = "vertex-project-12345-b6" private const val MODEL_ID = "claude-3-sonnet" private const val GC_REGION = "us-central1" + private const val CUSTOM_BASE_URL = "https://vertex.example.com" private const val ANTHROPIC_VERSION = "vertex-2023-10-16" } @@ -106,6 +107,19 @@ internal class VertexBackendTest { assertThat(backend2.baseUrl()).isEqualTo("https://$otherRegion-aiplatform.googleapis.com") } + @Test + fun baseUrlOverride() { + val backend = + VertexBackend.builder() + .googleCredentials(createCredentials()) + .region(GC_REGION) + .project(GC_PROJECT) + .baseUrl(CUSTOM_BASE_URL) + .build() + + assertThat(backend.baseUrl()).isEqualTo(CUSTOM_BASE_URL) + } + @Test fun prepareRequestMissingJSON() { val backend = createBackend()