feat: AWS Bedrock and Google Vertex adapters (P-b, stacked on #22) #23
No files matched your search
@@ -43,6 +43,8 @@ Always use the wrapper: Maven 3.10 cannot read the parent POM. The committed def
|
||||
- `engine/ProviderService` (active provider from `assistant_provider` rows, key references, prices), `UsageRecorder` (one `assistant_usage` row per model call), `ConnectionTest`.
|
||||
- `qbit-assistant-provider-anthropic`: `AnthropicModelProvider` and `AnthropicMapping` (anthropic-java streaming).
|
||||
- `qbit-assistant-provider-openai`: OpenAI-compatible Chat Completions, model discovery and health using JDK HTTP and Jackson.
|
||||
- `qbit-assistant-provider-bedrock`: Claude through the Anthropic Bedrock backend, AWS default credentials and SigV4; IRSA/web-identity hosts supply STS.
|
||||
- `qbit-assistant-provider-vertex`: Claude through the Anthropic Vertex backend and Gemini through OpenAI-compatible Chat Completions, with refreshed Google ADC tokens.
|
||||
|
||||
## Invariants (do not break)
|
||||
|
||||
@@ -91,4 +93,4 @@ QQQ style:
|
||||
- `develop` publishes SNAPSHOTs.
|
||||
- `release/*` publishes RCs; the orb bumps `<revision>` in the root POM.
|
||||
- `v*` tags publish releases.
|
||||
- Both modules and the parent POM publish. The flatten plugin resolves `${revision}`.
|
||||
- All five modules and the parent POM publish. The flatten plugin resolves `${revision}`.
|
||||
@@ -4,6 +4,10 @@ All notable changes to this project are documented here. The format follows Keep
|
||||
|
||||
## [0.1.0] - Unreleased
|
||||
|
||||
- Add optional AWS Bedrock (Claude) and Google Vertex (Claude/Gemini) adapters using server credentials, static model discovery and allow-listed endpoint overrides. AWS IRSA/web-identity hosts supply STS themselves; no additional dependency is bundled.
|
||||
- Add the provider `cloudProject` column, cloud settings validation, injected Anthropic clients and per-request OpenAI bearer-token suppliers. Vertex Gemini prefixes model ids and omits strict tools.
|
||||
- Add localhost cloud contract tests and golden tool turns, including AWS binary event-stream framing, SigV4, token refresh and cloud HTTP error mapping.
|
||||
|
||||
- Add named, server-only CA bundles for HTTPS providers, additive JVM trust and sanitized TLS diagnostics; tolerate repeated/missing streamed tool identities and invalid arguments, with a golden tool round-trip test.
|
||||
- Add the provider/model registry, per-model capabilities and prices, keyless local providers, discovery and health, and provider/model connection tests.
|
||||
- Add the optional JDK HTTP OpenAI-compatible adapter for OpenAI, Azure, OpenRouter, Ollama, LM Studio, vLLM and llama.cpp, with strict-schema safeguards and localhost contract tests.
|
||||
|
||||
@@ -43,6 +43,8 @@ Always use the wrapper: Maven 3.10 cannot read the parent POM. The committed def
|
||||
- `engine/ProviderService` (active provider from `assistant_provider` rows, key references, prices), `UsageRecorder` (one `assistant_usage` row per model call), `ConnectionTest`.
|
||||
- `qbit-assistant-provider-anthropic`: `AnthropicModelProvider` and `AnthropicMapping` (anthropic-java streaming).
|
||||
- `qbit-assistant-provider-openai`: OpenAI-compatible Chat Completions, model discovery and health using JDK HTTP and Jackson.
|
||||
- `qbit-assistant-provider-bedrock`: Claude through the Anthropic Bedrock backend, AWS default credentials and SigV4; IRSA/web-identity hosts supply STS.
|
||||
- `qbit-assistant-provider-vertex`: Claude through the Anthropic Vertex backend and Gemini through OpenAI-compatible Chat Completions, with refreshed Google ADC tokens.
|
||||
|
||||
## Invariants (do not break)
|
||||
|
||||
@@ -93,4 +95,4 @@ QQQ style:
|
||||
- `develop` publishes SNAPSHOTs.
|
||||
- `release/*` publishes RCs; the orb bumps `<revision>` in the root POM.
|
||||
- `v*` tags publish releases.
|
||||
- Both modules and the parent POM publish. The flatten plugin resolves `${revision}`.
|
||||
- All five modules and the parent POM publish. The flatten plugin resolves `${revision}`.
|
||||
@@ -311,6 +311,73 @@ the live test. Run:
|
||||
./mvnw verify -pl qbit-assistant-provider-openai -am -Dtest.excludedGroups= -Dgroups=live -Djacoco.skip=true
|
||||
```
|
||||
|
||||
### AWS Bedrock (Claude)
|
||||
|
||||
Add `com.kingsrook.qbits:qbit-assistant-provider-bedrock` at the QBit version and
|
||||
register `com.kingsrook.qbits.assistant.provider.bedrock.BedrockProviderType`:
|
||||
|
||||
```java
|
||||
config.withProviderType(new BedrockProviderType());
|
||||
```
|
||||
|
||||
Create a provider row with type `bedrock`, kind `BEDROCK`, and the required AWS
|
||||
`region` (for example `us-east-1`). Leave `apiKeyEnvName` empty. The server uses
|
||||
the AWS default credential chain: environment variables, profiles and container
|
||||
credentials (ECS/EKS Pod Identity) work with the transitive dependencies.
|
||||
**For IRSA / web-identity credentials, the host application must add
|
||||
`software.amazon.awssdk:sts` to its runtime classpath**, matching the transitive
|
||||
AWS SDK version. The default chain picks it up when present; this adapter does
|
||||
not include STS or the full AWS SDK. Credentials are never stored in tables.
|
||||
|
||||
Leave the base URL empty for the SDK's
|
||||
`https://bedrock-runtime.<region>.amazonaws.com` endpoint. Overrides, including
|
||||
VPC endpoints, must be listed in `withAllowedBaseUrls`. Model ids are Bedrock
|
||||
Claude ids or inference profile ids, for example
|
||||
`us.anthropic.claude-sonnet-4-5-20250929-v1:0`. The model must be available to the
|
||||
server's AWS role in that region.
|
||||
|
||||
### Google Vertex AI (Claude, Gemini)
|
||||
|
||||
Add `com.kingsrook.qbits:qbit-assistant-provider-vertex` at the QBit version and
|
||||
register `com.kingsrook.qbits.assistant.provider.vertex.VertexProviderType`:
|
||||
|
||||
```java
|
||||
config.withProviderType(new VertexProviderType());
|
||||
```
|
||||
|
||||
Create a row with type `vertex`, kind `VERTEX`, required `region` (for example
|
||||
`us-east5`) and `cloudProject` (the GCP project id). Leave `apiKeyEnvName` empty.
|
||||
The server uses Google Application Default Credentials, such as
|
||||
`GOOGLE_APPLICATION_CREDENTIALS` or workload identity. Access tokens are refreshed
|
||||
when expired, held only in memory, and obtained afresh for each Gemini request.
|
||||
|
||||
Models beginning `claude-` (for example `claude-sonnet-4-5@20250929`) use the
|
||||
Anthropic Vertex backend's `streamRawPredict` API. All other model ids use Vertex's
|
||||
OpenAI-compatible endpoint:
|
||||
|
||||
```text
|
||||
https://{region}-aiplatform.googleapis.com/v1/projects/{project}/locations/{region}/endpoints/openapi
|
||||
```
|
||||
|
||||
Unprefixed ids, such as `gemini-2.5-flash`, are sent as `google/gemini-2.5-flash`;
|
||||
already-prefixed ids are preserved. Gemini requests omit `strict`, including when
|
||||
a tool opts in. An allow-listed base URL overrides the Claude API origin or the
|
||||
complete OpenAI-compatible base URL, according to the selected model. Use separate
|
||||
provider rows when an override needs different URL prefixes for the two APIs.
|
||||
|
||||
Both cloud types default to CLOUD and expose a short static model catalog.
|
||||
**Discover models** returns those defaults; it does not call Bedrock's
|
||||
`ListFoundationModels` (which needs the control-plane SDK) or Vertex Model Garden.
|
||||
Health probes are unsupported; **Test connection** performs the existing one-shot
|
||||
completion. Model access and availability depend on the configured region/project.
|
||||
`hipaaEligible` is an administrator-recorded fact: a BAA and the associated legal
|
||||
and operational requirements remain owner responsibilities. The flag does not
|
||||
establish eligibility or implement PHI routing.
|
||||
|
||||
The default build tests both clouds using localhost only, fake credentials, AWS
|
||||
binary event-stream framing, Claude SSE and Gemini Chat Completions, including
|
||||
engine/tool golden turns. No cloud credentials or live endpoints are required.
|
||||
|
||||
## License
|
||||
|
||||
Apache 2.0.
|
||||
+23
@@ -185,3 +185,26 @@ startup, then set the HTTPS provider row's `tlsCaBundle` to that name. Empty
|
||||
references retain JVM default trust. The bundle adds trust without replacing JVM
|
||||
roots; hostname verification remains enabled. If the P-a columns were already
|
||||
applied, add only `tls_ca_bundle VARCHAR(100)` to `assistant_provider`.
|
||||
|
||||
|
||||
## P-b: AWS Bedrock and Google Vertex
|
||||
|
||||
Before starting this version, database-backed hosts that already applied P-a add:
|
||||
|
||||
```sql
|
||||
ALTER TABLE assistant_provider ADD COLUMN cloud_project VARCHAR(100);
|
||||
```
|
||||
|
||||
Fresh installations use `docs/schema-postgres.sql`; memory-backed hosts need no
|
||||
migration. Existing providers leave the new column null. Vertex requires the GCP
|
||||
project id in `cloudProject`; Bedrock and Vertex both require `region` and reject
|
||||
`apiKeyEnvName`. They default to CLOUD; `hipaaEligible` remains a manual admin fact.
|
||||
|
||||
Add the optional `qbit-assistant-provider-bedrock` and/or
|
||||
`qbit-assistant-provider-vertex` module at the QBit version and register
|
||||
`BedrockProviderType` / `VertexProviderType`. Credentials come from the server's
|
||||
AWS default chain / Google Application Default Credentials. No credentials are
|
||||
stored in provider rows. Hosts using AWS IRSA or web identity supply
|
||||
`software.amazon.awssdk:sts` on their own runtime classpath (matching the
|
||||
transitive AWS SDK version); this QBit does not add it. See README Providers for
|
||||
configuration, endpoint overrides and static discovery limitations.
|
||||
@@ -333,6 +333,25 @@ AI exposure NONE is not enough, because the assistant acts as the user and an ad
|
||||
errors, timeouts, cancellation, discovery and a golden tool round trip. The live
|
||||
smoke is environment-driven and excluded by default.
|
||||
|
||||
### P-b as built
|
||||
|
||||
- Optional Bedrock and Vertex modules reuse the Anthropic mapping and streaming
|
||||
adapter through an injected SDK client. Bedrock retains SDK SigV4 and binary
|
||||
event-stream decoding; a small Backend delegate supplies allow-listed endpoints
|
||||
because SDK 2.71.0 has no Bedrock endpoint setter.
|
||||
- Bedrock uses the server's AWS default credential chain; IRSA/web-identity hosts
|
||||
supply STS on their runtime classpath. No STS or control-plane SDK is bundled.
|
||||
- Vertex uses scoped Google ADC. `claude-*` routes to the Anthropic backend; other
|
||||
ids route to the OpenAI-compatible endpoint with a per-request token supplier,
|
||||
`google/` for unprefixed ids, and no strict tools.
|
||||
- Both types require region, reject key references and default to CLOUD. Vertex
|
||||
additionally requires `cloudProject` (`cloud_project VARCHAR(100)`); it is part
|
||||
of resolved settings and the provider cache key. Type hooks own validation.
|
||||
- Discovery is a static catalog, health is unsupported, and Test connection uses
|
||||
the existing completion. HIPAA eligibility remains an admin fact, not a BAA.
|
||||
- Localhost contract tests cover wire paths, authentication, text/tools/usage,
|
||||
expired-token refresh, errors, allowlists and golden turns for all three routes.
|
||||
|
||||
### Phases (each a green PR)
|
||||
| | Scope | Depends on |
|
||||
|---|---|---|
|
||||
|
||||
@@ -177,6 +177,7 @@ CREATE TABLE assistant_provider
|
||||
location VARCHAR(10),
|
||||
hipaa_eligible BOOLEAN DEFAULT FALSE,
|
||||
region VARCHAR(50),
|
||||
cloud_project VARCHAR(100),
|
||||
api_version VARCHAR(30),
|
||||
connect_timeout_seconds INTEGER DEFAULT 10 CHECK (connect_timeout_seconds BETWEEN 1 AND 600),
|
||||
request_timeout_seconds INTEGER DEFAULT 120 CHECK (request_timeout_seconds BETWEEN 1 AND 600),
|
||||
|
||||
@@ -53,6 +53,8 @@
|
||||
<module>qbit-assistant</module>
|
||||
<module>qbit-assistant-provider-anthropic</module>
|
||||
<module>qbit-assistant-provider-openai</module>
|
||||
<module>qbit-assistant-provider-bedrock</module>
|
||||
<module>qbit-assistant-provider-vertex</module>
|
||||
</modules>
|
||||
|
||||
<licenses>
|
||||
@@ -82,6 +84,21 @@
|
||||
|
||||
<dependencyManagement>
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>com.kingsrook.qbits</groupId>
|
||||
<artifactId>qbit-assistant-provider-anthropic</artifactId>
|
||||
<version>${revision}</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.anthropic</groupId>
|
||||
<artifactId>anthropic-java-bedrock</artifactId>
|
||||
<version>${anthropic.version}</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.anthropic</groupId>
|
||||
<artifactId>anthropic-java-vertex</artifactId>
|
||||
<version>${anthropic.version}</version>
|
||||
</dependency>
|
||||
<!-- QQQ 4.1 is built on Jackson 2.21; the Anthropic SDK asks for 2.19. -->
|
||||
<!-- Pin the newer one so both resolve to the same line. -->
|
||||
<dependency>
|
||||
|
||||
+24
-6
@@ -98,15 +98,16 @@ public class AnthropicModelProvider implements AssistantModelProvider
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Builds the provider. apiKey is required (a literal from a secret store,
|
||||
** or "${env.NAME}").
|
||||
** Builds the provider. Unless a client is supplied, apiKey is required
|
||||
** (a literal from a secret store, or "${env.NAME}").
|
||||
*******************************************************************************/
|
||||
public static class Builder
|
||||
{
|
||||
private String apiKey;
|
||||
private String baseUrl;
|
||||
private Duration timeout = Duration.ofMinutes(10);
|
||||
private Integer maxRetries = 2;
|
||||
private String apiKey;
|
||||
private AnthropicClient client;
|
||||
private String baseUrl;
|
||||
private Duration timeout = Duration.ofMinutes(10);
|
||||
private Integer maxRetries = 2;
|
||||
|
||||
|
||||
|
||||
@@ -121,6 +122,19 @@ public class AnthropicModelProvider implements AssistantModelProvider
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Use an already configured client, including a cloud backend. The provider
|
||||
** owns and closes this client; key, URL, timeout and retry builder options
|
||||
** apply only when the builder creates its own client.
|
||||
*******************************************************************************/
|
||||
public Builder client(AnthropicClient client)
|
||||
{
|
||||
this.client = client;
|
||||
return (this);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Another API base URL (a gateway, or a test server).
|
||||
*******************************************************************************/
|
||||
@@ -159,6 +173,10 @@ public class AnthropicModelProvider implements AssistantModelProvider
|
||||
*******************************************************************************/
|
||||
public AnthropicModelProvider build()
|
||||
{
|
||||
if(client != null)
|
||||
{
|
||||
return (new AnthropicModelProvider(client, baseUrl));
|
||||
}
|
||||
String key = resolveSecret(apiKey);
|
||||
if(!StringUtils.hasContent(key))
|
||||
{
|
||||
|
||||
+24
@@ -342,4 +342,28 @@ class AnthropicModelProviderTest
|
||||
.hasMessageContaining("ANTHROPIC_API_KEY");
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** An injected cloud-capable client wins over the key and HTTP builder options.
|
||||
*******************************************************************************/
|
||||
@Test
|
||||
void injectedClientTakesPrecedence() throws Exception
|
||||
{
|
||||
server.thenStream("text-with-thinking.sse", 0);
|
||||
var client = com.anthropic.client.okhttp.AnthropicOkHttpClient.builder().apiKey(TEST_KEY).baseUrl(server.baseUrl()).maxRetries(0).build();
|
||||
var injected = AnthropicModelProvider.builder().client(client).apiKey(null).baseUrl("https://unused.invalid").build();
|
||||
try
|
||||
{
|
||||
var response = injected.stream(request(), delta -> {}).await(Duration.ofSeconds(10));
|
||||
assertThat(response.getText()).isEqualTo("There are 4 published posts.");
|
||||
assertThat(server.getRequests()).hasSize(1);
|
||||
assertThat(server.getRequests().getFirst().get("apiKey")).isEqualTo(TEST_KEY);
|
||||
}
|
||||
finally
|
||||
{
|
||||
injected.close();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<!--
|
||||
~ Copyright 2026 Kingsrook, LLC
|
||||
~
|
||||
~ Licensed under the Apache License, Version 2.0 (the "License");
|
||||
~ you may not use this file except in compliance with the License.
|
||||
~ You may obtain a copy of the License at
|
||||
~
|
||||
~ http://www.apache.org/licenses/LICENSE-2.0
|
||||
~
|
||||
~ Unless required by applicable law or agreed to in writing, software
|
||||
~ distributed under the License is distributed on an "AS IS" BASIS,
|
||||
~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
~ See the License for the specific language governing permissions and
|
||||
~ limitations under the License.
|
||||
-->
|
||||
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
|
||||
<parent>
|
||||
<groupId>com.kingsrook.qbits</groupId>
|
||||
<artifactId>qbit-assistant-parent</artifactId>
|
||||
<version>${revision}</version>
|
||||
</parent>
|
||||
|
||||
<artifactId>qbit-assistant-provider-bedrock</artifactId>
|
||||
<packaging>jar</packaging>
|
||||
|
||||
<name>QBit Assistant - Bedrock provider</name>
|
||||
<description>Bedrock provider using the Anthropic cloud backend</description>
|
||||
<url>https://github.com/QRun-IO/qbit-assistant</url>
|
||||
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>com.anthropic</groupId>
|
||||
<artifactId>anthropic-java-bedrock</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.kingsrook.qbits</groupId>
|
||||
<artifactId>qbit-assistant-provider-anthropic</artifactId>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>com.kingsrook.qbits</groupId>
|
||||
<artifactId>qbit-assistant</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.kingsrook.qqq</groupId>
|
||||
<artifactId>qqq-backend-core</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.fasterxml.jackson.core</groupId>
|
||||
<artifactId>jackson-databind</artifactId>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>com.kingsrook.qbits</groupId>
|
||||
<artifactId>qbit-assistant</artifactId>
|
||||
<type>test-jar</type>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.kingsrook.qqq</groupId>
|
||||
<artifactId>qqq-middleware-javalin</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.junit.jupiter</groupId>
|
||||
<artifactId>junit-jupiter</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.assertj</groupId>
|
||||
<artifactId>assertj-core</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
<build>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-surefire-plugin</artifactId>
|
||||
<configuration>
|
||||
<!-- live evals call the real API: opt-in only -->
|
||||
<excludedGroups>${test.excludedGroups}</excludedGroups>
|
||||
<redirectTestOutputToFile>true</redirectTestOutputToFile>
|
||||
</configuration>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.jacoco</groupId>
|
||||
<artifactId>jacoco-maven-plugin</artifactId>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>prepare-agent</id>
|
||||
<goals>
|
||||
<goal>prepare-agent</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
<execution>
|
||||
<id>report</id>
|
||||
<phase>verify</phase>
|
||||
<goals>
|
||||
<goal>report</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
<execution>
|
||||
<id>check</id>
|
||||
<goals>
|
||||
<goal>check</goal>
|
||||
</goals>
|
||||
<configuration>
|
||||
<rules>
|
||||
<rule>
|
||||
<element>BUNDLE</element>
|
||||
<limits>
|
||||
<limit>
|
||||
<counter>INSTRUCTION</counter>
|
||||
<value>COVEREDRATIO</value>
|
||||
<minimum>0.70</minimum>
|
||||
</limit>
|
||||
<limit>
|
||||
<counter>CLASS</counter>
|
||||
<value>COVEREDRATIO</value>
|
||||
<minimum>0.90</minimum>
|
||||
</limit>
|
||||
</limits>
|
||||
</rule>
|
||||
</rules>
|
||||
</configuration>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</build>
|
||||
</project>
|
||||
+73
@@ -0,0 +1,73 @@
|
||||
/*
|
||||
* Copyright 2026 Kingsrook, LLC
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package com.kingsrook.qbits.assistant.provider.bedrock;
|
||||
|
||||
|
||||
import com.anthropic.backends.Backend;
|
||||
import com.anthropic.core.http.HttpRequest;
|
||||
import com.anthropic.core.http.HttpResponse;
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** SDK 2.71 has no Bedrock endpoint setter. Override only the endpoint, keeping
|
||||
** the SDK's request conversion, SigV4 signing of that endpoint, and decoding.
|
||||
** The caller must validate the URL against the host's allowlist first.
|
||||
*******************************************************************************/
|
||||
record BedrockEndpointBackend(Backend delegate, String baseUrl) implements Backend
|
||||
{
|
||||
/*******************************************************************************
|
||||
** Preserve the Bedrock model-path and body conversion.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public HttpRequest prepareRequest(HttpRequest request)
|
||||
{
|
||||
return (delegate.prepareRequest(request));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Sign the actual request URL, including the overridden host.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public HttpRequest authorizeRequest(HttpRequest request)
|
||||
{
|
||||
return (delegate.authorizeRequest(request));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Decode the AWS event stream through the SDK.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public HttpResponse prepareResponse(HttpResponse response)
|
||||
{
|
||||
return (delegate.prepareResponse(response));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Release the SDK's stream conversion resources.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public void close()
|
||||
{
|
||||
delegate.close();
|
||||
}
|
||||
}
|
||||
+185
@@ -0,0 +1,185 @@
|
||||
/*
|
||||
* Copyright 2026 Kingsrook, LLC
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package com.kingsrook.qbits.assistant.provider.bedrock;
|
||||
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import com.anthropic.backends.Backend;
|
||||
import com.anthropic.bedrock.backends.BedrockBackend;
|
||||
import com.anthropic.client.okhttp.AnthropicOkHttpClient;
|
||||
import com.anthropic.core.Timeout;
|
||||
import com.kingsrook.qbits.assistant.model.AssistantProvider;
|
||||
import com.kingsrook.qbits.assistant.provider.AssistantModelProvider;
|
||||
import com.kingsrook.qbits.assistant.provider.AssistantProviderType;
|
||||
import com.kingsrook.qbits.assistant.provider.DiscoveredModel;
|
||||
import com.kingsrook.qbits.assistant.provider.ProviderSettings;
|
||||
import com.kingsrook.qbits.assistant.provider.anthropic.AnthropicModelProvider;
|
||||
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
|
||||
import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** AWS Bedrock (Claude), using server credentials and region-derived endpoints.
|
||||
** Environment, profile and container (ECS/EKS Pod Identity) credentials use
|
||||
** the transitive SDK dependencies. Hosts using IRSA / web identity must add
|
||||
** software.amazon.awssdk:sts to their runtime classpath; the default credential
|
||||
** chain discovers it when present. This adapter does not depend on STS.
|
||||
*******************************************************************************/
|
||||
public class BedrockProviderType implements AssistantProviderType
|
||||
{
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Provider registry contract.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public String getType()
|
||||
{
|
||||
return ("bedrock");
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Provider registry contract.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public String getLabel()
|
||||
{
|
||||
return ("AWS Bedrock (Claude)");
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Provider registry contract.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public List<String> getModels()
|
||||
{
|
||||
return (List.of("us.anthropic.claude-sonnet-4-5-20250929-v1:0", "us.anthropic.claude-haiku-4-5-20251001-v1:0"));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Provider registry contract.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public String getDefaultApiKeyEnvName()
|
||||
{
|
||||
return (null);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Provider registry contract.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public String getDefaultBaseUrl()
|
||||
{
|
||||
return (null);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Provider registry contract.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public boolean requiresBaseUrl(String kind)
|
||||
{
|
||||
return (false);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Validate cloud settings without loading credentials.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public List<String> validate(AssistantProvider row)
|
||||
{
|
||||
List<String> errors = new ArrayList<>();
|
||||
if(row.getRegion() == null || row.getRegion().isBlank())
|
||||
{
|
||||
errors.add("Region is required for BEDROCK.");
|
||||
}
|
||||
if(row.getApiKeyEnvName() != null && !row.getApiKeyEnvName().isBlank())
|
||||
{
|
||||
errors.add("Leave API key variable empty: AWS Bedrock (Claude) uses the AWS default credential chain (environment, profile or server role).");
|
||||
}
|
||||
return (errors);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Cloud providers require resolved region/project settings.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public AssistantModelProvider create(String apiKey, String baseUrl)
|
||||
{
|
||||
throw (new IllegalArgumentException("Use ProviderSettings with the required cloud region."));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Static defaults: ListFoundationModels requires the Bedrock control-plane SDK,
|
||||
** which this adapter deliberately does not add.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public List<DiscoveredModel> discoverModels(ProviderSettings settings)
|
||||
{
|
||||
return (getModels().stream().map(id -> new DiscoveredModel(id, id, null, null)).toList());
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Create the provider with server credentials.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public AssistantModelProvider create(ProviderSettings settings)
|
||||
{
|
||||
return (create(settings, DefaultCredentialsProvider.builder().build()));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Credential injection for contract tests; no credential values are stored.
|
||||
*******************************************************************************/
|
||||
AnthropicModelProvider create(ProviderSettings settings, AwsCredentialsProvider credentials)
|
||||
{
|
||||
Backend backend = BedrockBackend.builder().region(Region.of(settings.getRegion())).awsCredentialsProvider(credentials).build();
|
||||
if(settings.getBaseUrl() != null)
|
||||
{
|
||||
backend = new BedrockEndpointBackend(backend, settings.getBaseUrl());
|
||||
}
|
||||
var client = AnthropicOkHttpClient.builder().backend(backend)
|
||||
.timeout(Timeout.builder().connect(Duration.ofSeconds(settings.getConnectTimeoutSeconds()))
|
||||
.request(Duration.ofSeconds(settings.getRequestTimeoutSeconds())).build());
|
||||
return (AnthropicModelProvider.builder().client(client.build()).build());
|
||||
}
|
||||
}
|
||||
+213
@@ -0,0 +1,213 @@
|
||||
/*
|
||||
* Copyright 2026 Kingsrook, LLC
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package com.kingsrook.qbits.assistant.provider.bedrock;
|
||||
|
||||
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.time.Clock;
|
||||
import java.time.Duration;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Set;
|
||||
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
|
||||
import com.kingsrook.qbits.assistant.engine.ProviderService;
|
||||
import com.kingsrook.qbits.assistant.model.AssistantProvider;
|
||||
import com.kingsrook.qbits.assistant.provider.ModelEventListener;
|
||||
import com.kingsrook.qbits.assistant.provider.ModelMessage;
|
||||
import com.kingsrook.qbits.assistant.provider.ModelProviderException;
|
||||
import com.kingsrook.qbits.assistant.provider.ModelRequest;
|
||||
import com.kingsrook.qbits.assistant.provider.ModelToolDefinition;
|
||||
import com.kingsrook.qbits.assistant.provider.ProviderSettings;
|
||||
import com.kingsrook.qbits.assistant.provider.anthropic.AnthropicModelProvider;
|
||||
import com.kingsrook.qbits.assistant.testing.CloudProviderTestServer;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.params.ParameterizedTest;
|
||||
import org.junit.jupiter.params.provider.CsvSource;
|
||||
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
|
||||
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
|
||||
import static com.kingsrook.qbits.assistant.testing.CloudProviderTestServer.JSON;
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Real Anthropic Bedrock SDK against signed loopback requests and binary streams.
|
||||
*******************************************************************************/
|
||||
class BedrockContractTest
|
||||
{
|
||||
static final String MODEL = "us.anthropic.claude-sonnet-4-5-20250929-v1:0";
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Explicit fake credentials avoid all default-chain network lookup in tests.
|
||||
*******************************************************************************/
|
||||
static AnthropicModelProvider provider(CloudProviderTestServer server)
|
||||
{
|
||||
return (new BedrockProviderType().create(new ProviderSettings().withRegion("us-east-1").withBaseUrl(server.url())
|
||||
.withConnectTimeoutSeconds(2).withRequestTimeoutSeconds(5),
|
||||
StaticCredentialsProvider.create(AwsBasicCredentials.create("AKIDEXAMPLE", "synthetic-secret"))));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Engine-shaped request with an opted-in strict tool.
|
||||
*******************************************************************************/
|
||||
static ModelRequest request()
|
||||
{
|
||||
return (new ModelRequest().withModel(MODEL).withMaxTokens(1024).withMessages(List.of(ModelMessage.user("Count posts")))
|
||||
.withTools(List.of(new ModelToolDefinition().withName("count_records").withDescription("Count records").withStrict(true)
|
||||
.withInputSchema(Map.of("type", "object", "properties", Map.of("table", Map.of("type", "string")), "required", List.of("table"), "additionalProperties", false)))));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Binary chunk framing, SigV4, request rewriting, text, tools and usage.
|
||||
*******************************************************************************/
|
||||
@Test
|
||||
void streamingContract() throws Exception
|
||||
{
|
||||
try(CloudProviderTestServer server = new CloudProviderTestServer())
|
||||
{
|
||||
server.enqueue(200, "application/vnd.amazon.eventstream", BedrockEventStream.encode(CloudProviderTestServer.anthropicEvents(MODEL, "Checking.", Map.of("table", "blog_post"))));
|
||||
AnthropicModelProvider provider = provider(server);
|
||||
try
|
||||
{
|
||||
StringBuilder text = new StringBuilder();
|
||||
List<String> tools = new ArrayList<>();
|
||||
var response = provider.stream(request(), new ModelEventListener()
|
||||
{
|
||||
/*******************************************************************************
|
||||
** Capture only visible text.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public void onTextDelta(String delta)
|
||||
{
|
||||
text.append(delta);
|
||||
}
|
||||
|
||||
/*******************************************************************************
|
||||
** Confirm tool-start callbacks pass through the reused adapter.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public void onToolUseStart(String id, String name)
|
||||
{
|
||||
tools.add(id + ":" + name);
|
||||
}
|
||||
}).await(Duration.ofSeconds(10));
|
||||
assertThat(text.toString()).isEqualTo("Checking.");
|
||||
assertThat(response.getText()).isEqualTo("Checking.");
|
||||
assertThat(tools).containsExactly("count:count_records");
|
||||
assertThat(response.getToolUses()).singleElement().satisfies(tool ->
|
||||
{
|
||||
assertThat(tool.getToolUseId()).isEqualTo("count");
|
||||
assertThat(tool.getInput()).containsEntry("table", "blog_post");
|
||||
});
|
||||
assertThat(response.getStopReason()).isEqualTo("tool_use");
|
||||
assertThat(response.getUsage().getInputTokens()).isEqualTo(12);
|
||||
assertThat(response.getUsage().getOutputTokens()).isEqualTo(7);
|
||||
var sent = server.requests.getFirst();
|
||||
assertThat(sent.path()).isEqualTo("/model/" + MODEL + "/invoke-with-response-stream");
|
||||
assertThat(sent.authorization()).startsWith("AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/").contains("/us-east-1/bedrock/aws4_request", "Signature=");
|
||||
var body = JSON.readTree(sent.body());
|
||||
assertThat(body.path("anthropic_version").asText()).isEqualTo("bedrock-2023-05-31");
|
||||
assertThat(body.has("model")).isFalse();
|
||||
assertThat(body.has("stream")).isFalse();
|
||||
assertThat(body.path("max_tokens").asInt()).isEqualTo(1024);
|
||||
assertThat(body.path("messages").get(0).path("role").asText()).isEqualTo("user");
|
||||
assertThat(body.path("tools").get(0).path("strict").asBoolean()).isTrue();
|
||||
}
|
||||
finally
|
||||
{
|
||||
provider.close();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** The cloud backend keeps the Anthropic adapter's HTTP error mapping.
|
||||
*******************************************************************************/
|
||||
@ParameterizedTest
|
||||
@CsvSource({ "403,AUTHENTICATION", "429,RATE_LIMITED", "500,SERVER" })
|
||||
void errors(Integer status, ModelProviderException.Kind kind) throws Exception
|
||||
{
|
||||
try(CloudProviderTestServer server = new CloudProviderTestServer())
|
||||
{
|
||||
for(int attempt = 0; attempt < 3; attempt++)
|
||||
{
|
||||
server.enqueue(status, "application/json", "{\"message\":\"Synthetic failure\"}".getBytes(StandardCharsets.UTF_8));
|
||||
}
|
||||
AnthropicModelProvider provider = provider(server);
|
||||
try
|
||||
{
|
||||
assertThatThrownBy(() -> provider.stream(request(), delta -> {}).await(Duration.ofSeconds(20)))
|
||||
.isInstanceOfSatisfying(ModelProviderException.class, error ->
|
||||
{
|
||||
assertThat(error.getKind()).isEqualTo(kind);
|
||||
assertThat(error.getHttpStatus()).isEqualTo(status);
|
||||
});
|
||||
}
|
||||
finally
|
||||
{
|
||||
provider.close();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Registry defaults and validation run without resolving AWS credentials.
|
||||
*******************************************************************************/
|
||||
@Test
|
||||
void settingsValidationAndDefaults() throws Exception
|
||||
{
|
||||
var type = new BedrockProviderType();
|
||||
var config = new AssistantQBitConfig().withProviderType(type).withAllowedBaseUrls(Set.of("http://127.0.0.1:1234"));
|
||||
var service = new ProviderService(config, null, Clock.systemUTC());
|
||||
var row = new AssistantProvider().withProviderType("bedrock");
|
||||
assertThat(service.validate(row)).anyMatch(error -> error.contains("Region is required"));
|
||||
row.withRegion(" ");
|
||||
assertThat(service.validate(row)).anyMatch(error -> error.contains("Region is required"));
|
||||
row.withRegion("us-east-1");
|
||||
assertThat(service.validate(row)).isEmpty();
|
||||
row.withApiKeyEnvName("AWS_API_KEY");
|
||||
assertThat(service.validate(row)).anyMatch(error -> error.contains("default credential chain"));
|
||||
row.withApiKeyEnvName(" ").withBaseUrl("https://unlisted.example");
|
||||
assertThat(service.validate(row)).anyMatch(error -> error.contains("not allowed"));
|
||||
row.withBaseUrl("http://127.0.0.1:1234");
|
||||
assertThat(service.validate(row)).isEmpty();
|
||||
assertThat(type.getType()).isEqualTo("bedrock");
|
||||
assertThat(type.getKinds()).containsExactly("BEDROCK");
|
||||
assertThat(type.getLabel()).isEqualTo("AWS Bedrock (Claude)");
|
||||
assertThat(type.getDefaultLocation("BEDROCK")).isEqualTo("CLOUD");
|
||||
assertThat(type.getDefaultApiKeyEnvName()).isNull();
|
||||
assertThat(type.getDefaultBaseUrl()).isNull();
|
||||
assertThat(type.discoverModels(new ProviderSettings())).extracting(model -> model.modelId()).containsExactlyElementsOf(type.getModels());
|
||||
assertThat(type.checkHealth(new ProviderSettings()).healthy()).isNull();
|
||||
assertThatThrownBy(() -> type.create("key", null)).isInstanceOf(IllegalArgumentException.class);
|
||||
// Client creation is lazy: no credentials are resolved and no request sent.
|
||||
var defaultChain = (AnthropicModelProvider) type.create(new ProviderSettings().withRegion("us-east-1"));
|
||||
defaultChain.close();
|
||||
}
|
||||
}
|
||||
|
||||
+90
@@ -0,0 +1,90 @@
|
||||
/*
|
||||
* Copyright 2026 Kingsrook, LLC
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package com.kingsrook.qbits.assistant.provider.bedrock;
|
||||
|
||||
|
||||
import java.io.ByteArrayOutputStream;
|
||||
import java.io.DataOutputStream;
|
||||
import java.io.IOException;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.Base64;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.zip.CRC32;
|
||||
import com.kingsrook.qbits.assistant.testing.CloudProviderTestServer;
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Test-side AWS event-stream encoder: prelude, typed headers and both CRC32s.
|
||||
*******************************************************************************/
|
||||
final class BedrockEventStream
|
||||
{
|
||||
/*******************************************************************************
|
||||
** Each chunk carries base64 Anthropic event JSON, not an SSE envelope.
|
||||
*******************************************************************************/
|
||||
static byte[] encode(List<String> events) throws IOException
|
||||
{
|
||||
ByteArrayOutputStream stream = new ByteArrayOutputStream();
|
||||
for(String event : events)
|
||||
{
|
||||
byte[] payload = CloudProviderTestServer.JSON.writeValueAsBytes(Map.of("bytes", Base64.getEncoder().encodeToString(event.getBytes(StandardCharsets.UTF_8))));
|
||||
ByteArrayOutputStream headers = new ByteArrayOutputStream();
|
||||
DataOutputStream headerData = new DataOutputStream(headers);
|
||||
header(headerData, ":message-type", "event");
|
||||
header(headerData, ":event-type", "chunk");
|
||||
header(headerData, ":content-type", "application/json");
|
||||
ByteArrayOutputStream frame = new ByteArrayOutputStream();
|
||||
DataOutputStream data = new DataOutputStream(frame);
|
||||
data.writeInt(16 + headers.size() + payload.length);
|
||||
data.writeInt(headers.size());
|
||||
data.writeInt(crc(frame.toByteArray()));
|
||||
data.write(headers.toByteArray());
|
||||
data.write(payload);
|
||||
data.writeInt(crc(frame.toByteArray()));
|
||||
stream.write(frame.toByteArray());
|
||||
}
|
||||
return (stream.toByteArray());
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** AWS header type 7 is a UTF-8 string with a two-byte length.
|
||||
*******************************************************************************/
|
||||
private static void header(DataOutputStream data, String name, String value) throws IOException
|
||||
{
|
||||
byte[] nameBytes = name.getBytes(StandardCharsets.UTF_8);
|
||||
byte[] valueBytes = value.getBytes(StandardCharsets.UTF_8);
|
||||
data.writeByte(nameBytes.length);
|
||||
data.write(nameBytes);
|
||||
data.writeByte(7);
|
||||
data.writeShort(valueBytes.length);
|
||||
data.write(valueBytes);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Network-order CRC32 for the prelude and complete message respectively.
|
||||
*******************************************************************************/
|
||||
private static int crc(byte[] bytes)
|
||||
{
|
||||
CRC32 crc = new CRC32();
|
||||
crc.update(bytes);
|
||||
return ((int) crc.getValue());
|
||||
}
|
||||
}
|
||||
+87
@@ -0,0 +1,87 @@
|
||||
/*
|
||||
* Copyright 2026 Kingsrook, LLC
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package com.kingsrook.qbits.assistant.provider.bedrock;
|
||||
|
||||
|
||||
import com.kingsrook.qbits.assistant.AssistantModelRouting;
|
||||
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
|
||||
import com.kingsrook.qbits.assistant.testing.AssistantTestFixture;
|
||||
import com.kingsrook.qbits.assistant.testing.CloudProviderTestServer;
|
||||
import com.kingsrook.qbits.assistant.testing.GoldenTasks;
|
||||
import com.kingsrook.qbits.assistant.testing.TurnRunner;
|
||||
import com.kingsrook.qqq.backend.core.context.QContext;
|
||||
import org.junit.jupiter.api.Tag;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import static com.kingsrook.qbits.assistant.testing.CloudProviderTestServer.JSON;
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Golden task through the cloud adapter, turn engine and permissioned UI tools.
|
||||
*******************************************************************************/
|
||||
@Tag("golden")
|
||||
class BedrockGoldenTurnTest
|
||||
{
|
||||
/*******************************************************************************
|
||||
** A real count_records round trip returns the tool result to the model.
|
||||
*******************************************************************************/
|
||||
@Test
|
||||
void goldenToolRoundTrip() throws Exception
|
||||
{
|
||||
try(CloudProviderTestServer server = new CloudProviderTestServer())
|
||||
{
|
||||
String answer = "There are " + AssistantTestFixture.PUBLISHED_POSTS + " published posts. /app/blog_post";
|
||||
String model = BedrockContractTest.MODEL;
|
||||
server.enqueue(200, "application/vnd.amazon.eventstream", BedrockEventStream.encode(CloudProviderTestServer.anthropicEvents(model, "", CloudProviderTestServer.countInput())));
|
||||
server.enqueue(200, "application/vnd.amazon.eventstream", BedrockEventStream.encode(CloudProviderTestServer.anthropicEvents(model, answer, null)));
|
||||
var provider = BedrockContractTest.provider(server);
|
||||
var config = new AssistantQBitConfig().withProviderName(provider.getName()).withModelProvider(provider)
|
||||
.withModels(new AssistantModelRouting().withMainModel(model));
|
||||
try
|
||||
{
|
||||
var instance = AssistantTestFixture.build(config);
|
||||
var task = GoldenTasks.all().stream().filter(t -> t.id().equals("count-published-posts")).findFirst().orElseThrow();
|
||||
var outcome = new TurnRunner(instance, config).run(task.userId(), task.prompt());
|
||||
assertThat(outcome.errorCode()).isNull();
|
||||
task.check().accept(outcome);
|
||||
assertThat(outcome.text()).endsWith(answer);
|
||||
assertThat(outcome.end()).isNotEmpty();
|
||||
assertThat(server.requests).hasSize(2);
|
||||
var messages = JSON.readTree(server.requests.get(1).body()).path("messages");
|
||||
var assistant = messages.get(messages.size() - 2);
|
||||
var result = messages.get(messages.size() - 1);
|
||||
assertThat(assistant.path("role").asText()).isEqualTo("assistant");
|
||||
assertThat(assistant.path("content").toString()).contains("count_records", "count");
|
||||
assertThat(result.path("role").asText()).isEqualTo("user");
|
||||
var toolResult = result.path("content").get(0);
|
||||
assertThat(toolResult.path("type").asText()).isEqualTo("tool_result");
|
||||
assertThat(toolResult.path("tool_use_id").asText()).isEqualTo("count");
|
||||
assertThat(toolResult.path("is_error").asBoolean()).isFalse();
|
||||
assertThat(toolResult.path("content").toString()).contains("count");
|
||||
}
|
||||
finally
|
||||
{
|
||||
if(config.getRuntime() != null)
|
||||
{
|
||||
config.getRuntime().close();
|
||||
}
|
||||
provider.close();
|
||||
QContext.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+21
-1
@@ -62,6 +62,7 @@ final class OpenAiHttp
|
||||
settings = new ProviderSettings().withKind(source.getKind() == null ? "OPENAI_COMPATIBLE" : source.getKind())
|
||||
.withTlsCaBundle(source.getTlsCaBundle()).withSslContext(source.getSslContext())
|
||||
.withApiKey(source.getApiKey()).withBaseUrl(source.getBaseUrl()).withApiVersion(source.getApiVersion()).withRegion(source.getRegion())
|
||||
.withBearerTokenSupplier(source.getBearerTokenSupplier()).withCloudProject(source.getCloudProject())
|
||||
.withConnectTimeoutSeconds(source.getConnectTimeoutSeconds()).withRequestTimeoutSeconds(source.getRequestTimeoutSeconds());
|
||||
String url = settings.getBaseUrl() == null ? new OpenAiCompatibleProviderType().getDefaultBaseUrl(settings.getKind()) : settings.getBaseUrl();
|
||||
if(url == null)
|
||||
@@ -93,7 +94,16 @@ final class OpenAiHttp
|
||||
path += "?api-version=" + encode(settings.getApiVersion() == null || settings.getApiVersion().isBlank() ? "2024-10-21" : settings.getApiVersion());
|
||||
}
|
||||
HttpRequest.Builder builder = HttpRequest.newBuilder(URI.create(base + path)).timeout(Duration.ofSeconds(seconds)).header("Content-Type", "application/json");
|
||||
if(settings.getApiKey() != null && !settings.getApiKey().isBlank())
|
||||
if(settings.getBearerTokenSupplier() != null)
|
||||
{
|
||||
String token = settings.getBearerTokenSupplier().get();
|
||||
if(token == null || token.isBlank())
|
||||
{
|
||||
throw (new ModelProviderException(ModelProviderException.Kind.AUTHENTICATION, null, "The provider access token is unavailable."));
|
||||
}
|
||||
builder.header("Authorization", "Bearer " + token);
|
||||
}
|
||||
else if(settings.getApiKey() != null && !settings.getApiKey().isBlank())
|
||||
{
|
||||
if("AZURE_OPENAI".equals(settings.getKind()))
|
||||
{
|
||||
@@ -110,6 +120,10 @@ final class OpenAiHttp
|
||||
{
|
||||
throw (new ModelProviderException(ModelProviderException.Kind.INVALID_REQUEST, null, "The provider connection settings are invalid."));
|
||||
}
|
||||
catch(RuntimeException e)
|
||||
{
|
||||
throw (new ModelProviderException(ModelProviderException.Kind.AUTHENTICATION, null, "Could not obtain the provider access token."));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -203,6 +217,12 @@ final class OpenAiHttp
|
||||
*******************************************************************************/
|
||||
private String safe(String value)
|
||||
{
|
||||
if(settings.getBearerTokenSupplier() != null)
|
||||
{
|
||||
// Dynamic credentials can be echoed by a remote error. Do not retain
|
||||
// tokens or include remote diagnostics for bearer-supplier requests.
|
||||
return (null);
|
||||
}
|
||||
String key = settings.getApiKey();
|
||||
return (ModelProviderException.sanitizeProviderMessage(value == null || key == null || key.isEmpty() ? value : value.replace(key, "[redacted]"), 300));
|
||||
}
|
||||
|
||||
+2
-1
@@ -54,7 +54,8 @@ final class OpenAiMapping
|
||||
ObjectNode body = OpenAiHttp.JSON.createObjectNode();
|
||||
boolean streaming = Boolean.TRUE.equals(capabilities.getStreaming());
|
||||
boolean nativeOpenAi = List.of("OPENAI", "AZURE_OPENAI").contains(kind);
|
||||
body.put("model", request.getModel());
|
||||
String model = request.getModel();
|
||||
body.put("model", "VERTEX".equals(kind) && model != null && !model.contains("/") ? "google/" + model : model);
|
||||
body.put("stream", streaming);
|
||||
if(streaming)
|
||||
{
|
||||
|
||||
+66
@@ -458,4 +458,70 @@ class OpenAiContractTest
|
||||
assertThat(response.getToolUses().getFirst().getInput()).containsExactlyEntriesOf(Map.of("_invalid_arguments", arguments.substring(0, Math.min(200, arguments.length()))));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Supplier tokens take precedence over static keys and refresh per request.
|
||||
*******************************************************************************/
|
||||
@Test
|
||||
void bearerSupplierPrecedenceAndRedaction() throws Exception
|
||||
{
|
||||
try(MockOpenAiServer server = new MockOpenAiServer())
|
||||
{
|
||||
var token = new java.util.concurrent.atomic.AtomicReference<>("first-private-token");
|
||||
ProviderSettings settings = settings(server, "VERTEX").withBearerTokenSupplier(token::get);
|
||||
var provider = new OpenAiCompatibleModelProvider(settings);
|
||||
server.enqueue(200, TEXT);
|
||||
provider.stream(request(), delta -> {}).await(Duration.ofSeconds(5));
|
||||
assertThat(server.requests.getFirst().authorization()).isEqualTo("Bearer first-private-token");
|
||||
token.set("next-private-token");
|
||||
server.enqueue(403, "{\"error\":{\"type\":\"next-private-token\",\"message\":\"next-private-token\"}}");
|
||||
assertThatThrownBy(() -> provider.stream(request(), delta -> {}).await(Duration.ofSeconds(5)))
|
||||
.isInstanceOfSatisfying(ModelProviderException.class, error ->
|
||||
{
|
||||
assertThat(error.getKind()).isEqualTo(ModelProviderException.Kind.AUTHENTICATION);
|
||||
assertThat(error.getProviderMessage()).isNull();
|
||||
assertThat(error.getProviderErrorType()).isNull();
|
||||
assertThat(error.toString()).doesNotContain("private-token");
|
||||
});
|
||||
assertThat(server.requests.get(1).authorization()).isEqualTo("Bearer next-private-token");
|
||||
assertThat(settings.toString()).doesNotContain("private-token");
|
||||
assertThat(OpenAiHttp.JSON.writeValueAsString(settings)).doesNotContain("bearerTokenSupplier", "private-token");
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Missing or failed dynamic credentials never fall back to a static secret.
|
||||
*******************************************************************************/
|
||||
@Test
|
||||
void unavailableBearerSupplier() throws Exception
|
||||
{
|
||||
try(MockOpenAiServer server = new MockOpenAiServer())
|
||||
{
|
||||
ProviderSettings settings = settings(server, "VERTEX");
|
||||
for(String value : java.util.Arrays.asList(null, "", " "))
|
||||
{
|
||||
settings.setBearerTokenSupplier(() -> value);
|
||||
var provider = new OpenAiCompatibleModelProvider(settings);
|
||||
assertThatThrownBy(() -> provider.stream(request(), delta -> {}).await(Duration.ofSeconds(5)))
|
||||
.isInstanceOfSatisfying(ModelProviderException.class, error -> assertThat(error.getKind()).isEqualTo(ModelProviderException.Kind.AUTHENTICATION));
|
||||
}
|
||||
settings.withBearerTokenSupplier(() ->
|
||||
{
|
||||
throw (new IllegalStateException("private-token"));
|
||||
});
|
||||
var provider = new OpenAiCompatibleModelProvider(settings);
|
||||
assertThatThrownBy(() -> provider.stream(request(), delta -> {}).await(Duration.ofSeconds(5)))
|
||||
.isInstanceOfSatisfying(ModelProviderException.class, error ->
|
||||
{
|
||||
assertThat(error.getKind()).isEqualTo(ModelProviderException.Kind.AUTHENTICATION);
|
||||
assertThat(error.toString()).doesNotContain("private-token");
|
||||
assertThat(error.getCause()).isNull();
|
||||
});
|
||||
assertThat(server.requests).isEmpty();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<!--
|
||||
~ Copyright 2026 Kingsrook, LLC
|
||||
~
|
||||
~ Licensed under the Apache License, Version 2.0 (the "License");
|
||||
~ you may not use this file except in compliance with the License.
|
||||
~ You may obtain a copy of the License at
|
||||
~
|
||||
~ http://www.apache.org/licenses/LICENSE-2.0
|
||||
~
|
||||
~ Unless required by applicable law or agreed to in writing, software
|
||||
~ distributed under the License is distributed on an "AS IS" BASIS,
|
||||
~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
~ See the License for the specific language governing permissions and
|
||||
~ limitations under the License.
|
||||
-->
|
||||
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
|
||||
<parent>
|
||||
<groupId>com.kingsrook.qbits</groupId>
|
||||
<artifactId>qbit-assistant-parent</artifactId>
|
||||
<version>${revision}</version>
|
||||
</parent>
|
||||
|
||||
<artifactId>qbit-assistant-provider-vertex</artifactId>
|
||||
<packaging>jar</packaging>
|
||||
|
||||
<name>QBit Assistant - Vertex provider</name>
|
||||
<description>Vertex provider using the Anthropic cloud backend</description>
|
||||
<url>https://github.com/QRun-IO/qbit-assistant</url>
|
||||
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>com.anthropic</groupId>
|
||||
<artifactId>anthropic-java-vertex</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.kingsrook.qbits</groupId>
|
||||
<artifactId>qbit-assistant-provider-anthropic</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.kingsrook.qbits</groupId>
|
||||
<artifactId>qbit-assistant-provider-openai</artifactId>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>com.kingsrook.qbits</groupId>
|
||||
<artifactId>qbit-assistant</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.kingsrook.qqq</groupId>
|
||||
<artifactId>qqq-backend-core</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.fasterxml.jackson.core</groupId>
|
||||
<artifactId>jackson-databind</artifactId>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>com.kingsrook.qbits</groupId>
|
||||
<artifactId>qbit-assistant</artifactId>
|
||||
<type>test-jar</type>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.kingsrook.qqq</groupId>
|
||||
<artifactId>qqq-middleware-javalin</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.junit.jupiter</groupId>
|
||||
<artifactId>junit-jupiter</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.assertj</groupId>
|
||||
<artifactId>assertj-core</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
<build>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-surefire-plugin</artifactId>
|
||||
<configuration>
|
||||
<!-- live evals call the real API: opt-in only -->
|
||||
<excludedGroups>${test.excludedGroups}</excludedGroups>
|
||||
<redirectTestOutputToFile>true</redirectTestOutputToFile>
|
||||
</configuration>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.jacoco</groupId>
|
||||
<artifactId>jacoco-maven-plugin</artifactId>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>prepare-agent</id>
|
||||
<goals>
|
||||
<goal>prepare-agent</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
<execution>
|
||||
<id>report</id>
|
||||
<phase>verify</phase>
|
||||
<goals>
|
||||
<goal>report</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
<execution>
|
||||
<id>check</id>
|
||||
<goals>
|
||||
<goal>check</goal>
|
||||
</goals>
|
||||
<configuration>
|
||||
<rules>
|
||||
<rule>
|
||||
<element>BUNDLE</element>
|
||||
<limits>
|
||||
<limit>
|
||||
<counter>INSTRUCTION</counter>
|
||||
<value>COVEREDRATIO</value>
|
||||
<minimum>0.70</minimum>
|
||||
</limit>
|
||||
<limit>
|
||||
<counter>CLASS</counter>
|
||||
<value>COVEREDRATIO</value>
|
||||
<minimum>0.90</minimum>
|
||||
</limit>
|
||||
</limits>
|
||||
</rule>
|
||||
</rules>
|
||||
</configuration>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</build>
|
||||
</project>
|
||||
+166
@@ -0,0 +1,166 @@
|
||||
/*
|
||||
* Copyright 2026 Kingsrook, LLC
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package com.kingsrook.qbits.assistant.provider.vertex;
|
||||
|
||||
|
||||
import java.io.IOException;
|
||||
import java.time.Duration;
|
||||
import com.anthropic.client.okhttp.AnthropicOkHttpClient;
|
||||
import com.anthropic.core.Timeout;
|
||||
import com.anthropic.vertex.backends.VertexBackend;
|
||||
import com.google.auth.oauth2.GoogleCredentials;
|
||||
import com.kingsrook.qbits.assistant.provider.AssistantModelProvider;
|
||||
import com.kingsrook.qbits.assistant.provider.ModelCapabilities;
|
||||
import com.kingsrook.qbits.assistant.provider.ModelEventListener;
|
||||
import com.kingsrook.qbits.assistant.provider.ModelRequest;
|
||||
import com.kingsrook.qbits.assistant.provider.ProviderSettings;
|
||||
import com.kingsrook.qbits.assistant.provider.StreamHandle;
|
||||
import com.kingsrook.qbits.assistant.provider.anthropic.AnthropicModelProvider;
|
||||
import com.kingsrook.qbits.assistant.provider.openai.OpenAiCompatibleModelProvider;
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Route Claude to the Anthropic backend and other models to Vertex OpenAI.
|
||||
** One scoped ADC credential supplies both clients; access tokens stay in memory.
|
||||
*******************************************************************************/
|
||||
public class VertexModelProvider implements AssistantModelProvider
|
||||
{
|
||||
private final AnthropicModelProvider claude;
|
||||
private final OpenAiCompatibleModelProvider openai;
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Load Application Default Credentials without exposing credential failures.
|
||||
*******************************************************************************/
|
||||
public VertexModelProvider(ProviderSettings settings)
|
||||
{
|
||||
this(settings, applicationCredentials());
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Credential injection for deterministic, offline contract tests.
|
||||
*******************************************************************************/
|
||||
VertexModelProvider(ProviderSettings settings, GoogleCredentials credentials)
|
||||
{
|
||||
GoogleCredentials scoped = credentials.createScoped("https://www.googleapis.com/auth/cloud-platform");
|
||||
var backend = VertexBackend.builder().googleCredentials(scoped).region(settings.getRegion()).project(settings.getCloudProject());
|
||||
if(settings.getBaseUrl() != null)
|
||||
{
|
||||
backend.baseUrl(settings.getBaseUrl());
|
||||
}
|
||||
claude = AnthropicModelProvider.builder().client(AnthropicOkHttpClient.builder().backend(backend.build())
|
||||
.timeout(Timeout.builder().connect(Duration.ofSeconds(settings.getConnectTimeoutSeconds()))
|
||||
.request(Duration.ofSeconds(settings.getRequestTimeoutSeconds())).build()).build()).build();
|
||||
String base = settings.getBaseUrl() == null ? "https://" + settings.getRegion() + "-aiplatform.googleapis.com/v1/projects/"
|
||||
+ settings.getCloudProject() + "/locations/" + settings.getRegion() + "/endpoints/openapi" : settings.getBaseUrl();
|
||||
openai = new OpenAiCompatibleModelProvider(new ProviderSettings().withKind("VERTEX").withBaseUrl(base)
|
||||
.withConnectTimeoutSeconds(settings.getConnectTimeoutSeconds()).withRequestTimeoutSeconds(settings.getRequestTimeoutSeconds())
|
||||
.withTlsCaBundle(settings.getTlsCaBundle()).withSslContext(settings.getSslContext())
|
||||
.withBearerTokenSupplier(() -> accessToken(scoped)));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Resolve ADC lazily when a provider is created, not during registration.
|
||||
*******************************************************************************/
|
||||
private static GoogleCredentials applicationCredentials()
|
||||
{
|
||||
try
|
||||
{
|
||||
return (GoogleCredentials.getApplicationDefault());
|
||||
}
|
||||
catch(IOException e)
|
||||
{
|
||||
throw (new IllegalStateException("Google Application Default Credentials are unavailable. Configure GOOGLE_APPLICATION_CREDENTIALS or workload identity."));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Refresh expired credentials on every request; never retain a token in settings.
|
||||
*******************************************************************************/
|
||||
private static String accessToken(GoogleCredentials credentials)
|
||||
{
|
||||
try
|
||||
{
|
||||
credentials.refreshIfExpired();
|
||||
return (credentials.getAccessToken().getTokenValue());
|
||||
}
|
||||
catch(IOException e)
|
||||
{
|
||||
throw (new IllegalStateException("Could not refresh Google Application Default Credentials."));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Provider identity.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public String getName()
|
||||
{
|
||||
return ("vertex");
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Route using the registry's model id; caller requests are never mutated.
|
||||
*******************************************************************************/
|
||||
private AssistantModelProvider route(String model)
|
||||
{
|
||||
return (model != null && model.startsWith("claude-") ? claude : openai);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Capabilities follow the selected backend.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public ModelCapabilities getCapabilities(String modelId)
|
||||
{
|
||||
return (route(modelId).getCapabilities(modelId));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Stream through the existing adapters, preserving their error and tool mapping.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public StreamHandle stream(ModelRequest request, ModelEventListener listener)
|
||||
{
|
||||
return (route(request.getModel()).stream(request, listener));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Release the Anthropic client and its executor.
|
||||
*******************************************************************************/
|
||||
public void close()
|
||||
{
|
||||
claude.close();
|
||||
}
|
||||
}
|
||||
+157
@@ -0,0 +1,157 @@
|
||||
/*
|
||||
* Copyright 2026 Kingsrook, LLC
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package com.kingsrook.qbits.assistant.provider.vertex;
|
||||
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import com.kingsrook.qbits.assistant.model.AssistantProvider;
|
||||
import com.kingsrook.qbits.assistant.provider.AssistantModelProvider;
|
||||
import com.kingsrook.qbits.assistant.provider.AssistantProviderType;
|
||||
import com.kingsrook.qbits.assistant.provider.DiscoveredModel;
|
||||
import com.kingsrook.qbits.assistant.provider.ProviderSettings;
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Google Vertex AI (Claude, Gemini), using server credentials and region-derived endpoints.
|
||||
*******************************************************************************/
|
||||
public class VertexProviderType implements AssistantProviderType
|
||||
{
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Provider registry contract.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public String getType()
|
||||
{
|
||||
return ("vertex");
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Provider registry contract.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public String getLabel()
|
||||
{
|
||||
return ("Google Vertex AI (Claude, Gemini)");
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Provider registry contract.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public List<String> getModels()
|
||||
{
|
||||
return (List.of("claude-sonnet-4-5@20250929", "claude-haiku-4-5@20251001", "gemini-2.5-pro", "gemini-2.5-flash"));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Provider registry contract.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public String getDefaultApiKeyEnvName()
|
||||
{
|
||||
return (null);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Provider registry contract.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public String getDefaultBaseUrl()
|
||||
{
|
||||
return (null);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Provider registry contract.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public boolean requiresBaseUrl(String kind)
|
||||
{
|
||||
return (false);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Validate cloud settings without loading credentials.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public List<String> validate(AssistantProvider row)
|
||||
{
|
||||
List<String> errors = new ArrayList<>();
|
||||
if(row.getRegion() == null || row.getRegion().isBlank())
|
||||
{
|
||||
errors.add("Region is required for VERTEX.");
|
||||
}
|
||||
if(row.getApiKeyEnvName() != null && !row.getApiKeyEnvName().isBlank())
|
||||
{
|
||||
errors.add("Leave API key variable empty: Google Vertex AI (Claude, Gemini) uses Google Application Default Credentials (GOOGLE_APPLICATION_CREDENTIALS or workload identity).");
|
||||
}
|
||||
if(row.getCloudProject() == null || row.getCloudProject().isBlank())
|
||||
{
|
||||
errors.add("GCP project id (cloudProject) is required for VERTEX.");
|
||||
}
|
||||
return (errors);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Cloud providers require resolved region/project settings.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public AssistantModelProvider create(String apiKey, String baseUrl)
|
||||
{
|
||||
throw (new IllegalArgumentException("Use ProviderSettings with the required cloud region and project."));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Static defaults; this adapter does not query the Vertex Model Garden.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public List<DiscoveredModel> discoverModels(ProviderSettings settings)
|
||||
{
|
||||
return (getModels().stream().map(id -> new DiscoveredModel(id, id, null, null)).toList());
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Create the provider with server credentials.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public AssistantModelProvider create(ProviderSettings settings)
|
||||
{
|
||||
return (new VertexModelProvider(settings));
|
||||
}
|
||||
}
|
||||
+309
@@ -0,0 +1,309 @@
|
||||
/*
|
||||
* Copyright 2026 Kingsrook, LLC
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package com.kingsrook.qbits.assistant.provider.vertex;
|
||||
|
||||
|
||||
import java.io.IOException;
|
||||
import java.time.Clock;
|
||||
import java.time.Duration;
|
||||
import java.util.Date;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Set;
|
||||
import com.google.auth.oauth2.AccessToken;
|
||||
import com.google.auth.oauth2.GoogleCredentials;
|
||||
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
|
||||
import com.kingsrook.qbits.assistant.engine.ProviderService;
|
||||
import com.kingsrook.qbits.assistant.model.AssistantProvider;
|
||||
import com.kingsrook.qbits.assistant.provider.ModelMessage;
|
||||
import com.kingsrook.qbits.assistant.provider.ModelProviderException;
|
||||
import com.kingsrook.qbits.assistant.provider.ModelRequest;
|
||||
import com.kingsrook.qbits.assistant.provider.ModelToolDefinition;
|
||||
import com.kingsrook.qbits.assistant.provider.ProviderSettings;
|
||||
import com.kingsrook.qbits.assistant.testing.CloudProviderTestServer;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.params.ParameterizedTest;
|
||||
import org.junit.jupiter.params.provider.ValueSource;
|
||||
import static com.kingsrook.qbits.assistant.testing.CloudProviderTestServer.JSON;
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Real Vertex Claude and OpenAI-compatible transports using fake ADC tokens.
|
||||
*******************************************************************************/
|
||||
class VertexContractTest
|
||||
{
|
||||
static final String CLAUDE = "claude-sonnet-4-5@20250929";
|
||||
static final String GEMINI = "gemini-2.5-flash";
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Fixed fake token, valid well beyond the test duration.
|
||||
*******************************************************************************/
|
||||
static GoogleCredentials credentials()
|
||||
{
|
||||
return (GoogleCredentials.create(new AccessToken("fake-google-token", new Date(System.currentTimeMillis() + 3_600_000))));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Scoped connection settings with a loopback override.
|
||||
*******************************************************************************/
|
||||
static ProviderSettings settings(CloudProviderTestServer server)
|
||||
{
|
||||
return (new ProviderSettings().withKind("VERTEX").withRegion("us-east5").withCloudProject("test-project").withBaseUrl(server.url())
|
||||
.withConnectTimeoutSeconds(2).withRequestTimeoutSeconds(5));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** A strict request verifies that only the Claude backend preserves strict.
|
||||
*******************************************************************************/
|
||||
static ModelRequest request(String model)
|
||||
{
|
||||
return (new ModelRequest().withModel(model).withMaxTokens(1024).withMessages(List.of(ModelMessage.user("Count posts")))
|
||||
.withTools(List.of(new ModelToolDefinition().withName("count_records").withDescription("Count records").withStrict(true)
|
||||
.withInputSchema(Map.of("type", "object", "properties", Map.of("table", Map.of("type", "string")), "required", List.of("table"), "additionalProperties", false)))));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Text and a fragmented tool call in a Vertex OpenAI-compatible stream.
|
||||
*******************************************************************************/
|
||||
static String geminiSse(String text, Map<String, Object> input) throws IOException
|
||||
{
|
||||
StringBuilder events = new StringBuilder();
|
||||
events.append("data: ").append(JSON.writeValueAsString(Map.of("choices", List.of(Map.of("delta", Map.of("content", text)))))).append("\n\n");
|
||||
if(input != null)
|
||||
{
|
||||
String arguments = JSON.writeValueAsString(input);
|
||||
Integer middle = arguments.length() / 2;
|
||||
events.append("data: ").append(JSON.writeValueAsString(Map.of("choices", List.of(Map.of("delta", Map.of("tool_calls", List.of(
|
||||
Map.of("index", 0, "id", "count", "type", "function", "function", Map.of("name", "count_records", "arguments", arguments.substring(0, middle)))))))))).append("\n\n");
|
||||
events.append("data: ").append(JSON.writeValueAsString(Map.of("choices", List.of(Map.of("delta", Map.of("tool_calls", List.of(
|
||||
Map.of("index", 0, "function", Map.of("arguments", arguments.substring(middle)))))))))).append("\n\n");
|
||||
}
|
||||
events.append("data: ").append(JSON.writeValueAsString(Map.of("choices", List.of(Map.of("delta", Map.of(), "finish_reason", input == null ? "stop" : "tool_calls")),
|
||||
"usage", Map.of("prompt_tokens", 12, "completion_tokens", 7)))).append("\n\ndata: [DONE]\n\n");
|
||||
return (events.toString());
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Claude path, ADC bearer, Anthropic version, streaming text, tools and usage.
|
||||
*******************************************************************************/
|
||||
@Test
|
||||
void claudeContract() throws Exception
|
||||
{
|
||||
try(CloudProviderTestServer server = new CloudProviderTestServer())
|
||||
{
|
||||
server.sse(CloudProviderTestServer.anthropicSse(CLAUDE, "Checking.", Map.of("table", "blog_post")));
|
||||
var provider = new VertexModelProvider(settings(server), credentials());
|
||||
try
|
||||
{
|
||||
StringBuilder text = new StringBuilder();
|
||||
var response = provider.stream(request(CLAUDE), text::append).await(Duration.ofSeconds(10));
|
||||
assertThat(provider.getName()).isEqualTo("vertex");
|
||||
assertThat(provider.getCapabilities(CLAUDE).getPromptCaching()).isTrue();
|
||||
assertThat(text.toString()).isEqualTo("Checking.");
|
||||
assertThat(response.getToolUses()).singleElement().satisfies(tool ->
|
||||
{
|
||||
assertThat(tool.getToolUseId()).isEqualTo("count");
|
||||
assertThat(tool.getToolName()).isEqualTo("count_records");
|
||||
assertThat(tool.getInput()).containsEntry("table", "blog_post");
|
||||
});
|
||||
assertThat(response.getUsage().getInputTokens()).isEqualTo(12);
|
||||
assertThat(response.getUsage().getOutputTokens()).isEqualTo(7);
|
||||
var sent = server.requests.getFirst();
|
||||
assertThat(sent.path()).isEqualTo("/v1/projects/test-project/locations/us-east5/publishers/anthropic/models/" + CLAUDE + ":streamRawPredict");
|
||||
assertThat(sent.authorization()).isEqualTo("Bearer fake-google-token");
|
||||
var body = JSON.readTree(sent.body());
|
||||
assertThat(body.path("anthropic_version").asText()).isEqualTo("vertex-2023-10-16");
|
||||
assertThat(body.has("model")).isFalse();
|
||||
assertThat(body.path("tools").get(0).path("strict").asBoolean()).isTrue();
|
||||
}
|
||||
finally
|
||||
{
|
||||
provider.close();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Gemini, already-prefixed and other ids all use OpenAI without strict.
|
||||
*******************************************************************************/
|
||||
@ParameterizedTest
|
||||
@ValueSource(strings = { "gemini-2.5-flash", "google/gemini-2.5-pro", "other-model" })
|
||||
void geminiContractAndTokenRefresh(String model) throws Exception
|
||||
{
|
||||
try(CloudProviderTestServer server = new CloudProviderTestServer())
|
||||
{
|
||||
var credentials = new RefreshingCredentials();
|
||||
String endpoint = "/v1/projects/test-project/locations/us-east5/endpoints/openapi";
|
||||
var provider = new VertexModelProvider(settings(server).withBaseUrl(server.url() + endpoint), credentials);
|
||||
try
|
||||
{
|
||||
for(int call = 1; call <= 2; call++)
|
||||
{
|
||||
server.sse(geminiSse("Checking.", Map.of("table", "blog_post")));
|
||||
StringBuilder text = new StringBuilder();
|
||||
ModelRequest request = request(model);
|
||||
var response = provider.stream(request, text::append).await(Duration.ofSeconds(10));
|
||||
assertThat(text.toString()).isEqualTo("Checking.");
|
||||
assertThat(request.getModel()).isEqualTo(model);
|
||||
assertThat(response.getToolUses()).singleElement().satisfies(tool ->
|
||||
{
|
||||
assertThat(tool.getToolName()).isEqualTo("count_records");
|
||||
assertThat(tool.getInput()).containsEntry("table", "blog_post");
|
||||
});
|
||||
assertThat(response.getUsage().getInputTokens()).isEqualTo(12);
|
||||
assertThat(response.getUsage().getOutputTokens()).isEqualTo(7);
|
||||
var sent = server.requests.get(call - 1);
|
||||
assertThat(sent.path()).isEqualTo(endpoint + "/chat/completions");
|
||||
assertThat(sent.authorization()).isEqualTo("Bearer refreshed-" + call);
|
||||
var body = JSON.readTree(sent.body());
|
||||
assertThat(body.path("model").asText()).isEqualTo(model.contains("/") ? model : "google/" + model);
|
||||
assertThat(body.path("tools").get(0).path("function").has("strict")).isFalse();
|
||||
if(call == 1)
|
||||
{
|
||||
credentials.refresh();
|
||||
}
|
||||
}
|
||||
assertThat(credentials.refreshes).isEqualTo(2);
|
||||
assertThat(provider.getCapabilities(model).getPromptCaching()).isFalse();
|
||||
assertThat(provider.getCapabilities(null).getTools()).isTrue();
|
||||
}
|
||||
finally
|
||||
{
|
||||
provider.close();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Start expired so refreshIfExpired must obtain the first token itself.
|
||||
*******************************************************************************/
|
||||
static class RefreshingCredentials extends GoogleCredentials
|
||||
{
|
||||
private Integer refreshes = 0;
|
||||
|
||||
/*******************************************************************************
|
||||
** An expired fake token cannot authorize the first request.
|
||||
*******************************************************************************/
|
||||
RefreshingCredentials()
|
||||
{
|
||||
super(new AccessToken("expired", new Date(0)));
|
||||
}
|
||||
|
||||
/*******************************************************************************
|
||||
** Deterministic replacement tokens, without a token-server request.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public AccessToken refreshAccessToken() throws IOException
|
||||
{
|
||||
refreshes++;
|
||||
return (new AccessToken("refreshed-" + refreshes, new Date(System.currentTimeMillis() + 3_600_000)));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** ADC refresh failures cannot leak the credential exception's details.
|
||||
*******************************************************************************/
|
||||
@Test
|
||||
void refreshFailureIsSanitized() throws Exception
|
||||
{
|
||||
try(CloudProviderTestServer server = new CloudProviderTestServer())
|
||||
{
|
||||
var credentials = new RefreshingCredentials()
|
||||
{
|
||||
/*******************************************************************************
|
||||
** Simulate a token service failure containing confidential diagnostics.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public AccessToken refreshAccessToken() throws IOException
|
||||
{
|
||||
throw (new IOException("synthetic-private-token"));
|
||||
}
|
||||
};
|
||||
var provider = new VertexModelProvider(settings(server), credentials);
|
||||
try
|
||||
{
|
||||
assertThatThrownBy(() -> provider.stream(request(GEMINI), delta -> {}).await(Duration.ofSeconds(10)))
|
||||
.isInstanceOfSatisfying(ModelProviderException.class, error ->
|
||||
{
|
||||
assertThat(error.getKind()).isEqualTo(ModelProviderException.Kind.AUTHENTICATION);
|
||||
assertThat(error.toString()).doesNotContain("synthetic-private-token");
|
||||
assertThat(error.getCause()).isNull();
|
||||
});
|
||||
assertThat(server.requests).isEmpty();
|
||||
}
|
||||
finally
|
||||
{
|
||||
provider.close();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Cloud requirements, static discovery and URL allowlists need no ADC lookup.
|
||||
*******************************************************************************/
|
||||
@Test
|
||||
void settingsValidationAndDefaults() throws Exception
|
||||
{
|
||||
var type = new VertexProviderType();
|
||||
var config = new AssistantQBitConfig().withProviderType(type).withAllowedBaseUrls(Set.of("http://127.0.0.1:1234"));
|
||||
var service = new ProviderService(config, null, Clock.systemUTC());
|
||||
var row = new AssistantProvider().withProviderType("vertex");
|
||||
assertThat(service.validate(row)).anyMatch(error -> error.contains("Region is required")).anyMatch(error -> error.contains("project id"));
|
||||
row.withRegion(" ").withCloudProject(" ");
|
||||
assertThat(service.validate(row)).hasSize(2);
|
||||
row.withRegion("us-east5").withCloudProject("test-project");
|
||||
assertThat(service.validate(row)).isEmpty();
|
||||
row.withApiKeyEnvName("GOOGLE_API_KEY");
|
||||
assertThat(service.validate(row)).anyMatch(error -> error.contains("Application Default Credentials"));
|
||||
row.withApiKeyEnvName(" ").withBaseUrl("https://unlisted.example");
|
||||
assertThat(service.validate(row)).anyMatch(error -> error.contains("not allowed"));
|
||||
row.withBaseUrl("http://127.0.0.1:1234");
|
||||
assertThat(service.validate(row)).isEmpty();
|
||||
assertThat(type.getType()).isEqualTo("vertex");
|
||||
assertThat(type.getKinds()).containsExactly("VERTEX");
|
||||
assertThat(type.getLabel()).isEqualTo("Google Vertex AI (Claude, Gemini)");
|
||||
assertThat(type.getDefaultLocation("VERTEX")).isEqualTo("CLOUD");
|
||||
assertThat(type.getDefaultApiKeyEnvName()).isNull();
|
||||
assertThat(type.getDefaultBaseUrl()).isNull();
|
||||
assertThat(type.discoverModels(new ProviderSettings())).extracting(model -> model.modelId()).containsExactlyElementsOf(type.getModels());
|
||||
assertThat(type.checkHealth(new ProviderSettings()).healthy()).isNull();
|
||||
assertThatThrownBy(() -> type.create("key", null)).isInstanceOf(IllegalArgumentException.class);
|
||||
// Exercise endpoint derivation without sending any request to that endpoint.
|
||||
var provider = new VertexModelProvider(new ProviderSettings().withRegion("us-east5").withCloudProject("test-project"), credentials());
|
||||
provider.close();
|
||||
}
|
||||
}
|
||||
+100
@@ -0,0 +1,100 @@
|
||||
/*
|
||||
* Copyright 2026 Kingsrook, LLC
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package com.kingsrook.qbits.assistant.provider.vertex;
|
||||
|
||||
|
||||
import com.kingsrook.qbits.assistant.AssistantModelRouting;
|
||||
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
|
||||
import com.kingsrook.qbits.assistant.testing.AssistantTestFixture;
|
||||
import com.kingsrook.qbits.assistant.testing.CloudProviderTestServer;
|
||||
import com.kingsrook.qbits.assistant.testing.GoldenTasks;
|
||||
import com.kingsrook.qbits.assistant.testing.TurnRunner;
|
||||
import com.kingsrook.qqq.backend.core.context.QContext;
|
||||
import org.junit.jupiter.api.Tag;
|
||||
import org.junit.jupiter.params.ParameterizedTest;
|
||||
import org.junit.jupiter.params.provider.ValueSource;
|
||||
import static com.kingsrook.qbits.assistant.testing.CloudProviderTestServer.JSON;
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Golden task through the cloud adapter, turn engine and permissioned UI tools.
|
||||
*******************************************************************************/
|
||||
@Tag("golden")
|
||||
class VertexGoldenTurnTest
|
||||
{
|
||||
/*******************************************************************************
|
||||
** A real count_records round trip returns the tool result to the model.
|
||||
*******************************************************************************/
|
||||
@ParameterizedTest
|
||||
@ValueSource(strings = { VertexContractTest.CLAUDE, VertexContractTest.GEMINI })
|
||||
void goldenToolRoundTrip(String model) throws Exception
|
||||
{
|
||||
try(CloudProviderTestServer server = new CloudProviderTestServer())
|
||||
{
|
||||
String answer = "There are " + AssistantTestFixture.PUBLISHED_POSTS + " published posts. /app/blog_post";
|
||||
boolean claude = model.startsWith("claude-");
|
||||
server.sse(claude ? CloudProviderTestServer.anthropicSse(model, "", CloudProviderTestServer.countInput()) : VertexContractTest.geminiSse("", CloudProviderTestServer.countInput()));
|
||||
server.sse(claude ? CloudProviderTestServer.anthropicSse(model, answer, null) : VertexContractTest.geminiSse(answer, null));
|
||||
var provider = new VertexModelProvider(VertexContractTest.settings(server), VertexContractTest.credentials());
|
||||
var config = new AssistantQBitConfig().withProviderName(provider.getName()).withModelProvider(provider)
|
||||
.withModels(new AssistantModelRouting().withMainModel(model));
|
||||
try
|
||||
{
|
||||
var instance = AssistantTestFixture.build(config);
|
||||
var task = GoldenTasks.all().stream().filter(t -> t.id().equals("count-published-posts")).findFirst().orElseThrow();
|
||||
var outcome = new TurnRunner(instance, config).run(task.userId(), task.prompt());
|
||||
assertThat(outcome.errorCode()).isNull();
|
||||
task.check().accept(outcome);
|
||||
assertThat(outcome.text()).endsWith(answer);
|
||||
assertThat(outcome.end()).isNotEmpty();
|
||||
assertThat(server.requests).hasSize(2);
|
||||
var messages = JSON.readTree(server.requests.get(1).body()).path("messages");
|
||||
var assistant = messages.get(messages.size() - 2);
|
||||
var result = messages.get(messages.size() - 1);
|
||||
assertThat(assistant.path("role").asText()).isEqualTo("assistant");
|
||||
if(claude)
|
||||
{
|
||||
assertThat(assistant.path("content").toString()).contains("count_records", "count");
|
||||
assertThat(result.path("role").asText()).isEqualTo("user");
|
||||
var toolResult = result.path("content").get(0);
|
||||
assertThat(toolResult.path("type").asText()).isEqualTo("tool_result");
|
||||
assertThat(toolResult.path("tool_use_id").asText()).isEqualTo("count");
|
||||
assertThat(toolResult.path("is_error").asBoolean()).isFalse();
|
||||
assertThat(toolResult.path("content").toString()).contains("count");
|
||||
}
|
||||
else
|
||||
{
|
||||
assertThat(assistant.path("tool_calls")).hasSize(1);
|
||||
assertThat(assistant.path("tool_calls").get(0).path("id").asText()).isEqualTo("count");
|
||||
assertThat(result.path("role").asText()).isEqualTo("tool");
|
||||
assertThat(result.path("tool_call_id").asText()).isEqualTo("count");
|
||||
assertThat(result.path("content").asText()).contains("count").doesNotStartWith("ERROR:");
|
||||
}
|
||||
}
|
||||
finally
|
||||
{
|
||||
if(config.getRuntime() != null)
|
||||
{
|
||||
config.getRuntime().close();
|
||||
}
|
||||
provider.close();
|
||||
QContext.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+5
-4
@@ -214,7 +214,7 @@ public class ProviderService
|
||||
throw (new ProviderSetupException("The selected model is disabled in this provider's model registry."));
|
||||
}
|
||||
String cacheKey = row.getId() + "|" + row.getProviderType() + "|" + settings.getKind() + "|" + settings.getBaseUrl()
|
||||
+ "|" + settings.getApiVersion() + "|" + settings.getRegion() + "|" + settings.getConnectTimeoutSeconds()
|
||||
+ "|" + settings.getApiVersion() + "|" + settings.getRegion() + "|" + settings.getCloudProject() + "|" + settings.getConnectTimeoutSeconds()
|
||||
+ "|" + settings.getTlsCaBundle() + "|" + settings.getRequestTimeoutSeconds() + "|" + fingerprint(settings.getApiKey());
|
||||
AssistantModelProvider provider = built.computeIfAbsent(cacheKey, k -> type.create(settings));
|
||||
return (new Active(row, row.getProviderType(), model, registry == null ? provider : new RegistryModelProvider(provider, registry)));
|
||||
@@ -236,6 +236,7 @@ public class ProviderService
|
||||
return (errors);
|
||||
}
|
||||
|
||||
errors.addAll(type.validate(row));
|
||||
String envName = keyEnvName(row);
|
||||
if(envName != null && !AssistantQBitConfig.ENV_NAME.matcher(envName).matches())
|
||||
{
|
||||
@@ -276,7 +277,7 @@ public class ProviderService
|
||||
{
|
||||
errors.add("Location must be LOCAL or CLOUD.");
|
||||
}
|
||||
if(!StringUtils.hasContent(row.getBaseUrl()) && type.getDefaultBaseUrl(kind(row)) == null)
|
||||
if(!StringUtils.hasContent(row.getBaseUrl()) && type.requiresBaseUrl(kind(row)))
|
||||
{
|
||||
errors.add("An allow-listed base URL is required for this provider kind.");
|
||||
}
|
||||
@@ -313,7 +314,7 @@ public class ProviderService
|
||||
String envName = keyEnvName(row);
|
||||
if(envName == null)
|
||||
{
|
||||
return ("Not needed (local, no key)");
|
||||
return ("Not needed (no API key)");
|
||||
}
|
||||
if(!AssistantQBitConfig.ENV_NAME.matcher(envName).matches() || !config.getEffectiveAllowedApiKeyEnvNames().contains(envName))
|
||||
{
|
||||
@@ -653,7 +654,7 @@ public class ProviderService
|
||||
}
|
||||
return (new ProviderSettings().withKind(kind(row)).withApiKey(key)
|
||||
.withBaseUrl(StringUtils.hasContent(row.getBaseUrl()) ? row.getBaseUrl().trim() : null)
|
||||
.withApiVersion(row.getApiVersion()).withRegion(row.getRegion())
|
||||
.withApiVersion(row.getApiVersion()).withRegion(row.getRegion()).withCloudProject(row.getCloudProject())
|
||||
.withTlsCaBundle(row.getTlsCaBundle()).withSslContext(StringUtils.hasContent(row.getTlsCaBundle()) ? config.getTrustedCaBundles().get(row.getTlsCaBundle()) : null)
|
||||
.withConnectTimeoutSeconds(row.getConnectTimeoutSeconds() == null ? 10 : row.getConnectTimeoutSeconds())
|
||||
.withRequestTimeoutSeconds(row.getRequestTimeoutSeconds() == null ? 120 : row.getRequestTimeoutSeconds()));
|
||||
|
||||
+7
-3
@@ -252,7 +252,7 @@ public final class AssistantRecordScreens
|
||||
types.add(new QPossibleValue<>(config.getProviderName(), config.getProviderName()));
|
||||
}
|
||||
List<QPossibleValue<?>> kinds = new ArrayList<>();
|
||||
Set<String> kindNames = new LinkedHashSet<>(List.of("ANTHROPIC", "OPENAI", "AZURE_OPENAI", "OPENROUTER", "OLLAMA", "LM_STUDIO", "VLLM", "LLAMA_CPP", "OPENAI_COMPATIBLE"));
|
||||
Set<String> kindNames = new LinkedHashSet<>(List.of("ANTHROPIC", "OPENAI", "AZURE_OPENAI", "OPENROUTER", "OLLAMA", "LM_STUDIO", "VLLM", "LLAMA_CPP", "OPENAI_COMPATIBLE", "BEDROCK", "VERTEX"));
|
||||
config.getProviderTypes().values().forEach(t -> kindNames.addAll(t.getKinds()));
|
||||
kindNames.forEach(k -> kinds.add(new QPossibleValue<>(k, providerKindLabel(k))));
|
||||
List<QPossibleValue<?>> modelValues = new ArrayList<>();
|
||||
@@ -292,6 +292,8 @@ public final class AssistantRecordScreens
|
||||
case "VLLM" -> "vLLM";
|
||||
case "LLAMA_CPP" -> "llama.cpp server";
|
||||
case "OPENAI_COMPATIBLE" -> "OpenAI-compatible (other)";
|
||||
case "BEDROCK" -> "AWS Bedrock";
|
||||
case "VERTEX" -> "Google Vertex AI";
|
||||
default -> kind;
|
||||
});
|
||||
}
|
||||
@@ -316,6 +318,8 @@ public final class AssistantRecordScreens
|
||||
table.getField("apiKeyEnvName").withLabel("API Key Variable");
|
||||
table.getField("tlsCaBundle").withLabel("TLS CA bundle");
|
||||
table.getField("baseUrl").withLabel("Base URL Override");
|
||||
table.getField("cloudProject").withLabel("GCP Project ID");
|
||||
help(table, "cloudProject", "Required for Vertex: the Google Cloud project id, not a credential.");
|
||||
table.getField("monthlyCostCapUsd").withLabel("Monthly Cost Cap (USD)").withDisplayFormat(DisplayFormat.CURRENCY);
|
||||
counts(table, "maxOutputTokensPerRequest", "lastTestLatencyMs");
|
||||
table.getField("lastTestLatencyMs").withLabel("Last test latency (ms)");
|
||||
@@ -325,7 +329,7 @@ public final class AssistantRecordScreens
|
||||
help(table, "enabled", "Off: the assistant never uses this provider.");
|
||||
help(table, "isDefault", "The assistant uses the default enabled provider (else the first enabled one). Marking one default unmarks the others.");
|
||||
help(table, "apiKeyEnvName", "The NAME of the environment variable that holds the API key (e.g. ANTHROPIC_API_KEY), never the key itself. "
|
||||
+ "The key stays in the server's secret store; only allowed variable names are accepted. Empty uses the type's default variable. To replace the key, change it in the secret store.");
|
||||
+ "The key stays in the server's secret store; only allowed variable names are accepted. Empty uses the type's default variable. Bedrock and Vertex require this field to be empty and use server credentials. To replace the key, change it in the secret store.");
|
||||
help(table, "tlsCaBundle", "Name of a CA bundle configured on the server (withTrustedCaBundle). Leave empty to use the JVM trust store.");
|
||||
help(table, "baseUrl", "Leave empty to use the provider's own API. Other URLs must be allowed in the server's configuration.");
|
||||
help(table, "monthlyCostCapUsd", "When this provider's cost this month (UTC) reaches the cap, new requests are refused. Empty: no cap.");
|
||||
@@ -334,7 +338,7 @@ public final class AssistantRecordScreens
|
||||
table.getField(ProviderService.KEY_STATUS_FIELD).withLabel("API Key").withHelpContent(new QHelpContent("Whether the key variable is set on this server, with the key's last four characters. The key itself is never shown or returned."));
|
||||
sections(table,
|
||||
section("provider", "Provider", "settings_input_component", Tier.T1, "id", "name", "providerType", "model", "enabled", "isDefault"),
|
||||
section("connection", "Connection", "key", Tier.T2, "apiKeyEnvName", ProviderService.KEY_STATUS_FIELD, "baseUrl", "tlsCaBundle"),
|
||||
section("connection", "Connection", "key", Tier.T2, "apiKeyEnvName", ProviderService.KEY_STATUS_FIELD, "baseUrl", "tlsCaBundle", "cloudProject"),
|
||||
section("limits", "Limits", "speed", Tier.T2, "monthlyCostCapUsd", "maxOutputTokensPerRequest", "maxConcurrentRequests"),
|
||||
section("registry", "Provider settings", "settings", Tier.T2, "kind", "location", "hipaaEligible", "region", "apiVersion", "connectTimeoutSeconds", "requestTimeoutSeconds"),
|
||||
section("discovery", "Last discovery", "memory", Tier.T2, "lastDiscoveryAt", "lastDiscoveryStatus", "lastDiscoveryError"),
|
||||
|
||||
+34
@@ -35,6 +35,9 @@ import com.kingsrook.qqq.backend.core.model.data.QRecordEntity;
|
||||
*******************************************************************************/
|
||||
public class AssistantProvider extends QRecordEntity
|
||||
{
|
||||
@QField(maxLength = 100)
|
||||
private String cloudProject;
|
||||
|
||||
@QField(maxLength = 30)
|
||||
private String kind;
|
||||
|
||||
@@ -1082,4 +1085,35 @@ public class AssistantProvider extends QRecordEntity
|
||||
this.tlsCaBundle = tlsCaBundle;
|
||||
return (this);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Getter for cloudProject
|
||||
*******************************************************************************/
|
||||
public String getCloudProject()
|
||||
{
|
||||
return (this.cloudProject);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Setter for cloudProject
|
||||
*******************************************************************************/
|
||||
public void setCloudProject(String cloudProject)
|
||||
{
|
||||
this.cloudProject = cloudProject;
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Fluent setter for cloudProject
|
||||
*******************************************************************************/
|
||||
public AssistantProvider withCloudProject(String cloudProject)
|
||||
{
|
||||
this.cloudProject = cloudProject;
|
||||
return (this);
|
||||
}
|
||||
}
|
||||
+22
@@ -18,6 +18,7 @@ package com.kingsrook.qbits.assistant.provider;
|
||||
|
||||
|
||||
import java.util.List;
|
||||
import com.kingsrook.qbits.assistant.model.AssistantProvider;
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
@@ -135,4 +136,25 @@ public interface AssistantProviderType
|
||||
return (new HealthResult(null, null, "Health check not supported"));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Additional settings errors, without resolving credentials or doing I/O.
|
||||
*******************************************************************************/
|
||||
default List<String> validate(AssistantProvider row)
|
||||
{
|
||||
return (List.of());
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Whether this kind needs an explicit URL when it has no fixed default.
|
||||
** Cloud SDKs can instead derive their endpoint from region and project.
|
||||
*******************************************************************************/
|
||||
default boolean requiresBaseUrl(String kind)
|
||||
{
|
||||
return (getDefaultBaseUrl(kind) == null);
|
||||
}
|
||||
|
||||
}
|
||||
+11
-2
@@ -79,20 +79,29 @@ public class CompletableStreamHandle implements StreamHandle
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Finish with a response.
|
||||
** Finish with a response (ignored once cancelled: cancel() answers with a
|
||||
** cancelled response, and closing the stream may wake the worker first).
|
||||
*******************************************************************************/
|
||||
public void complete(ModelResponse response)
|
||||
{
|
||||
if(cancelled.get())
|
||||
{
|
||||
return;
|
||||
}
|
||||
future.complete(response);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Finish with an error.
|
||||
** Finish with an error (ignored once cancelled, see complete()).
|
||||
*******************************************************************************/
|
||||
public void fail(ModelProviderException exception)
|
||||
{
|
||||
if(cancelled.get())
|
||||
{
|
||||
return;
|
||||
}
|
||||
future.completeExceptionally(exception);
|
||||
}
|
||||
|
||||
|
||||
+71
-1
@@ -17,15 +17,22 @@
|
||||
package com.kingsrook.qbits.assistant.provider;
|
||||
|
||||
|
||||
import java.util.function.Supplier;
|
||||
import javax.net.ssl.SSLContext;
|
||||
import com.fasterxml.jackson.annotation.JsonIgnore;
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Resolved connection settings. The API key is never included in diagnostics.
|
||||
** Resolved connection settings. API keys and bearer-token suppliers are never
|
||||
** included in diagnostics. A supplier resolves a fresh token for each request.
|
||||
*******************************************************************************/
|
||||
public class ProviderSettings
|
||||
{
|
||||
@JsonIgnore
|
||||
private Supplier<String> bearerTokenSupplier;
|
||||
|
||||
private String cloudProject;
|
||||
|
||||
private String tlsCaBundle;
|
||||
@JsonIgnore
|
||||
private transient SSLContext sslContext;
|
||||
@@ -327,4 +334,67 @@ public class ProviderSettings
|
||||
this.sslContext = sslContext;
|
||||
return (this);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Getter for cloudProject
|
||||
*******************************************************************************/
|
||||
public String getCloudProject()
|
||||
{
|
||||
return (this.cloudProject);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Setter for cloudProject
|
||||
*******************************************************************************/
|
||||
public void setCloudProject(String cloudProject)
|
||||
{
|
||||
this.cloudProject = cloudProject;
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Fluent setter for cloudProject
|
||||
*******************************************************************************/
|
||||
public ProviderSettings withCloudProject(String cloudProject)
|
||||
{
|
||||
this.cloudProject = cloudProject;
|
||||
return (this);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Getter for bearerTokenSupplier
|
||||
*******************************************************************************/
|
||||
@JsonIgnore
|
||||
public Supplier<String> getBearerTokenSupplier()
|
||||
{
|
||||
return (this.bearerTokenSupplier);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Setter for bearerTokenSupplier
|
||||
*******************************************************************************/
|
||||
public void setBearerTokenSupplier(Supplier<String> bearerTokenSupplier)
|
||||
{
|
||||
this.bearerTokenSupplier = bearerTokenSupplier;
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Fluent setter for bearerTokenSupplier
|
||||
*******************************************************************************/
|
||||
public ProviderSettings withBearerTokenSupplier(Supplier<String> bearerTokenSupplier)
|
||||
{
|
||||
this.bearerTokenSupplier = bearerTokenSupplier;
|
||||
return (this);
|
||||
}
|
||||
}
|
||||
+27
-1
@@ -262,7 +262,7 @@ class ProviderRegistryTest extends BaseAssistantTest
|
||||
assertThat(row.getHipaaEligible()).isFalse();
|
||||
assertThat(row.getConnectTimeoutSeconds()).isEqualTo(10);
|
||||
assertThat(row.getRequestTimeoutSeconds()).isEqualTo(120);
|
||||
assertThat(service.keyStatus(row)).isEqualTo("Not needed (local, no key)");
|
||||
assertThat(service.keyStatus(row)).isEqualTo("Not needed (no API key)");
|
||||
assertThat(service.forRow(row, "legacy").model()).isEqualTo("legacy");
|
||||
service.forRow(row, "legacy");
|
||||
assertThat(creates).isEqualTo(1);
|
||||
@@ -429,4 +429,30 @@ class ProviderRegistryTest extends BaseAssistantTest
|
||||
input.setFrontendStepBehavior(RunProcessInput.FrontendStepBehavior.SKIP);
|
||||
new RunProcessAction().execute(input);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Cloud project survives record mapping and changing it rebuilds the client.
|
||||
*******************************************************************************/
|
||||
@Test
|
||||
void cloudProjectRoundTripAndCache() throws Exception
|
||||
{
|
||||
AssistantProvider row = row();
|
||||
row.setCloudProject("project-one");
|
||||
var recordCopy = new AssistantProvider(row.toQRecord());
|
||||
assertThat(recordCopy.getCloudProject()).isEqualTo("project-one");
|
||||
var service = config.getRuntime().getProviderService();
|
||||
assertThat(service.settings(row).getCloudProject()).isEqualTo("project-one");
|
||||
assertThat(qInstance.getTable(AssistantProvider.TABLE_NAME).getField("cloudProject").getMaxLength()).isEqualTo(100);
|
||||
service.forRow(row, "legacy");
|
||||
service.forRow(row, "legacy");
|
||||
assertThat(creates).isEqualTo(1);
|
||||
row.withCloudProject("project-two");
|
||||
service.forRow(row, "legacy");
|
||||
assertThat(creates).isEqualTo(2);
|
||||
var settings = new ProviderSettings();
|
||||
settings.setCloudProject("project-three");
|
||||
assertThat(settings.getCloudProject()).isEqualTo("project-three");
|
||||
}
|
||||
}
|
||||
+200
@@ -0,0 +1,200 @@
|
||||
/*
|
||||
* Copyright 2026 Kingsrook, LLC
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package com.kingsrook.qbits.assistant.testing;
|
||||
|
||||
|
||||
import java.io.IOException;
|
||||
import java.net.InetSocketAddress;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.CopyOnWriteArrayList;
|
||||
import java.util.concurrent.LinkedBlockingQueue;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import com.sun.net.httpserver.HttpExchange;
|
||||
import com.sun.net.httpserver.HttpServer;
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Loopback-only cloud protocol fixture shared through the QBit test jar.
|
||||
*******************************************************************************/
|
||||
public class CloudProviderTestServer implements AutoCloseable
|
||||
{
|
||||
public static final ObjectMapper JSON = new ObjectMapper();
|
||||
public final List<Captured> requests = new CopyOnWriteArrayList<>();
|
||||
private final LinkedBlockingQueue<Reply> replies = new LinkedBlockingQueue<>();
|
||||
private final HttpServer server;
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Only synthetic credentials are captured by cloud tests.
|
||||
*******************************************************************************/
|
||||
public record Captured(String path, String authorization, String body)
|
||||
{
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Binary event-stream and SSE use the same scripted HTTP transport.
|
||||
*******************************************************************************/
|
||||
private record Reply(Integer status, String contentType, byte[] body)
|
||||
{
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Listen exclusively on loopback with an ephemeral port.
|
||||
*******************************************************************************/
|
||||
public CloudProviderTestServer() throws IOException
|
||||
{
|
||||
server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
|
||||
server.createContext("/", this::serve);
|
||||
server.start();
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** The root URL; each backend appends its own API path.
|
||||
*******************************************************************************/
|
||||
public String url()
|
||||
{
|
||||
return ("http://127.0.0.1:" + server.getAddress().getPort());
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Queue a response without contacting a real cloud service.
|
||||
*******************************************************************************/
|
||||
public void enqueue(Integer status, String contentType, byte[] body)
|
||||
{
|
||||
replies.add(new Reply(status, contentType, body));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Queue a successful SSE stream.
|
||||
*******************************************************************************/
|
||||
public void sse(String events)
|
||||
{
|
||||
enqueue(200, "text/event-stream", events.getBytes(StandardCharsets.UTF_8));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Capture the request before returning the scripted response.
|
||||
*******************************************************************************/
|
||||
private void serve(HttpExchange exchange) throws IOException
|
||||
{
|
||||
try(exchange)
|
||||
{
|
||||
requests.add(new Captured(exchange.getRequestURI().getPath(), exchange.getRequestHeaders().getFirst("Authorization"),
|
||||
new String(exchange.getRequestBody().readAllBytes(), StandardCharsets.UTF_8)));
|
||||
Reply reply = replies.poll();
|
||||
if(reply == null)
|
||||
{
|
||||
exchange.sendResponseHeaders(500, -1);
|
||||
return;
|
||||
}
|
||||
exchange.getResponseHeaders().set("Content-Type", reply.contentType());
|
||||
if("application/vnd.amazon.eventstream".equals(reply.contentType()))
|
||||
{
|
||||
exchange.getResponseHeaders().set("x-amzn-bedrock-content-type", "application/json");
|
||||
}
|
||||
exchange.sendResponseHeaders(reply.status(), reply.body().length);
|
||||
exchange.getResponseBody().write(reply.body());
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Anthropic event JSON, including fragmented tool arguments and usage.
|
||||
*******************************************************************************/
|
||||
public static List<String> anthropicEvents(String model, String text, Map<String, Object> toolInput) throws IOException
|
||||
{
|
||||
List<Map<String, Object>> events = new ArrayList<>();
|
||||
events.add(Map.of("type", "message_start", "message", Map.of("id", "msg_mock", "type", "message", "role", "assistant", "model", model,
|
||||
"content", List.of(), "usage", Map.of("input_tokens", 12, "output_tokens", 1))));
|
||||
events.add(Map.of("type", "content_block_start", "index", 0, "content_block", Map.of("type", "text", "text", "")));
|
||||
events.add(Map.of("type", "content_block_delta", "index", 0, "delta", Map.of("type", "text_delta", "text", text)));
|
||||
events.add(Map.of("type", "content_block_stop", "index", 0));
|
||||
if(toolInput != null)
|
||||
{
|
||||
events.add(Map.of("type", "content_block_start", "index", 1, "content_block", Map.of("type", "tool_use", "id", "count", "name", "count_records", "input", Map.of())));
|
||||
String arguments = JSON.writeValueAsString(toolInput);
|
||||
Integer middle = arguments.length() / 2;
|
||||
for(String fragment : List.of(arguments.substring(0, middle), arguments.substring(middle)))
|
||||
{
|
||||
events.add(Map.of("type", "content_block_delta", "index", 1, "delta", Map.of("type", "input_json_delta", "partial_json", fragment)));
|
||||
}
|
||||
events.add(Map.of("type", "content_block_stop", "index", 1));
|
||||
}
|
||||
events.add(Map.of("type", "message_delta", "delta", Map.of("stop_reason", toolInput == null ? "end_turn" : "tool_use"), "usage", Map.of("output_tokens", 7)));
|
||||
events.add(Map.of("type", "message_stop"));
|
||||
List<String> result = new ArrayList<>();
|
||||
for(Map<String, Object> event : events)
|
||||
{
|
||||
result.add(JSON.writeValueAsString(event));
|
||||
}
|
||||
return (result);
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Wrap Anthropic JSON events in standard SSE framing for Vertex.
|
||||
*******************************************************************************/
|
||||
public static String anthropicSse(String model, String text, Map<String, Object> toolInput) throws IOException
|
||||
{
|
||||
StringBuilder result = new StringBuilder();
|
||||
for(String event : anthropicEvents(model, text, toolInput))
|
||||
{
|
||||
result.append("event: ").append(JSON.readTree(event).path("type").asText()).append("\ndata: ").append(event).append("\n\n");
|
||||
}
|
||||
return (result.toString());
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Published-post filter shared by the cloud golden turns.
|
||||
*******************************************************************************/
|
||||
public static Map<String, Object> countInput()
|
||||
{
|
||||
return (Map.of("table", "blog_post", "filter", Map.of("booleanOperator", "AND", "criteria",
|
||||
List.of(Map.of("field", "status", "operator", "EQUALS", "values", List.of("PUBLISHED"))))));
|
||||
}
|
||||
|
||||
|
||||
|
||||
/*******************************************************************************
|
||||
** Stop the mock listener.
|
||||
*******************************************************************************/
|
||||
@Override
|
||||
public void close()
|
||||
{
|
||||
server.stop(0);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user