diff --git a/src/main/java/org/sagebionetworks/template/TemplateGuiceModule.java b/src/main/java/org/sagebionetworks/template/TemplateGuiceModule.java index d081361a..0a74d30a 100644 --- a/src/main/java/org/sagebionetworks/template/TemplateGuiceModule.java +++ b/src/main/java/org/sagebionetworks/template/TemplateGuiceModule.java @@ -11,6 +11,7 @@ import static org.sagebionetworks.template.TemplateUtils.loadFromJsonFile; import java.io.IOException; +import java.time.Duration; import org.apache.http.client.HttpClient; import org.apache.http.impl.client.HttpClientBuilder; diff --git a/src/main/java/org/sagebionetworks/template/repo/agent/BedrockAgentContextProvider.java b/src/main/java/org/sagebionetworks/template/repo/agent/BedrockAgentContextProvider.java index d3d70849..b0766fd7 100644 --- a/src/main/java/org/sagebionetworks/template/repo/agent/BedrockAgentContextProvider.java +++ b/src/main/java/org/sagebionetworks/template/repo/agent/BedrockAgentContextProvider.java @@ -4,6 +4,7 @@ import static org.sagebionetworks.template.Constants.PROPERTY_KEY_STACK; import java.util.StringJoiner; +import java.util.UUID; import org.apache.velocity.VelocityContext; import org.json.JSONArray; @@ -12,16 +13,19 @@ import org.sagebionetworks.template.config.RepoConfiguration; import org.sagebionetworks.template.repo.VelocityContextProvider; +import com.amazonaws.services.s3.AmazonS3Client; import com.google.inject.Inject; public class BedrockAgentContextProvider implements VelocityContextProvider { private final RepoConfiguration repoConfig; + private final AmazonS3Client s3Cient; @Inject - public BedrockAgentContextProvider(RepoConfiguration repoConfig) { + public BedrockAgentContextProvider(RepoConfiguration repoConfig, AmazonS3Client s3Client) { super(); this.repoConfig = repoConfig; + this.s3Cient = s3Client; } @Override @@ -29,6 +33,13 @@ public void addToContext(VelocityContext context) { String stack = repoConfig.getProperty(PROPERTY_KEY_STACK); String instance = repoConfig.getProperty(PROPERTY_KEY_INSTANCE); String agentName = new StringJoiner("-").add(stack).add(instance).add("agent").toString(); + + String openApiSchemaBucket = String.format("%s-configuration.sagebase.org", stack); + String openApiSchemakey = String.format("chat/openapi/%s/%s.json", instance, UUID.randomUUID().toString()); + + String openApiSchemJsonString = TemplateUtils.loadContentFromFile("templates/repo/agent/agent_open_api.json"); + s3Cient.putObject(openApiSchemaBucket, openApiSchemakey, openApiSchemJsonString); + JSONObject baseTemplate = new JSONObject(TemplateUtils.loadContentFromFile("templates/repo/agent/bedrock_agent_template.json")); @@ -60,6 +71,11 @@ public void addToContext(VelocityContext context) { kbProperty.getJSONObject("KnowledgeBaseId").put("Ref", "SynapseHelpKnowledgeBase"); kbProperty.put("Description", baseTemplate.getJSONObject("Parameters").getJSONObject("knowledgeBaseDescription").getString("Default")); + JSONObject s3 = bedrockAgentProps.getJSONArray("ActionGroups").getJSONObject(1).getJSONObject("ApiSchema") + .getJSONObject("S3"); + s3.put("S3BucketName", openApiSchemaBucket); + s3.put("S3ObjectKey", openApiSchemakey); + bedrockAgentProps.put("AgentName", agentName); String json = resources.toString(); context.put("bedrock_agent_resouces", "," + json.substring(1, json.length()-1)); diff --git a/src/main/resources/templates/repo/agent/agent_open_api.json b/src/main/resources/templates/repo/agent/agent_open_api.json new file mode 100644 index 00000000..a23d3daa --- /dev/null +++ b/src/main/resources/templates/repo/agent/agent_open_api.json @@ -0,0 +1,137 @@ +{ + "openapi": "3.0.0", + "info": { + "title": "CRUD Opperations Schema", + "version": "v1" + }, + "paths": { + "/entity/{entityId}/annotations": { + "get": { + "description": "Get the annotations of an entity.", + "operationId": "orgSageOneEntityAnnotationsGet", + "parameters": [ + { + "name": "entityId", + "description": "The 'syn' Id of the entity to get the annotations for.", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "Auto-generated description", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/org.sagebionetworks.repo.model.annotation.v2.Annotations" + } + } + } + } + } + }, + "put": { + "description": "Set the annotations of an entity.", + "operationId": "orgSageOneEntityAnnotationsPut", + "x-requireConfirmation": "ENABLED", + "parameters": [ + { + "name": "entityId", + "description": "The 'syn' of the target Entity.", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/org.sagebionetworks.repo.model.annotation.v2.Annotations" + } + } + }, + "required": true + }, + "responses": { + "201": { + "description": "the updated annotations", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/org.sagebionetworks.repo.model.annotation.v2.Annotations" + } + } + } + } + } + } + } + }, + "components": { + "schemas": { + "org.sagebionetworks.repo.model.annotation.v2.Annotations": { + "type": "object", + "description": "Each annotation is a key-value pair of metadata for an Entity.", + "properties": { + "id": { + "type": "string", + "description": "The 'syn' ID of the Entity" + }, + "etag": { + "type": "string", + "description": "The current etag of the Entity. The etag of an Entity will change each time an Entity is updated. An annotations update request will be rejected if provided etag does not match the current etag of the Entity." + }, + "annotations": { + "type": "object", + "description": "Annotations are a map of key-value pairs. The key is the name of the annotation. The value is an object.", + "additionalProperties": { + "$ref": "#/components/schemas/org.sagebionetworks.repo.model.annotation.v2.AnnotationsValue" + } + } + }, + "required": [ + "id", + "etag", + "annotations" + ] + }, + "org.sagebionetworks.repo.model.annotation.v2.AnnotationsValue": { + "type": "object", + "description": "This object defines the value of an annotation.", + "properties": { + "type": { + "$ref": "#/components/schemas/org.sagebionetworks.repo.model.annotation.v2.AnnotationsValueType" + }, + "value": { + "type": "array", + "description": "The annotation's value must include one or more strings. See the 'type' to determine how each string should be parsed. An empty array is not allowed.", + "items": { + "type": "string" + } + } + }, + "required": [ + "type", + "value" + ] + }, + "org.sagebionetworks.repo.model.annotation.v2.AnnotationsValueType": { + "type": "string", + "description": "The annotation's type defines how each string value should be parsed.", + "enum": [ + "STRING", + "DOUBLE", + "LONG", + "TIMESTAMP_MS", + "BOOLEAN" + ] + } + } + } +} diff --git a/src/main/resources/templates/repo/agent/bedrock_agent_template.json b/src/main/resources/templates/repo/agent/bedrock_agent_template.json index 675ea471..420e69e8 100644 --- a/src/main/resources/templates/repo/agent/bedrock_agent_template.json +++ b/src/main/resources/templates/repo/agent/bedrock_agent_template.json @@ -18,11 +18,32 @@ "Type": "String", "Default": "This knowledge base contains the Synapse help documentation. You can use it to answer questions around Synapse usage and its features.", "MaxLength": 200 + }, + "openApiSchemaS3Bucket": { + "Description": "The name of the S3 bucket containing the Open API Schema definition JSON file.", + "Type": "String", + "AllowedPattern": "^([0-9a-zA-Z][_-]?){1,100}$" + }, + "openApiSchemaS3Key": { + "Description": "The S3 key of the Open API Schema definition JSON file.", + "Type": "String", + "AllowedPattern": "^([0-9a-zA-Z][_-]?){1,100}$" } }, "Conditions": { - "AttachKnowledgeBase": {"Fn::Not": [{"Fn::Equals" : [{"Ref" : "knowledgeBaseId"}, ""]}]} - }, + "AttachKnowledgeBase": { + "Fn::Not": [ + { + "Fn::Equals": [ + { + "Ref": "knowledgeBaseId" + }, + "" + ] + } + ] + } + }, "Resources": { "bedrockAgentRole": { "Type": "AWS::IAM::Role", @@ -59,7 +80,7 @@ "Statement": [ { "Effect": "Allow", - "Action": [ + "Action": [ "bedrock:InvokeModel", "bedrock:InvokeModelWithResponseStream" ], @@ -70,23 +91,29 @@ ] }, { - "Fn::If" : [ "AttachKnowledgeBase", + "Fn::If": [ + "AttachKnowledgeBase", { - "Effect": "Allow", - "Action": [ - "bedrock:Retrieve", - "bedrock:RetrieveAndGenerate" - ], - "Resource": [ + "Effect": "Allow", + "Action": [ + "bedrock:Retrieve", + "bedrock:RetrieveAndGenerate" + ], + "Resource": [ { "Fn::Sub": "arn:aws:bedrock:${AWS::Region}:${AWS::AccountId}:knowledge-base/${knowledgeBaseId}" } - ] - }, + ] + }, { - "Ref" : "AWS::NoValue" + "Ref": "AWS::NoValue" } ] + }, + { + "Effect": "Allow", + "Action": "s3:*", + "Resource": "*" } ] } @@ -164,9 +191,29 @@ ] }, "SkipResourceInUseCheckOnDelete": false + }, + { + "ActionGroupExecutor": { + "CustomControl": "RETURN_CONTROL" + }, + "ActionGroupName": "org_sage_one", + "ActionGroupState": "ENABLED", + "Description": "These functions can be used to make data changes in Synapse.", + "ApiSchema": { + "S3": { + "S3BucketName": { + "Ref": "openApiSchemaS3Bucket" + }, + "S3ObjectKey": { + "Ref": "openApiSchemaS3Key" + } + } + } } ], - "AgentName": { "Ref" : "agentName" }, + "AgentName": { + "Ref": "agentName" + }, "AgentResourceRoleArn": { "Fn::GetAtt": [ "bedrockAgentRole", @@ -174,15 +221,22 @@ ] }, "KnowledgeBases": { - "Fn::If" : [ "AttachKnowledgeBase", - [{ - "KnowledgeBaseId" : { "Ref" : "knowledgeBaseId" }, - "Description" : { "Ref" : "knowledgeBaseDescription" } - }], + "Fn::If": [ + "AttachKnowledgeBase", + [ + { + "KnowledgeBaseId": { + "Ref": "knowledgeBaseId" + }, + "Description": { + "Ref": "knowledgeBaseDescription" + } + } + ], { - "Ref" : "AWS::NoValue" - } - ] + "Ref": "AWS::NoValue" + } + ] }, "AutoPrepare": true, "Description": "Test of the use of actions groups to allow the agent to make Synapse API calls.", diff --git a/src/test/java/org/sagebionetworks/template/repo/RepositoryTemplateBuilderImplTest.java b/src/test/java/org/sagebionetworks/template/repo/RepositoryTemplateBuilderImplTest.java index cc98e525..8a97bdad 100644 --- a/src/test/java/org/sagebionetworks/template/repo/RepositoryTemplateBuilderImplTest.java +++ b/src/test/java/org/sagebionetworks/template/repo/RepositoryTemplateBuilderImplTest.java @@ -8,6 +8,7 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.never; @@ -127,6 +128,7 @@ import com.amazonaws.services.elasticbeanstalk.model.ListPlatformVersionsResult; import com.amazonaws.services.elasticbeanstalk.model.PlatformFilter; import com.amazonaws.services.elasticbeanstalk.model.PlatformSummary; +import com.amazonaws.services.s3.AmazonS3Client; import com.amazonaws.services.securitytoken.AWSSecurityTokenService; import com.amazonaws.services.securitytoken.model.GetCallerIdentityResult; import com.google.common.collect.Lists; @@ -167,8 +169,14 @@ public class RepositoryTemplateBuilderImplTest { private AWSSecurityTokenService mockStsClient; @Mock private WaitConditionHandler mockWaitConditionHandler; + @Mock + private AmazonS3Client mockS3Client; + @Captor private ArgumentCaptor requestCaptor; + + @Captor + private ArgumentCaptor jsonStringCaptor; private VelocityEngine velocityEngine; private RepositoryTemplateBuilderImpl builder; @@ -201,7 +209,7 @@ public void before() throws InterruptedException { when(mockLoggerFactory.getLogger(any())).thenReturn(mockLogger); builder = new RepositoryTemplateBuilderImpl(mockCloudFormationClient, velocityEngine, config, mockLoggerFactory, - mockArtifactCopy, mockSecretBuilder, Sets.newHashSet(mockContextProvider1, mockContextProvider2, new BedrockAgentContextProvider(config)), + mockArtifactCopy, mockSecretBuilder, Sets.newHashSet(mockContextProvider1, mockContextProvider2, new BedrockAgentContextProvider(config, mockS3Client)), mockElasticBeanstalkSolutionStackNameProvider, mockStackTagsProvider, mockCwlContextProvider, mockEc2Client, mockBeanstalkClient, mockTimeToLive, mockStsClient, Set.of(mockWaitConditionHandler)); @@ -386,7 +394,27 @@ public void testBuildAndDeployProd() throws InterruptedException { assertTrue(resources.has("SynapseHelpKnowledgeBase")); assertTrue(resources.has("bedrockAgentRole")); assertTrue(resources.has("bedrockAgent")); - assertEquals("prod-101-agent", resources.getJSONObject("bedrockAgent").getJSONObject("Properties").get("AgentName")); + JSONObject bedrockAgentProps = resources.getJSONObject("bedrockAgent").getJSONObject("Properties"); + + assertEquals("prod-101-agent", bedrockAgentProps.get("AgentName")); + + validateOpenApiSchema(bedrockAgentProps); + + } + + void validateOpenApiSchema(JSONObject bedrockAgentProps) { + JSONObject s3 = bedrockAgentProps.getJSONArray("ActionGroups").getJSONObject(1).getJSONObject("ApiSchema") + .getJSONObject("S3"); + String openApiBucket = s3.getString("S3BucketName"); + assertEquals("prod-configuration.sagebase.org", openApiBucket); + String openApiKey = s3.getString("S3ObjectKey"); + assertTrue(openApiKey.startsWith("chat/openapi/101/")); + verify(mockS3Client).putObject(eq(openApiBucket), eq(openApiKey), jsonStringCaptor.capture()); + + JSONObject openApiSchema = new JSONObject(jsonStringCaptor.getValue()); + assertTrue(openApiSchema.has("openapi")); + assertTrue(openApiSchema.has("info")); + assertTrue(openApiSchema.has("paths")); } @Test