Skip to content

Commit

Permalink
Add Bedrock contract tests. (#855)
Browse files Browse the repository at this point in the history
* Add Bedrock contract tests.

* address comments.

* Address comments.

* Delete un-useful text file

* Clean up.

* Addressing non-blocking comments.

* Apply spotless.

---------

Co-authored-by: Zhonghao Zhao <zzhlogin@amazon.com>
  • Loading branch information
mxiamxia and zzhlogin authored Aug 16, 2024
1 parent faa9a08 commit 5d211eb
Show file tree
Hide file tree
Showing 12 changed files with 1,164 additions and 7 deletions.

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,25 @@ protected String getKinesisSpanNamePrefix() {
return "Kinesis";
}

protected String getBedrockSpanNamePrefix() {
return "Bedrock";
}

@Override
protected String getBedrockAgentSpanNamePrefix() {
return "AWSBedrockAgent";
}

@Override
protected String getBedrockRuntimeSpanNamePrefix() {
return "BedrockRuntime";
}

@Override
protected String getBedrockAgentRuntimeSpanNamePrefix() {
return "AWSBedrockAgentRuntime";
}

protected String getS3RpcServiceName() {
return "Amazon S3";
}
Expand All @@ -75,6 +94,22 @@ protected String getKinesisRpcServiceName() {
return "AmazonKinesis";
}

protected String getBedrockRpcServiceName() {
return "AmazonBedrock";
}

protected String getBedrockAgentRpcServiceName() {
return "AWSBedrockAgent";
}

protected String getBedrockRuntimeRpcServiceName() {
return "AmazonBedrockRuntime";
}

protected String getBedrockAgentRuntimeRpcServiceName() {
return "AWSBedrockAgentRuntime";
}

@Test
void testS3CreateBucket() throws Exception {
doTestS3CreateBucket();
Expand Down Expand Up @@ -165,4 +200,39 @@ void testKinsesisError() throws Exception {
void testKinesisFault() throws Exception {
doTestKinesisFault();
}

@Test
void testBedrockAgentGetKnowledgeBaseId() {
doTestBedrockAgentKnowledgeBaseId();
}

@Test
void testBedrockAgentAgentId() {
doTestBedrockAgentAgentId();
}

@Test
void testBedrockAgentDataSourceId() {
doTestBedrockAgentDataSourceId();
}

@Test
void testBedrockRuntimeModelId() {
doTestBedrockRuntimeModelId();
}

@Test
void testBedrockGuardrailId() {
doTestBedrockGuardrailId();
}

@Test
void testBedrockAgentRuntimeAgentId() {
doTestBedrockAgentRuntimeAgentId();
}

@Test
void testBedrockAgentRuntimeKnowledgeBaseId() {
doTestBedrockAgentRuntimeKnowledgeBaseId();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,25 @@ protected String getKinesisSpanNamePrefix() {
return "Kinesis";
}

protected String getBedrockSpanNamePrefix() {
return "Bedrock";
}

@Override
protected String getBedrockAgentSpanNamePrefix() {
return "BedrockAgent";
}

@Override
protected String getBedrockRuntimeSpanNamePrefix() {
return "BedrockRuntime";
}

@Override
protected String getBedrockAgentRuntimeSpanNamePrefix() {
return "BedrockAgentRuntime";
}

@Override
protected String getS3RpcServiceName() {
return "S3";
Expand All @@ -75,6 +94,26 @@ protected String getKinesisRpcServiceName() {
return "Kinesis";
}

@Override
protected String getBedrockRpcServiceName() {
return "Bedrock";
}

@Override
protected String getBedrockAgentRpcServiceName() {
return "BedrockAgent";
}

@Override
protected String getBedrockRuntimeRpcServiceName() {
return "BedrockRuntime";
}

@Override
protected String getBedrockAgentRuntimeRpcServiceName() {
return "BedrockAgentRuntime";
}

@Test
void testS3CreateBucket() throws Exception {
doTestS3CreateBucket();
Expand Down Expand Up @@ -164,4 +203,41 @@ void testKinesisError() throws Exception {
void testKinesisFault() throws Exception {
doTestKinesisFault();
}

@Test
void testBedrockAgentGetKnowlesgeBaseId() {
doTestBedrockAgentKnowledgeBaseId();
}

@Test
void testBedrockAgentAgentId() {
doTestBedrockAgentAgentId();
}

@Test
void testBedrockAgentDataSourceId() {
doTestBedrockAgentDataSourceId();
}

@Test
void testBedrockRuntimeModelId() {
doTestBedrockRuntimeModelId();
}

@Test
void testBedrockGuardrailId() {
doTestBedrockGuardrailId();
}

@Test
void testBedrockAgentRuntimeAgentId() {
doTestBedrockAgentRuntimeAgentId();
}

// TODO: Enable testBedrockAgentRuntimeKnowledgeBaseId test after KnowledgeBaseId is supported in
// OTEL BedrockAgentRuntime instrumentation
// @Test
// void testBedrockAgentRuntimeKnowledgeBaseId() {
// doTestBedrockAgentRuntimeKnowledgeBaseId();
// }
}
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,11 @@ public class SemanticConventionsConstants {
public static final String AWS_QUEUE_URL = "aws.queue.url";
public static final String AWS_QUEUE_NAME = "aws.queue.name";
public static final String AWS_STREAM_NAME = "aws.stream.name";
public static final String AWS_KNOWLEDGE_BASE_ID = "aws.bedrock.knowledge_base.id";
public static final String AWS_DATA_SOURCE_ID = "aws.bedrock.data_source.id";
public static final String AWS_AGENT_ID = "aws.bedrock.agent.id";
public static final String AWS_GUARDRAIL_ID = "aws.bedrock.guardrail.id";
public static final String GEN_AI_REQUEST_MODEL = "gen_ai.request.model";

// kafka
public static final String MESSAGING_CLIENT_ID = "messaging.client_id";
Expand Down
29 changes: 29 additions & 0 deletions appsignals-tests/images/aws-sdk/aws-sdk-base/build.gradle.kts
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
/*
* Copyright Amazon.com, Inc. or its affiliates.
*
* Licensed under the Apache License, Version 2.0 (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/apache2.0
*
* or in the "license" file accompanying this file. This file 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.
*/

plugins {
`java-library`
}

java {
sourceCompatibility = JavaVersion.VERSION_11
targetCompatibility = JavaVersion.VERSION_11
}

dependencies {
implementation("com.google.code.gson:gson:2.10.1")
implementation("com.sparkjava:spark-core")
implementation("com.fasterxml.jackson.core:jackson-databind:2.13.3")
}
Loading

0 comments on commit 5d211eb

Please sign in to comment.