diff --git a/packages/blueprints/gen-ai-chatbot/README.md b/packages/blueprints/gen-ai-chatbot/README.md index e7337d4df..4ecbacf97 100644 --- a/packages/blueprints/gen-ai-chatbot/README.md +++ b/packages/blueprints/gen-ai-chatbot/README.md @@ -40,6 +40,9 @@ The following languages are supported for a custom chatbot: * Japanese (日本語) * Korean (한국어) * Chinese (中文) +* Français +* Deutsch +* Español ## Deployment After building your chatbot, you can also deploy it with this blueprint. Before a chatbot can be deployed with a CodeCatalyst workflow, you must enable model access. diff --git a/packages/blueprints/gen-ai-chatbot/src/blueprint.ts b/packages/blueprints/gen-ai-chatbot/src/blueprint.ts index 3b3a56a30..3d12c1943 100644 --- a/packages/blueprints/gen-ai-chatbot/src/blueprint.ts +++ b/packages/blueprints/gen-ai-chatbot/src/blueprint.ts @@ -179,6 +179,13 @@ export interface Options extends ParentOptions { * @defaultEntropy 8 */ stackDisambiguator?: string; + + /** + * This project is focused on Antrophic Claude models; limit support is provided for Mistral models. Enabling this option + * means that only Mistral models will be used for *all* the chat features. + * @displayName Enable Mistral + */ + enableMistral?: boolean; }; } @@ -221,6 +228,7 @@ export class Blueprint extends ParentBlueprint { bucketNamePrefix: options.code.bucketNamePrefix, enableSelfRegistration: options.enableSelfRegistration === 'Enabled', stackDisambiguator: options.code.stackDisambiguator, + enableMistral: options.code.enableMistral ?? false, }), ); diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/cdk.json b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/cdk.json index 1f2614dff..8939955c6 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/cdk.json +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/cdk.json @@ -6,6 +6,8 @@ ], "exclude": [ "README.md", + "./backend/**/__pycache__", + "./backend/tests/**", "cdk*.json", "**/*.d.ts", "**/*.js", @@ -29,6 +31,7 @@ "allowedIpV4AddressRanges": [{{#allowedIpV4AddressRanges}}"{{{value}}}"{{#comma}},{{/comma}}{{/allowedIpV4AddressRanges}}], "allowedIpV6AddressRanges": [{{#allowedIpV6AddressRanges}}"{{{value}}}"{{#comma}},{{/comma}}{{/allowedIpV6AddressRanges}}], "enableSelfRegistration": {{enableSelfRegistration}}, + "enableMistral": {{enableMistral}}, "publishedApiAllowedIpV4AddressRanges": ["0.0.0.0/1", "128.0.0.0/1"], "publishedApiAllowedIpV6AddressRanges": [ "0000:0000:0000:0000:0000:0000:0000:0000/1", diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/chatbot-genai-cdk-stack.ts b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/chatbot-genai-cdk-stack.ts index fdfcf36f8..a8eff8d46 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/chatbot-genai-cdk-stack.ts +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/chatbot-genai-cdk-stack.ts @@ -34,6 +34,7 @@ export interface ChatbotGenAiCdkStackProps extends StackProps { readonly publishedApiAllowedIpV6AddressRanges: string[]; readonly allowedSignUpEmailDomains: string[]; readonly rdsSchedules: CronScheduleProps; + readonly enableMistral: boolean; } export class ChatbotGenAiCdkStack extends cdk.Stack { @@ -93,6 +94,7 @@ export class ChatbotGenAiCdkStack extends cdk.Stack { prefix: `${this.node.tryGetContext('bucketNamePrefix') ?? 'chatbot-frontend-assets'}-${this.account}-${this.region}`, removalPolicy: this.node.tryGetContext('bucketRemovalPolicy') ?? 'DESTROY', }, + enableMistral: props.enableMistral, }); const auth = new Auth(this, "Auth", { @@ -118,21 +120,23 @@ export class ChatbotGenAiCdkStack extends cdk.Stack { objectOwnership: ObjectOwnership.OBJECT_WRITER, autoDeleteObjects: true, versioned: true, - serverAccessLogsBucket: new Bucket(this, 'DdbBucketLogs', { - encryption: BucketEncryption.S3_MANAGED, - blockPublicAccess: BlockPublicAccess.BLOCK_ALL, - enforceSSL: true, - removalPolicy: RemovalPolicy.DESTROY, - lifecycleRules: [ - { - enabled: true, - expiration: Duration.days(3653), - id: 'ExpireAfterTenYears', - }, - ], - versioned: true, - serverAccessLogsPrefix: 'self/', - }), + serverAccessLogsBucket: CF_SKIP_ACCESS_LOGGING_REGIONS.includes(this.region) + ? undefined + : new Bucket(this, "DdbBucketLogs", { + encryption: BucketEncryption.S3_MANAGED, + blockPublicAccess: BlockPublicAccess.BLOCK_ALL, + enforceSSL: true, + removalPolicy: RemovalPolicy.DESTROY, + lifecycleRules: [ + { + enabled: true, + expiration: Duration.days(3653), + id: "ExpireAfterTenYears", + }, + ], + versioned: true, + serverAccessLogsPrefix: "self/", + }), }); const database = new Database(this, "Database", { @@ -174,13 +178,14 @@ export class ChatbotGenAiCdkStack extends cdk.Stack { backendApiEndpoint: backendApi.api.apiEndpoint, webSocketApiEndpoint: websocket.apiEndpoint, userPoolDomainPrefix: props.userPoolDomainPrefix, + enableMistral: props.enableMistral, auth, idp, }); documentBucket.addCorsRule({ allowedMethods: [HttpMethods.PUT], - allowedOrigins: [frontend.getOrigin(), "http://localhost:5173"], + allowedOrigins: [frontend.getOrigin(), "http://localhost:5173", "*"], allowedHeaders: ["*"], maxAge: 3000, }); @@ -228,35 +233,51 @@ export class ChatbotGenAiCdkStack extends cdk.Stack { }); new CfnOutput(this, "AvailabilityZone0", { value: vpc.availabilityZones[0], - exportName: this.disambiguateIdentifier("BedrockClaudeChatAvailabilityZone0"), + exportName: this.disambiguateIdentifier( + "BedrockClaudeChatAvailabilityZone0" + ), }); new CfnOutput(this, "AvailabilityZone1", { value: vpc.availabilityZones[1], - exportName: this.disambiguateIdentifier("BedrockClaudeChatAvailabilityZone1"), + exportName: this.disambiguateIdentifier( + "BedrockClaudeChatAvailabilityZone1" + ), }); new CfnOutput(this, "PublicSubnetId0", { value: vpc.publicSubnets[0].subnetId, - exportName: this.disambiguateIdentifier("BedrockClaudeChatPublicSubnetId0"), + exportName: this.disambiguateIdentifier( + "BedrockClaudeChatPublicSubnetId0" + ), }); new CfnOutput(this, "PublicSubnetId1", { value: vpc.publicSubnets[1].subnetId, - exportName: this.disambiguateIdentifier("BedrockClaudeChatPublicSubnetId1"), + exportName: this.disambiguateIdentifier( + "BedrockClaudeChatPublicSubnetId1" + ), }); new CfnOutput(this, "PrivateSubnetId0", { value: vpc.privateSubnets[0].subnetId, - exportName: this.disambiguateIdentifier("BedrockClaudeChatPrivateSubnetId0"), + exportName: this.disambiguateIdentifier( + "BedrockClaudeChatPrivateSubnetId0" + ), }); new CfnOutput(this, "PrivateSubnetId1", { value: vpc.privateSubnets[1].subnetId, - exportName: this.disambiguateIdentifier("BedrockClaudeChatPrivateSubnetId1"), + exportName: this.disambiguateIdentifier( + "BedrockClaudeChatPrivateSubnetId1" + ), }); new CfnOutput(this, "DbConfigSecretArn", { value: vectorStore.secret.secretArn, - exportName: this.disambiguateIdentifier("BedrockClaudeChatDbConfigSecretArn"), + exportName: this.disambiguateIdentifier( + "BedrockClaudeChatDbConfigSecretArn" + ), }); new CfnOutput(this, "DbConfigHostname", { value: vectorStore.cluster.clusterEndpoint.hostname, - exportName: this.disambiguateIdentifier("BedrockClaudeChatDbConfigHostname"), + exportName: this.disambiguateIdentifier( + "BedrockClaudeChatDbConfigHostname" + ), }); new CfnOutput(this, "DbConfigPort", { value: vectorStore.cluster.clusterEndpoint.port.toString(), @@ -264,24 +285,32 @@ export class ChatbotGenAiCdkStack extends cdk.Stack { }); new CfnOutput(this, "ConversationTableName", { value: database.table.tableName, - exportName: this.disambiguateIdentifier("BedrockClaudeChatConversationTableName"), + exportName: this.disambiguateIdentifier( + "BedrockClaudeChatConversationTableName" + ), }); new CfnOutput(this, "TableAccessRoleArn", { value: database.tableAccessRole.roleArn, - exportName: this.disambiguateIdentifier("BedrockClaudeChatTableAccessRoleArn"), + exportName: this.disambiguateIdentifier( + "BedrockClaudeChatTableAccessRoleArn" + ), }); new CfnOutput(this, "DbSecurityGroupId", { value: vectorStore.securityGroup.securityGroupId, - exportName: this.disambiguateIdentifier("BedrockClaudeChatDbSecurityGroupId"), + exportName: this.disambiguateIdentifier( + "BedrockClaudeChatDbSecurityGroupId" + ), }); new CfnOutput(this, "LargeMessageBucketName", { value: largeMessageBucket.bucketName, - exportName: this.disambiguateIdentifier("BedrockClaudeChatLargeMessageBucketName"), + exportName: this.disambiguateIdentifier( + "BedrockClaudeChatLargeMessageBucketName" + ), }); } private disambiguateIdentifier(identifier: string) { - const disambiguator = this.node.tryGetContext('stackDisambiguator'); + const disambiguator = this.node.tryGetContext("stackDisambiguator"); return disambiguator ? `${identifier}-${disambiguator}` : identifier; } } diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/constructs/auth.ts b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/constructs/auth.ts index 40b2d217f..b6c7713e3 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/constructs/auth.ts +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/constructs/auth.ts @@ -99,6 +99,7 @@ export class Auth extends Construct { } ); client.node.addDependency(googleProvider); + break; } case "oidc": { const issuerUrl = secret @@ -124,6 +125,7 @@ export class Auth extends Construct { } ); client.node.addDependency(oidcProvider); + break; } } }; @@ -141,10 +143,16 @@ export class Auth extends Construct { } if (props.allowedSignUpEmailDomains.length >= 1) { - const checkEmailDomainFunction = new PythonFunction(this, 'CheckEmailDomain',{ + const checkEmailDomainFunction = new PythonFunction( + this, + "CheckEmailDomain", + { runtime: Runtime.PYTHON_3_12, - index: 'check_email_domain.py', - entry: path.join(__dirname, "../../backend/auth/check_email_domain"), + index: "check_email_domain.py", + entry: path.join( + __dirname, + "../../backend/auth/check_email_domain" + ), timeout: Duration.minutes(1), environment: { ALLOWED_SIGN_UP_EMAIL_DOMAINS_STR: JSON.stringify( diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/constructs/frontend.ts b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/constructs/frontend.ts index b5a1c29bc..7557ed6bf 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/constructs/frontend.ts +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/constructs/frontend.ts @@ -21,6 +21,7 @@ export interface FrontendProps { prefix: string; removalPolicy: string; }; + enableMistral: boolean; } export class Frontend extends Construct { @@ -34,7 +35,10 @@ export class Frontend extends Construct { encryption: BucketEncryption.S3_MANAGED, blockPublicAccess: BlockPublicAccess.BLOCK_ALL, enforceSSL: true, - removalPolicy: props.assetBucket.removalPolicy === 'RETAIN' ? RemovalPolicy.RETAIN : RemovalPolicy.DESTROY, + removalPolicy: + props.assetBucket.removalPolicy === "RETAIN" + ? RemovalPolicy.RETAIN + : RemovalPolicy.DESTROY, autoDeleteObjects: true, }); @@ -88,12 +92,14 @@ export class Frontend extends Construct { backendApiEndpoint, webSocketApiEndpoint, userPoolDomainPrefix, + enableMistral, auth, idp, }: { backendApiEndpoint: string; webSocketApiEndpoint: string; userPoolDomainPrefix: string; + enableMistral: boolean; auth: Auth; idp: Idp; }) { @@ -106,6 +112,7 @@ export class Frontend extends Construct { VITE_APP_WS_ENDPOINT: webSocketApiEndpoint, VITE_APP_USER_POOL_ID: auth.userPool.userPoolId, VITE_APP_USER_POOL_CLIENT_ID: auth.client.userPoolClientId, + VITE_APP_ENABLE_MISTRAL: enableMistral.toString(), VITE_APP_REGION: region, VITE_APP_USE_STREAMING: "true", }; diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/constructs/usage-analysis.ts b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/constructs/usage-analysis.ts index 9b299da62..ced8f7427 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/constructs/usage-analysis.ts +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/lib/constructs/usage-analysis.ts @@ -11,6 +11,7 @@ import { Runtime } from "aws-cdk-lib/aws-lambda"; import { aws_glue } from "aws-cdk-lib"; import { Database } from "./database"; import * as iam from "aws-cdk-lib/aws-iam"; +import { CF_SKIP_ACCESS_LOGGING_REGIONS } from "../utils/constants"; export interface UsageAnalysisProps { sourceDatabase: Database; @@ -31,7 +32,7 @@ export class UsageAnalysis extends Construct { ).stackName.toLowerCase()}_usage_analysis`; const DDB_EXPORT_TABLE_NAME = "ddb_export"; - // Bucket to export DynamoDB data + // Bucket to export DynamoDB data const ddbBucket = new s3.Bucket(this, "DdbBucket", { encryption: s3.BucketEncryption.S3_MANAGED, blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL, @@ -40,21 +41,23 @@ export class UsageAnalysis extends Construct { objectOwnership: s3.ObjectOwnership.OBJECT_WRITER, autoDeleteObjects: true, versioned: true, - serverAccessLogsBucket: new s3.Bucket(this, 'DdbBucketLogs', { - encryption: s3.BucketEncryption.S3_MANAGED, - blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL, - enforceSSL: true, - removalPolicy: RemovalPolicy.DESTROY, - lifecycleRules: [ - { - enabled: true, - expiration: Duration.days(3653), - id: 'ExpireAfterTenYears', - }, - ], - versioned: true, - serverAccessLogsPrefix: 'self/', - }), + serverAccessLogsBucket: CF_SKIP_ACCESS_LOGGING_REGIONS.includes(this.region) + ? undefined + : new s3.Bucket(this, "DdbBucketLogs", { + encryption: s3.BucketEncryption.S3_MANAGED, + blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL, + enforceSSL: true, + removalPolicy: RemovalPolicy.DESTROY, + lifecycleRules: [ + { + enabled: true, + expiration: Duration.days(3653), + id: "ExpireAfterTenYears", + }, + ], + versioned: true, + serverAccessLogsPrefix: "self/", + }), }); // Bucket for Athena query results @@ -94,7 +97,11 @@ export class UsageAnalysis extends Construct { }, { name: "MessageMap", - type: glue.Schema.STRING, + type: glue.Schema.struct([{ name: "S", type: glue.Schema.STRING }]), + }, + { + name: "IsLargeMessage", + type: glue.Schema.struct([{ name: "BOOL", type: glue.Schema.BOOLEAN }]), }, { name: "PK", diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/test/cdk.test.ts b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/test/cdk.test.ts index ef20d9aaa..6a00983fa 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/test/cdk.test.ts +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-cdk/test/cdk.test.ts @@ -28,6 +28,7 @@ describe("Fine-grained Assertions Test", () => { stop: {}, start: {}, }, + enableMistral: false, } ); const hasGoogleProviderTemplate = Template.fromStack( @@ -80,6 +81,7 @@ describe("Fine-grained Assertions Test", () => { stop: {}, start: {}, }, + enableMistral: false, } ); const hasOidcProviderTemplate = Template.fromStack(hasOidcProviderStack); @@ -121,10 +123,17 @@ describe("Fine-grained Assertions Test", () => { stop: {}, start: {}, }, + enableMistral: false, }); const template = Template.fromStack(stack); template.resourceCountIs("AWS::Cognito::UserPoolIdentityProvider", 0); + // verify the stack has environment variable VITE_APP_ENABLE_MISTRAL is set to "false" + template.hasResourceProperties("Custom::CDKNodejsBuild", { + environment: { + VITE_APP_ENABLE_MISTRAL: "false", + }, + }); }); }); @@ -156,6 +165,7 @@ describe("Scheduler Test", () => { year: "*", }, }, + enableMistral: false, }); const template = Template.fromStack(hasScheduleStack); template.hasResourceProperties("AWS::Scheduler::Schedule", { @@ -181,6 +191,7 @@ describe("Scheduler Test", () => { stop: {}, start: {}, }, + enableMistral: false, }); const template = Template.fromStack(defaultStack); // The stack should have only 1 rule for exporting the data from ddb to s3 diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/README.md b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/README.md index 43b5f088f..52be3a2fd 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/README.md +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/README.md @@ -23,6 +23,8 @@ export REGION=ap-northeast-1 export BEDROCK_REGION=us-east-1 export DOCUMENT_BUCKET=bedrockchatstack-documentbucketxxxxxxx export LARGE_MESSAGE_BUCKET=bedrockchatstack-largemessagebucketxxx +export USER_POOL_ID=xxxxxxxxx +export CLIENT_ID=xxxxxxxxx ``` ## Launch local server diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/bedrock.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/bedrock.py index 8c9c0938c..6064dcb5b 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/bedrock.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/bedrock.py @@ -3,9 +3,15 @@ import os from anthropic import AnthropicBedrock -from app.config import ANTHROPIC_PRICING, DEFAULT_EMBEDDING_CONFIG, GENERATION_CONFIG +from app.config import ( + BEDROCK_PRICING, + DEFAULT_EMBEDDING_CONFIG, + GENERATION_CONFIG, + MISTRAL_GENERATION_CONFIG, +) from app.repositories.models.conversation import MessageModel -from app.utils import get_bedrock_client +from app.utils import get_bedrock_client, is_anthropic_model +from pydantic import BaseModel logger = logging.getLogger(__name__) @@ -16,6 +22,58 @@ anthropic_client = AnthropicBedrock() +class InvocationMetrics(BaseModel): + input_tokens: int + output_tokens: int + + +def compose_args( + messages: list[MessageModel], + model: str, + instruction: str | None = None, + stream: bool = False, +) -> dict: + # if model is from Anthropic, use AnthropicBedrock + # otherwise, use bedrock client + model_id = get_model_id(model) + if is_anthropic_model(model_id): + return compose_args_for_anthropic_client(messages, model, instruction, stream) + else: + return compose_args_for_other_client(messages, model, instruction, stream) + + +def compose_args_for_other_client( + messages: list[MessageModel], + model: str, + instruction: str | None = None, + stream: bool = False, +) -> dict: + arg_messages = [] + for message in messages: + if message.role not in ["system", "instruction"]: + content: list[dict] = [] + for c in message.content: + if c.content_type == "text": + content.append( + { + "type": "text", + "text": c.body, + } + ) + m = {"role": message.role, "content": content} + arg_messages.append(m) + + args = { + **MISTRAL_GENERATION_CONFIG, + "model": get_model_id(model), + "messages": arg_messages, + "stream": stream, + } + if instruction: + args["system"] = instruction + return args + + def compose_args_for_anthropic_client( messages: list[MessageModel], model: str, @@ -66,14 +124,14 @@ def calculate_price( model: str, input_tokens: int, output_tokens: int, region: str = BEDROCK_REGION ) -> float: input_price = ( - ANTHROPIC_PRICING.get(region, {}) + BEDROCK_PRICING.get(region, {}) .get(model, {}) - .get("input", ANTHROPIC_PRICING["default"][model]["input"]) + .get("input", BEDROCK_PRICING["default"][model]["input"]) ) output_price = ( - ANTHROPIC_PRICING.get(region, {}) + BEDROCK_PRICING.get(region, {}) .get(model, {}) - .get("output", ANTHROPIC_PRICING["default"][model]["output"]) + .get("output", BEDROCK_PRICING["default"][model]["output"]) ) return input_price * input_tokens / 1000.0 + output_price * output_tokens / 1000.0 @@ -91,6 +149,12 @@ def get_model_id(model: str) -> str: return "anthropic.claude-3-haiku-20240307-v1:0" elif model == "claude-v3-opus": return "anthropic.claude-3-opus-20240229-v1:0" + elif model == "mistral-7b-instruct": + return "mistral.mistral-7b-instruct-v0:2" + elif model == "mixtral-8x7b-instruct": + return "mistral.mixtral-8x7b-instruct-v0:1" + elif model == "mistral-large": + return "mistral.mistral-large-2402-v1:0" else: raise NotImplementedError() @@ -141,3 +205,62 @@ def _calculate_document_embeddings(documents: list[str]) -> list[list[float]]: embeddings += _calculate_document_embeddings(batch) return embeddings + + +def get_bedrock_response(args: dict) -> dict: + + client = get_bedrock_client() + messages = args["messages"] + + prompt = "\n".join( + [ + message["content"][0]["text"] + for message in messages + if message["content"][0]["type"] == "text" + ] + ) + + model_id = args["model"] + is_mistral_model = model_id.startswith("mistral") + if is_mistral_model: + prompt = f"[INST] {prompt} [/INST]" + + logger.info(f"Final Prompt: {prompt}") + body = json.dumps( + { + "prompt": prompt, + "max_tokens": args["max_tokens"], + "temperature": args["temperature"], + "top_p": args["top_p"], + "top_k": args["top_k"], + } + ) + + logger.info(f"The args before invoke bedrock: {args}") + if args["stream"]: + try: + response = client.invoke_model_with_response_stream( + modelId=model_id, + body=body, + ) + # Ref: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/invoke_model_with_response_stream.html + response_body = response + except Exception as e: + logger.error(e) + else: + response = client.invoke_model( + modelId=model_id, + body=body, + ) + # Ref: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/invoke_model.html + response_body = json.loads(response.get("body").read()) + invocation_metrics = InvocationMetrics( + input_tokens=response["ResponseMetadata"]["HTTPHeaders"][ + "x-amzn-bedrock-input-token-count" + ], + output_tokens=response["ResponseMetadata"]["HTTPHeaders"][ + "x-amzn-bedrock-output-token-count" + ], + ) + response_body["amazon-bedrock-invocationMetrics"] = invocation_metrics + return response_body diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/config.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/config.py index 45cf38e97..693bb799a 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/config.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/config.py @@ -26,6 +26,15 @@ class EmbeddingConfig(TypedDict): "stop_sequences": ["Human: ", "Assistant: "], } +# Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mistral.html#model-parameters-mistral-request-response +MISTRAL_GENERATION_CONFIG: GenerationConfig = { + "max_tokens": 4096, + "top_k": 50, + "top_p": 0.9, + "temperature": 0.5, + "stop_sequences": ["[INST]", "[/INST]"], +} + # Configure embedding parameter. DEFAULT_EMBEDDING_CONFIG: EmbeddingConfig = { # DO NOT change `model_id` (currently other models are not supported) @@ -43,7 +52,7 @@ class EmbeddingConfig(TypedDict): # Used for price estimation. # NOTE: The following is based on 2024-03-07 # See: https://aws.amazon.com/bedrock/pricing/ -ANTHROPIC_PRICING = { +BEDROCK_PRICING = { "us-east-1": { "claude-instant-v1": { "input": 0.00080, @@ -55,6 +64,9 @@ class EmbeddingConfig(TypedDict): }, "claude-v3-haiku": {"input": 0.00025, "output": 0.00125}, "claude-v3-sonnet": {"input": 0.00300, "output": 0.01500}, + "mistral-7b-instruct": {"input": 0.00015, "output": 0.0002}, + "mixtral-8x7b-instruct": {"input": 0.00045, "output": 0.0007}, + "mistral-large": {"input": 0.008, "output": 0.024}, }, "us-west-2": { "claude-instant-v1": { @@ -67,6 +79,9 @@ class EmbeddingConfig(TypedDict): }, "claude-v3-sonnet": {"input": 0.00300, "output": 0.01500}, "claude-v3-opus": {"input": 0.01500, "output": 0.07500}, + "mistral-7b-instruct": {"input": 0.00015, "output": 0.0002}, + "mixtral-8x7b-instruct": {"input": 0.00045, "output": 0.0007}, + "mistral-large": {"input": 0.008, "output": 0.024}, }, "ap-northeast-1": { "claude-instant-v1": { @@ -90,5 +105,8 @@ class EmbeddingConfig(TypedDict): "claude-v3-haiku": {"input": 0.00025, "output": 0.00125}, "claude-v3-sonnet": {"input": 0.00300, "output": 0.01500}, "claude-v3-opus": {"input": 0.01500, "output": 0.07500}, + "mistral-7b-instruct": {"input": 0.00015, "output": 0.0002}, + "mixtral-8x7b-instruct": {"input": 0.00045, "output": 0.0007}, + "mistral-large": {"input": 0.008, "output": 0.024}, }, } diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/repositories/common.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/repositories/common.py index 881b4d7f5..1f362755e 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/repositories/common.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/repositories/common.py @@ -18,6 +18,7 @@ class RecordNotFoundError(Exception): class RecordAccessNotAllowedError(Exception): pass + class ResourceConflictError(Exception): pass diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/repositories/conversation.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/repositories/conversation.py index 05c5d74ba..52382d6fe 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/repositories/conversation.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/repositories/conversation.py @@ -14,9 +14,11 @@ decompose_conv_id, ) from app.repositories.models.conversation import ( + ChunkModel, ContentModel, ConversationMeta, ConversationModel, + FeedbackModel, MessageModel, ) from app.utils import get_current_time @@ -204,6 +206,27 @@ def find_conversation_by_id(user_id: str, conversation_id: str) -> ConversationM children=v["children"], parent=v["parent"], create_time=float(v["create_time"]), + feedback=( + FeedbackModel( + thumbs_up=v["feedback"]["thumbs_up"], + category=v["feedback"]["category"], + comment=v["feedback"]["comment"], + ) + if v.get("feedback") + else None + ), + used_chunks=( + [ + ChunkModel( + content=c["content"], + source=c["source"], + rank=c["rank"], + ) + for c in v["used_chunks"] + ] + if v.get("used_chunks") + else None + ), ) for k, v in message_map.items() }, @@ -325,3 +348,28 @@ def change_conversation_title(user_id: str, conversation_id: str, new_title: str logger.info(f"Updated conversation title response: {response}") return response + + +def update_feedback( + user_id: str, conversation_id: str, message_id: str, feedback: FeedbackModel +): + logger.info(f"Updating feedback for conversation: {conversation_id}") + table = _get_table_client(user_id) + conv = find_conversation_by_id(user_id, conversation_id) + message_map = conv.message_map + message_map[message_id].feedback = feedback + + response = table.update_item( + Key={ + "PK": user_id, + "SK": compose_conv_id(user_id, conversation_id), + }, + UpdateExpression="set MessageMap = :m", + ExpressionAttributeValues={ + ":m": json.dumps({k: v.model_dump() for k, v in message_map.items()}) + }, + ConditionExpression="attribute_exists(PK) AND attribute_exists(SK)", + ReturnValues="UPDATED_NEW", + ) + logger.info(f"Updated feedback response: {response}") + return response diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/repositories/models/conversation.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/repositories/models/conversation.py index bbe30b082..656dd58e5 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/repositories/models/conversation.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/repositories/models/conversation.py @@ -10,6 +10,18 @@ class ContentModel(BaseModel): body: str +class FeedbackModel(BaseModel): + thumbs_up: bool + category: str + comment: str + + +class ChunkModel(BaseModel): + content: str + source: str + rank: int + + class MessageModel(BaseModel): role: str content: list[ContentModel] @@ -17,6 +29,8 @@ class MessageModel(BaseModel): children: list[str] parent: str | None create_time: float + feedback: FeedbackModel | None + used_chunks: list[ChunkModel] | None class ConversationModel(BaseModel): diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/bot.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/bot.py index 5d2383b0a..af72d792c 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/bot.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/bot.py @@ -27,7 +27,6 @@ remove_bot_by_id, remove_uploaded_file, ) -from app.usecases.chat import chat, fetch_conversation, propose_conversation_title from app.user import User from fastapi import APIRouter, Request diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/conversation.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/conversation.py index 1d1fb7517..18500b68e 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/conversation.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/conversation.py @@ -3,12 +3,16 @@ delete_conversation_by_id, delete_conversation_by_user_id, find_conversation_by_user_id, + update_feedback, ) +from app.repositories.models.conversation import FeedbackModel from app.routes.schemas.conversation import ( ChatInput, ChatOutput, Conversation, ConversationMetaOutput, + FeedbackInput, + FeedbackOutput, NewTitleInput, ProposedTitle, RelatedDocumentsOutput, @@ -119,3 +123,33 @@ def get_proposed_title(request: Request, conversation_id: str): title = propose_conversation_title(current_user.id, conversation_id) return ProposedTitle(title=title) + + +@router.put( + "/conversation/{conversation_id}/{message_id}/feedback", + response_model=FeedbackOutput, +) +def put_feedback( + request: Request, + conversation_id: str, + message_id: str, + feedback_input: FeedbackInput, +): + """Send feedback.""" + current_user: User = request.state.current_user + + update_feedback( + user_id=current_user.id, + conversation_id=conversation_id, + message_id=message_id, + feedback=FeedbackModel( + thumbs_up=feedback_input.thumbs_up, + category=feedback_input.category if feedback_input.category else "", + comment=feedback_input.comment if feedback_input.comment else "", + ), + ) + return FeedbackOutput( + thumbs_up=feedback_input.thumbs_up, + category=feedback_input.category if feedback_input.category else "", + comment=feedback_input.comment if feedback_input.comment else "", + ) diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/schemas/bot.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/schemas/bot.py index 7514c385f..45c62d2ff 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/schemas/bot.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/schemas/bot.py @@ -1,8 +1,11 @@ -from typing import Literal - +from __future__ import annotations +from typing import Literal, TYPE_CHECKING from app.routes.schemas.base import BaseSchema from pydantic import Field +if TYPE_CHECKING: + from app.repositories.models.custom_bot import BotModel + # Knowledge sync status type # NOTE: `ORIGINAL_NOT_FOUND` is used when the original bot is removed. type_sync_status = Literal[ @@ -46,6 +49,42 @@ class BotModifyInput(BaseSchema): embedding_params: EmbeddingParams | None knowledge: KnowledgeDiffInput | None + def has_update_files(self) -> bool: + return self.knowledge is not None and ( + len(self.knowledge.added_filenames) > 0 + or len(self.knowledge.deleted_filenames) > 0 + ) + + def is_embedding_required(self, current_bot_model: BotModel) -> bool: + if self.has_update_files(): + return True + + if self.knowledge is not None and current_bot_model.has_knowledge(): + if set(self.knowledge.source_urls) == set( + current_bot_model.knowledge.source_urls + ) and set(self.knowledge.sitemap_urls) == set( + current_bot_model.knowledge.sitemap_urls + ): + pass + else: + return True + + if ( + self.embedding_params is not None + and current_bot_model.embedding_params is not None + ): + if ( + self.embedding_params.chunk_size + == current_bot_model.embedding_params.chunk_size + and self.embedding_params.chunk_overlap + == current_bot_model.embedding_params.chunk_overlap + ): + pass + else: + return True + + return False + class BotModifyOutput(BaseSchema): id: str diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/schemas/conversation.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/schemas/conversation.py index 790e7d96b..6198af411 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/schemas/conversation.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/routes/schemas/conversation.py @@ -1,7 +1,7 @@ from typing import Literal from app.routes.schemas.base import BaseSchema -from pydantic import Field +from pydantic import Field, root_validator, validator type_model_name = Literal[ "claude-instant-v1", @@ -9,6 +9,9 @@ "claude-v3-sonnet", "claude-v3-haiku", "claude-v3-opus", + "mistral-7b-instruct", + "mixtral-8x7b-instruct", + "mistral-large", ] @@ -23,6 +26,36 @@ class Content(BaseSchema): body: str = Field(..., description="Content body. Text or base64 encoded image.") +class FeedbackInput(BaseSchema): + thumbs_up: bool + category: str | None = Field( + None, description="Reason category. Required if thumbs_up is False." + ) + comment: str | None = Field(None, description="optional comment") + + @root_validator(pre=True) + def check_category(cls, values): + thumbs_up = values.get("thumbs_up") + category = values.get("category") + + if not thumbs_up and category is None: + raise ValueError("category is required if `thumbs_up` is `False`") + + return values + + +class FeedbackOutput(BaseSchema): + thumbs_up: bool + category: str + comment: str + + +class Chunk(BaseSchema): + content: str + source: str + rank: int + + class MessageInput(BaseSchema): role: str content: list[Content] @@ -38,6 +71,8 @@ class MessageOutput(BaseSchema): content: list[Content] model: type_model_name children: list[str] + feedback: FeedbackOutput | None + used_chunks: list[Chunk] | None parent: str | None diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/usecases/bot.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/usecases/bot.py index 650fe26ad..2bdb7c3cc 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/usecases/bot.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/usecases/bot.py @@ -172,6 +172,7 @@ def modify_owned_bot( source_urls = [] sitemap_urls = [] filenames = [] + sync_status: type_sync_status = "QUEUED" if modify_input.knowledge: source_urls = modify_input.knowledge.source_urls @@ -206,6 +207,12 @@ def modify_owned_bot( else DEFAULT_EMBEDDING_CONFIG["chunk_overlap"] ) + # if knowledge and embedding_params are not updated, skip embeding process. + # 'sync_status = "QUEUED"' will execute embeding process and update dynamodb record. + # 'sync_status= "SUCCEEDED"' will update only dynamodb record. + bot = find_private_bot_by_id(user_id, bot_id) + sync_status = "QUEUED" if modify_input.is_embedding_required(bot) else "SUCCEEDED" + update_bot( user_id, bot_id, @@ -221,9 +228,10 @@ def modify_owned_bot( sitemap_urls=sitemap_urls, filenames=filenames, ), - sync_status="QUEUED", + sync_status=sync_status, sync_status_reason="", ) + return BotModifyOutput( id=bot_id, title=modify_input.title, diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/usecases/chat.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/usecases/chat.py index 175ff8c41..0d3c6d0dc 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/usecases/chat.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/usecases/chat.py @@ -1,3 +1,12 @@ +from ulid import ULID +from app.vector_search import SearchResult, get_source_link, search_related_docs +from app.utils import ( + get_bedrock_client, + get_anthropic_client, + get_current_time, + is_running_on_lambda, + is_anthropic_model, +) import json import logging from copy import deepcopy @@ -5,7 +14,12 @@ from typing import Literal from anthropic.types import Message as AnthropicMessage -from app.bedrock import calculate_price, compose_args_for_anthropic_client, get_model_id +from app.bedrock import ( + calculate_price, + compose_args, + get_bedrock_response, + InvocationMetrics, +) from app.config import GENERATION_CONFIG, SEARCH_CONFIG from app.repositories.conversation import ( RecordNotFoundError, @@ -14,6 +28,7 @@ ) from app.repositories.custom_bot import find_alias_by_id, store_alias from app.repositories.models.conversation import ( + ChunkModel, ContentModel, ConversationModel, MessageModel, @@ -22,14 +37,21 @@ from app.routes.schemas.conversation import ( ChatInput, ChatOutput, + Chunk, Content, Conversation, + FeedbackOutput, MessageOutput, RelatedDocumentsOutput, ) from app.usecases.bot import fetch_bot, modify_bot_last_used_time from app.utils import get_anthropic_client, get_current_time, is_running_on_lambda -from app.vector_search import SearchResult, get_source_link, search_related_docs +from app.vector_search import ( + SearchResult, + filter_used_results, + get_source_link, + search_related_docs, +) from ulid import ULID logger = logging.getLogger(__name__) @@ -65,7 +87,7 @@ def prepare_conversation( ) initial_message_map = { - # Dummy system message + # Dummy system message, which is used for root node of the message tree. "system": MessageModel( role="system", content=[ @@ -79,6 +101,8 @@ def prepare_conversation( children=[], parent=None, create_time=current_time, + feedback=None, + used_chunks=None, ) } parent_id = "system" @@ -100,6 +124,8 @@ def prepare_conversation( children=[], parent="system", create_time=current_time, + feedback=None, + used_chunks=None, ) initial_message_map["system"].children.append("instruction") @@ -157,6 +183,8 @@ def prepare_conversation( children=[], parent=parent_id, create_time=current_time, + feedback=None, + used_chunks=None, ) conversation.message_map[message_id] = new_message conversation.message_map[parent_id].children.append(message_id) # type: ignore @@ -259,18 +287,19 @@ def chat(user_id: str, chat_input: ChatInput) -> ChatOutput: user_msg_id, conversation, bot = prepare_conversation(user_id, chat_input) message_map = conversation.message_map + search_results = [] if bot and is_running_on_lambda(): # NOTE: `is_running_on_lambda`is a workaround for local testing due to no postgres mock. # Fetch most related documents from vector store # NOTE: Currently embedding not support multi-modal. For now, use the last content. query = conversation.message_map[user_msg_id].content[-1].body - results = search_related_docs( + search_results = search_related_docs( bot_id=bot.id, limit=SEARCH_CONFIG["max_results"], query=query ) - logger.info(f"Search results from vector store: {results}") + logger.info(f"Search results from vector store: {search_results}") # Insert contexts to instruction - conversation_with_context = insert_knowledge(conversation, results) + conversation_with_context = insert_knowledge(conversation, search_results) message_map = conversation_with_context.message_map messages = trace_to_root( @@ -279,7 +308,7 @@ def chat(user_id: str, chat_input: ChatInput) -> ChatOutput: messages.append(chat_input.message) # type: ignore # Create payload to invoke Bedrock - args = compose_args_for_anthropic_client( + args = compose_args( messages=messages, model=chat_input.message.model, instruction=( @@ -288,8 +317,22 @@ def chat(user_id: str, chat_input: ChatInput) -> ChatOutput: else None ), ) - response: AnthropicMessage = client.messages.create(**args) - reply_txt = response.content[0].text + + if is_anthropic_model(args["model"]): + client = get_anthropic_client() + response: AnthropicMessage = client.messages.create(**args) + reply_txt = response.content[0].text + else: + response = get_bedrock_response(args) # type: ignore + reply_txt = response["outputs"][0]["text"] # type: ignore + + # Used chunks for RAG generation + used_chunks = None + if bot and is_running_on_lambda(): + used_chunks = [ + ChunkModel(content=r.content, source=r.source, rank=r.rank) + for r in filter_used_results(reply_txt, search_results) + ] # Issue id for new assistant message assistant_msg_id = str(ULID()) @@ -301,6 +344,8 @@ def chat(user_id: str, chat_input: ChatInput) -> ChatOutput: children=[], parent=user_msg_id, create_time=get_current_time(), + feedback=None, + used_chunks=used_chunks, ) conversation.message_map[assistant_msg_id] = message @@ -308,11 +353,14 @@ def chat(user_id: str, chat_input: ChatInput) -> ChatOutput: conversation.message_map[user_msg_id].children.append(assistant_msg_id) conversation.last_message_id = assistant_msg_id - # Update total pricing - input_tokens = response.usage.input_tokens - output_tokens = response.usage.output_tokens - - logger.debug(f"Input tokens: {input_tokens}, Output tokens: {output_tokens}") + if is_anthropic_model(args["model"]): + # Update total pricing + input_tokens = response.usage.input_tokens + output_tokens = response.usage.output_tokens + else: + metrics: InvocationMetrics = response["amazon-bedrock-invocationMetrics"] # type: ignore + input_tokens = metrics.input_tokens + output_tokens = metrics.output_tokens price = calculate_price(chat_input.message.model, input_tokens, output_tokens) conversation.total_price += price @@ -341,6 +389,19 @@ def chat(user_id: str, chat_input: ChatInput) -> ChatOutput: model=message.model, children=message.children, parent=message.parent, + feedback=None, + used_chunks=( + [ + Chunk( + content=c.content, + source=c.source, + rank=c.rank, + ) + for c in message.used_chunks + ] + if message.used_chunks + else None + ), ), bot_id=conversation.bot_id, ) @@ -357,6 +418,9 @@ def propose_conversation_title( "claude-v3-opus", "claude-v3-sonnet", "claude-v3-haiku", + "mistral-7b-instruct", + "mixtral-8x7b-instruct", + "mistral-large", ] = "claude-v3-haiku", ) -> str: PROMPT = """Reading the conversation above, what is the appropriate title for the conversation? When answering the title, please follow the rules below: @@ -389,16 +453,22 @@ def propose_conversation_title( children=[], parent=conversation.last_message_id, create_time=get_current_time(), + feedback=None, + used_chunks=None, ) messages.append(new_message) # Invoke Bedrock - args = compose_args_for_anthropic_client( + args = compose_args( messages=messages, model=model, ) - response = client.messages.create(**args) - reply_txt = response.content[0].text + if is_anthropic_model(args["model"]): + response = client.messages.create(**args) + reply_txt = response.content[0].text + else: + response: AnthropicMessage = get_bedrock_response(args)["outputs"][0] # type: ignore[no-redef] + reply_txt = response["text"] return reply_txt @@ -419,6 +489,27 @@ def fetch_conversation(user_id: str, conversation_id: str) -> Conversation: model=message.model, children=message.children, parent=message.parent, + feedback=( + FeedbackOutput( + thumbs_up=message.feedback.thumbs_up, + category=message.feedback.category, + comment=message.feedback.comment, + ) + if message.feedback + else None + ), + used_chunks=( + [ + Chunk( + content=c.content, + source=c.source, + rank=c.rank, + ) + for c in message.used_chunks + ] + if message.used_chunks + else None + ), ) for message_id, message in conversation.message_map.items() } diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/utils.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/utils.py index e3c14e838..7303c18e9 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/utils.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/utils.py @@ -19,6 +19,10 @@ def is_running_on_lambda(): return "AWS_EXECUTION_ENV" in os.environ +def is_anthropic_model(model_id: str) -> bool: + return model_id.startswith("anthropic") or False + + def get_bedrock_client(region=BEDROCK_REGION): client = boto3.client("bedrock-runtime", region) return client diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/vector_search.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/vector_search.py index c4d844c30..95d87e302 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/vector_search.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/vector_search.py @@ -1,6 +1,7 @@ import json import logging import os +import re from typing import Literal import pg8000 @@ -24,6 +25,31 @@ class SearchResult(BaseModel): rank: int +def filter_used_results( + generated_text: str, search_results: list[SearchResult] +) -> list[SearchResult]: + """Filter the search results based on the citations in the generated text. + Note that the citations in the generated text are in the format of [^rank]. + """ + used_results: list[SearchResult] = [] + + try: + # Extract citations from the generated text + citations = [ + citation.strip("[]^") + for citation in re.findall(r"\[\^(\d+)\]", generated_text) + ] + except Exception as e: + logger.error(f"Error extracting citations from the generated text: {e}") + return used_results + + for result in search_results: + if str(result.rank) in citations: + used_results.append(result) + + return used_results + + def get_source_link(source: str) -> tuple[Literal["s3", "url"], str]: if source.startswith("s3://"): s3_path = source[5:] # Remove "s3://" prefix @@ -71,10 +97,10 @@ def search_related_docs(bot_id: str, limit: int, query: str) -> list[SearchResul # It's important to choose the same distance metric as the one used for indexing. # Ref: https://github.com/pgvector/pgvector?tab=readme-ov-file#getting-started search_query = """ -SELECT id, botid, content, source, embedding -FROM items -WHERE botid = %s -ORDER BY embedding <-> %s +SELECT id, botid, content, source, embedding +FROM items +WHERE botid = %s +ORDER BY embedding <-> %s LIMIT %s """ cursor.execute(search_query, (bot_id, json.dumps(query_embedding), limit)) diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/websocket.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/websocket.py index 1f917a28e..77a4e6c5e 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/websocket.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/app/websocket.py @@ -7,16 +7,22 @@ import boto3 from anthropic.types import ContentBlockDeltaEvent, MessageDeltaEvent, MessageStopEvent +from anthropic.types import Message as AnthropicMessage from app.auth import verify_token -from app.bedrock import calculate_price, compose_args_for_anthropic_client +from app.bedrock import calculate_price, compose_args from app.config import GENERATION_CONFIG, SEARCH_CONFIG from app.repositories.conversation import RecordNotFoundError, store_conversation -from app.repositories.models.conversation import ContentModel, MessageModel +from app.repositories.models.conversation import ChunkModel, ContentModel, MessageModel from app.routes.schemas.conversation import ChatInputWithToken from app.usecases.bot import modify_bot_last_used_time -from app.usecases.chat import insert_knowledge, prepare_conversation, trace_to_root -from app.utils import get_anthropic_client, get_current_time -from app.vector_search import search_related_docs +from app.usecases.chat import ( + insert_knowledge, + prepare_conversation, + trace_to_root, + get_bedrock_response, +) +from app.utils import get_anthropic_client, get_current_time, is_anthropic_model +from app.vector_search import filter_used_results, search_related_docs from boto3.dynamodb.conditions import Key from ulid import ULID @@ -63,6 +69,7 @@ def process_chat_input( return {"statusCode": 400, "body": "Invalid request."} message_map = conversation.message_map + search_results = [] if bot and bot.has_knowledge(): gatewayapi.post_to_connection( ConnectionId=connection_id, @@ -75,13 +82,13 @@ def process_chat_input( # Fetch most related documents from vector store # NOTE: Currently embedding not support multi-modal. For now, use the last text content. query = conversation.message_map[user_msg_id].content[-1].body - results = search_related_docs( + search_results = search_related_docs( bot_id=bot.id, limit=SEARCH_CONFIG["max_results"], query=query ) - logger.info(f"Search results from vector store: {results}") + logger.info(f"Search results from vector store: {search_results}") # Insert contexts to instruction - conversation_with_context = insert_knowledge(conversation, results) + conversation_with_context = insert_knowledge(conversation, search_results) message_map = conversation_with_context.message_map messages = trace_to_root( @@ -90,8 +97,7 @@ def process_chat_input( ) messages.append(chat_input.message) # type: ignore - # Invoke Bedrock - args = compose_args_for_anthropic_client( + args = compose_args( messages, chat_input.message.model, instruction=( @@ -101,92 +107,177 @@ def process_chat_input( ), stream=True, ) + + is_anthropic = is_anthropic_model(args["model"]) # logger.debug(f"Invoking bedrock with args: {args}") try: - # Invoke bedrock streaming api - response = client.messages.create(**args) + if is_anthropic: + response = client.messages.create(**args) + else: + # Invoke bedrock streaming api + response = get_bedrock_response(args) except Exception as e: logger.error(f"Failed to invoke bedrock: {e}") return {"statusCode": 500, "body": "Failed to invoke bedrock."} completions: list[str] = [] last_data_to_send: bytes - for event in response: - # NOTE: following is the example of event sequence: - # MessageStartEvent(message=Message(id='compl_01GwmkwncsptaeBopeaR4eWE', content=[], model='claude-instant-1.2', role='assistant', stop_reason=None, stop_sequence=None, type='message', usage=Usage(input_tokens=21, output_tokens=1)), type='message_start') - # ContentBlockStartEvent(content_block=ContentBlock(text='', type='text'), index=0, type='content_block_start') - # ... - # ContentBlockDeltaEvent(delta=TextDelta(text='です', type='text_delta'), index=0, type='content_block_delta') - # ContentBlockStopEvent(index=0, type='content_block_stop') - # MessageDeltaEvent(delta=Delta(stop_reason='end_turn', stop_sequence=None), type='message_delta', usage=MessageDeltaUsage(output_tokens=26)) - # MessageStopEvent(type='message_stop', amazon-bedrock-invocationMetrics={'inputTokenCount': 21, 'outputTokenCount': 25, 'invocationLatency': 621, 'firstByteLatency': 279}) - - if isinstance(event, ContentBlockDeltaEvent): - completions.append(event.delta.text) - try: - # Send completion - data_to_send = json.dumps( + if is_anthropic: + for event in response: + # NOTE: following is the example of event sequence: + # MessageStartEvent(message=Message(id='compl_01GwmkwncsptaeBopeaR4eWE', content=[], model='claude-instant-1.2', role='assistant', stop_reason=None, stop_sequence=None, type='message', usage=Usage(input_tokens=21, output_tokens=1)), type='message_start') + # ContentBlockStartEvent(content_block=ContentBlock(text='', type='text'), index=0, type='content_block_start') + # ... + # ContentBlockDeltaEvent(delta=TextDelta(text='です', type='text_delta'), index=0, type='content_block_delta') + # ContentBlockStopEvent(index=0, type='content_block_stop') + # MessageDeltaEvent(delta=Delta(stop_reason='end_turn', stop_sequence=None), type='message_delta', usage=MessageDeltaUsage(output_tokens=26)) + # MessageStopEvent(type='message_stop', amazon-bedrock-invocationMetrics={'inputTokenCount': 21, 'outputTokenCount': 25, 'invocationLatency': 621, 'firstByteLatency': 279}) + + if isinstance(event, ContentBlockDeltaEvent): + completions.append(event.delta.text) + try: + # Send completion + data_to_send = json.dumps( + dict( + status="STREAMING", + completion=event.delta.text, + ) + ).encode("utf-8") + gatewayapi.post_to_connection( + ConnectionId=connection_id, Data=data_to_send + ) + except Exception as e: + logger.error(f"Failed to post message: {str(e)}") + return { + "statusCode": 500, + "body": "Failed to send message to connection.", + } + elif isinstance(event, MessageDeltaEvent): + logger.debug(f"Received message delta event: {event.delta}") + last_data_to_send = json.dumps( dict( - status="STREAMING", - completion=event.delta.text, + completion="", + stop_reason=event.delta.stop_reason, ) ).encode("utf-8") - gatewayapi.post_to_connection( - ConnectionId=connection_id, Data=data_to_send + elif isinstance(event, MessageStopEvent): + # Persist conversation before finish streaming so that front-end can avoid 404 issue + concatenated = "".join(completions) + + # Used chunks for RAG generation + used_chunks = None + if bot: + used_chunks = [ + ChunkModel(content=r.content, source=r.source, rank=r.rank) + for r in filter_used_results(concatenated, search_results) + ] + + # Append entire completion as the last message + assistant_msg_id = str(ULID()) + message = MessageModel( + role="assistant", + content=[ + ContentModel( + content_type="text", body=concatenated, media_type=None + ) + ], + model=chat_input.message.model, + children=[], + parent=user_msg_id, + create_time=get_current_time(), + feedback=None, + used_chunks=used_chunks, ) - except Exception as e: - logger.error(f"Failed to post message: {str(e)}") - return { - "statusCode": 500, - "body": "Failed to send message to connection.", - } - elif isinstance(event, MessageDeltaEvent): - logger.debug(f"Received message delta event: {event.delta}") - last_data_to_send = json.dumps( - dict( - completion="", - stop_reason=event.delta.stop_reason, + conversation.message_map[assistant_msg_id] = message + # Append children to parent + conversation.message_map[user_msg_id].children.append(assistant_msg_id) + conversation.last_message_id = assistant_msg_id + + # Update total pricing + metrics = event.model_dump()["amazon-bedrock-invocationMetrics"] + input_token_count = metrics.get("inputTokenCount") + output_token_count = metrics.get("outputTokenCount") + + logger.debug( + f"Input token count: {input_token_count}, output token count: {output_token_count}" ) - ).encode("utf-8") - elif isinstance(event, MessageStopEvent): - # Persist conversation before finish streaming so that front-end can avoid 404 issue - concatenated = "".join(completions) - # Append entire completion as the last message - assistant_msg_id = str(ULID()) - message = MessageModel( - role="assistant", - content=[ - ContentModel( - content_type="text", body=concatenated, media_type=None + + price = calculate_price( + chat_input.message.model, input_token_count, output_token_count + ) + conversation.total_price += price + + store_conversation(user_id, conversation) + else: + continue + else: + used_chunks = None + for event in response.get("body"): + chunk = event.get("chunk") + if chunk: + msg_chunk = json.loads(chunk.get("bytes").decode()) + is_stop = msg_chunk["outputs"][0]["stop_reason"] + if not is_stop: + msg = msg_chunk["outputs"][0]["text"] + completions.append(msg) + data_to_send = json.dumps( + dict( + status="STREAMING", + completion=msg, + ) + ).encode("utf-8") + gatewayapi.post_to_connection( + ConnectionId=connection_id, Data=data_to_send ) - ], - model=chat_input.message.model, - children=[], - parent=user_msg_id, - create_time=get_current_time(), - ) - conversation.message_map[assistant_msg_id] = message - # Append children to parent - conversation.message_map[user_msg_id].children.append(assistant_msg_id) - conversation.last_message_id = assistant_msg_id - - # Update total pricing - metrics = event.model_dump()["amazon-bedrock-invocationMetrics"] - input_token_count = metrics.get("inputTokenCount") - output_token_count = metrics.get("outputTokenCount") - - logger.debug( - f"Input token count: {input_token_count}, output token count: {output_token_count}" - ) + else: + last_data_to_send = json.dumps( + dict(completion="", stop_reason=is_stop) + ).encode("utf-8") + + concatenated = "".join(completions) + # Used chunks for RAG generation + if bot: + used_chunks = [ + ChunkModel(content=r.content, source=r.source, rank=r.rank) + for r in filter_used_results(concatenated, search_results) + ] + assistant_msg_id = str(ULID()) + message = MessageModel( + role="assistant", + content=[ + ContentModel( + content_type="text", body=concatenated, media_type=None + ) + ], + model=chat_input.message.model, + children=[], + parent=user_msg_id, + create_time=get_current_time(), + feedback=None, + used_chunks=used_chunks, + ) + conversation.message_map[assistant_msg_id] = message + # Append children to parent + conversation.message_map[user_msg_id].children.append( + assistant_msg_id + ) + conversation.last_message_id = assistant_msg_id - price = calculate_price( - chat_input.message.model, input_token_count, output_token_count - ) - conversation.total_price += price + # Update total pricing + metrics = msg_chunk["amazon-bedrock-invocationMetrics"] + input_token_count = metrics.get("inputTokenCount") + output_token_count = metrics.get("outputTokenCount") - store_conversation(user_id, conversation) - else: - continue + logger.debug( + f"Input token count: {input_token_count}, output token count: {output_token_count}" + ) + + price = calculate_price( + chat_input.message.model, input_token_count, output_token_count + ) + conversation.total_price += price + + store_conversation(user_id, conversation) # Send last completion after saving conversation try: diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/embedding.requirements.txt b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/embedding.requirements.txt index 5043c57f6..26a3e400e 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/embedding.requirements.txt +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/embedding.requirements.txt @@ -1,6 +1,6 @@ boto3==1.28.57 llama-index==0.10.19 -pydantic==2.1.1 +pydantic==2.4.0 youtube-transcript-api==0.6.1 playwright==1.40.0 requests==2.31.0 diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/requirements.txt b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/requirements.txt index dfbcdcf7e..2c1c54ef3 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/requirements.txt +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/requirements.txt @@ -1,7 +1,7 @@ fastapi==0.109.1 requests==2.31.0 types-requests==2.31.0 -pydantic==2.1.1 +pydantic==2.4.0 pyhumps==3.8.0 uvicorn==0.23.1 python-ulid==1.1.0 diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/tests/test_repositories/test_conversation.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/tests/test_repositories/test_conversation.py index f9b98a675..faa177aff 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/tests/test_repositories/test_conversation.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/tests/test_repositories/test_conversation.py @@ -1,31 +1,28 @@ import sys import unittest - sys.path.append(".") from app.config import DEFAULT_EMBEDDING_CONFIG - from app.repositories.conversation import ( ContentModel, ConversationModel, MessageModel, RecordNotFoundError, - _get_table_client, change_conversation_title, - compose_conv_id, delete_conversation_by_id, delete_conversation_by_user_id, find_conversation_by_id, find_conversation_by_user_id, store_conversation, + update_feedback, ) from app.repositories.custom_bot import ( delete_bot_by_id, - find_private_bot_by_id, find_private_bots_by_user_id, store_bot, ) +from app.repositories.models.conversation import FeedbackModel from app.repositories.models.custom_bot import ( BotModel, EmbeddingParamsModel, @@ -138,6 +135,8 @@ def test_store_and_find_conversation(self): children=["x", "y"], parent="z", create_time=1627984879.9, + feedback=None, + used_chunks=None, ) }, last_message_id="x", @@ -189,6 +188,25 @@ def test_store_and_find_conversation(self): ) self.assertEqual(found_conversation.title, "Updated title") + # Test give a feedback + self.assertIsNone(found_conversation.message_map["a"].feedback) + response = update_feedback( + user_id="user", + conversation_id="1", + message_id="a", + feedback=FeedbackModel( + thumbs_up=True, category="Good", comment="The response is pretty good." + ), + ) + found_conversation = find_conversation_by_id( + user_id="user", conversation_id="1" + ) + feedback = found_conversation.message_map["a"].feedback + self.assertIsNotNone(feedback) + self.assertEqual(feedback.thumbs_up, True) # type: ignore + self.assertEqual(feedback.category, "Good") # type: ignore + self.assertEqual(feedback.comment, "The response is pretty good.") # type: ignore + # Test deleting conversation by id delete_conversation_by_id(user_id="user", conversation_id="1") with self.assertRaises(RecordNotFoundError): @@ -217,6 +235,8 @@ def test_store_and_find_large_conversation(self): children=[], parent=None, create_time=1627984879.9, + feedback=None, + used_chunks=None, ) for i in range(10) # Create 10 large messages } @@ -297,6 +317,8 @@ def setUp(self) -> None: children=["x", "y"], parent="z", create_time=1627984879.9, + feedback=None, + used_chunks=None, ) }, last_message_id="x", @@ -324,6 +346,8 @@ def setUp(self) -> None: children=["x", "y"], parent="z", create_time=1627984879.9, + feedback=None, + used_chunks=None, ) }, last_message_id="x", diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/tests/test_routes/test_schemas/test_conversation.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/tests/test_routes/test_schemas/test_conversation.py new file mode 100644 index 000000000..c18f6d082 --- /dev/null +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/tests/test_routes/test_schemas/test_conversation.py @@ -0,0 +1,29 @@ +import sys + +sys.path.append(".") +import unittest + +from app.routes.schemas.conversation import FeedbackInput +from pydantic import ValidationError + + +class TestFeedbackInput(unittest.TestCase): + def test_create_input_valid_no_category(self): + obj = FeedbackInput(thumbs_up=True, category=None, comment="Excellent!") + self.assertTrue(obj.thumbs_up) + self.assertIsNone(obj.category) + self.assertEqual(obj.comment, "Excellent!") + + def test_create_input_invalid_no_category(self): + with self.assertRaises(ValidationError): + FeedbackInput(thumbs_up=False, category=None, comment="Needs improvement.") + + def test_create_input_valid_no_comment(self): + obj = FeedbackInput(thumbs_up=False, category="DISLIKE", comment=None) + self.assertFalse(obj.thumbs_up) + self.assertEqual(obj.category, "DISLIKE") + self.assertIsNone(obj.comment) + + +if __name__ == "__main__": + unittest.main() diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/tests/test_usecases/test_chat.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/tests/test_usecases/test_chat.py index 809fba2b4..ceb47bd20 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/tests/test_usecases/test_chat.py +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/tests/test_usecases/test_chat.py @@ -4,12 +4,6 @@ import unittest from pprint import pprint -from tests.test_usecases.utils.bot_factory import ( - create_test_private_bot, - create_test_public_bot, - create_test_instruction_template, -) - from anthropic.types import MessageStopEvent from app.bedrock import get_model_id from app.config import GENERATION_CONFIG @@ -30,7 +24,6 @@ ConversationModel, MessageModel, ) - from app.routes.schemas.conversation import ( ChatInput, ChatOutput, @@ -48,8 +41,14 @@ ) from app.utils import get_anthropic_client from app.vector_search import SearchResult +from tests.test_usecases.utils.bot_factory import ( + create_test_instruction_template, + create_test_private_bot, + create_test_public_bot, +) MODEL: type_model_name = "claude-instant-v1" +MISTRAL_MODEL: type_model_name = "mistral-7b-instruct" class TestTraceToRoot(unittest.TestCase): @@ -64,6 +63,8 @@ def test_trace_to_root(self): children=["bot_1"], parent=None, create_time=1627984879.9, + feedback=None, + used_chunks=None, ), "bot_1": MessageModel( role="assistant", @@ -74,6 +75,8 @@ def test_trace_to_root(self): children=["user_2"], parent="user_1", create_time=1627984879.9, + feedback=None, + used_chunks=None, ), "user_2": MessageModel( role="user", @@ -84,6 +87,8 @@ def test_trace_to_root(self): children=["bot_2"], parent="bot_1", create_time=1627984879.9, + feedback=None, + used_chunks=None, ), "bot_2": MessageModel( role="assistant", @@ -94,6 +99,8 @@ def test_trace_to_root(self): children=["user_3a", "user_3b"], parent="user_2", create_time=1627984879.9, + feedback=None, + used_chunks=None, ), "user_3a": MessageModel( role="user", @@ -104,6 +111,8 @@ def test_trace_to_root(self): children=[], parent="bot_2", create_time=1627984879.9, + feedback=None, + used_chunks=None, ), "user_3b": MessageModel( role="user", @@ -114,6 +123,8 @@ def test_trace_to_root(self): children=[], parent="bot_2", create_time=1627984879.9, + feedback=None, + used_chunks=None, ), } messages = trace_to_root("user_3a", message_map) @@ -177,6 +188,50 @@ def test_chat(self): self.assertEqual(conv.last_message_id, second_key) self.assertNotEqual(conv.total_price, 0) + def test_chat_mistral(self): + prompt = "あなたの名前は何ですか?" + body = f"[INST]{prompt}[/INST]" + + chat_input = ChatInput( + conversation_id="test_conversation_id", + message=MessageInput( + role="user", + content=[ + Content( + content_type="text", + body=body, + media_type=None, + ) + ], + model=MISTRAL_MODEL, + parent_message_id=None, + message_id=None, + ), + bot_id=None, + ) + output: ChatOutput = chat(user_id="user1", chat_input=chat_input) + self.output = output + + pprint(output.model_dump()) + self.assertNotEqual(output.conversation_id, "") + + conv = find_conversation_by_id( + user_id="user1", conversation_id=output.conversation_id + ) + self.assertEqual(len(conv.message_map), 3) + for k, v in conv.message_map.items(): + if v.parent == "system": + first_key = k + first_message = v + elif v.parent: + second_key = k + second_message = v + + self.assertEqual(second_message.parent, first_key) + self.assertEqual(first_message.children, [second_key]) + self.assertEqual(conv.last_message_id, second_key) + self.assertNotEqual(conv.total_price, 0) + def tearDown(self) -> None: delete_conversation_by_id("user1", self.output.conversation_id) @@ -241,6 +296,8 @@ def setUp(self) -> None: children=["1-assistant"], parent=None, create_time=1627984879.9, + feedback=None, + used_chunks=None, ), "1-assistant": MessageModel( role="assistant", @@ -255,6 +312,8 @@ def setUp(self) -> None: children=[], parent="1-user", create_time=1627984879.9, + feedback=None, + used_chunks=None, ), }, bot_id=None, @@ -325,6 +384,8 @@ def setUp(self) -> None: children=["a-2"], parent=None, create_time=1627984879.9, + feedback=None, + used_chunks=None, ), "a-2": MessageModel( role="assistant", @@ -339,6 +400,8 @@ def setUp(self) -> None: children=[], parent="a-1", create_time=1627984879.9, + feedback=None, + used_chunks=None, ), "b-1": MessageModel( role="user", @@ -353,6 +416,8 @@ def setUp(self) -> None: children=["b-2"], parent=None, create_time=1627984879.9, + feedback=None, + used_chunks=None, ), "b-2": MessageModel( role="assistant", @@ -367,6 +432,8 @@ def setUp(self) -> None: children=[], parent="b-1", create_time=1627984879.9, + feedback=None, + used_chunks=None, ), }, bot_id=None, @@ -451,14 +518,22 @@ def setUp(self) -> None: bot_id=None, ) output: ChatOutput = chat(user_id="user1", chat_input=chat_input) - print(output) self.output = output + chat_input.message.model = MISTRAL_MODEL + mistral_output: ChatOutput = chat(user_id="user1", chat_input=chat_input) + self.mistral_output = mistral_output + print(mistral_output) + def test_propose_title(self): title = propose_conversation_title("user1", self.output.conversation_id) print(f"[title]: {title}") + def test_propose_title_mistral(self): + title = propose_conversation_title("user1", self.mistral_output.conversation_id) + print(f"[title]: {title}") + def tearDown(self) -> None: delete_conversation_by_id("user1", self.output.conversation_id) @@ -519,7 +594,7 @@ def test_chat_with_private_bot(self): conv = find_conversation_by_id("user1", output.conversation_id) self.assertEqual(len(conv.message_map["system"].children), 1) - self.assertEqual(conv.message_map["system"].children[0], "instruction") + self.assertEqual(conv.message_map["system"].children, ["instruction"]) self.assertEqual(len(conv.message_map["instruction"].children), 1) # Second message @@ -687,6 +762,8 @@ def test_insert_knowledge(self): children=["1-user"], parent=None, create_time=1627984879.9, + feedback=None, + used_chunks=None, ), "1-user": MessageModel( role="user", @@ -701,6 +778,8 @@ def test_insert_knowledge(self): children=[], parent="instruction", create_time=1627984879.9, + feedback=None, + used_chunks=None, ), }, bot_id="bot1", diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/tests/test_vector_search/test_vector_search.py b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/tests/test_vector_search/test_vector_search.py new file mode 100644 index 000000000..a6fa52ada --- /dev/null +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/backend/python/tests/test_vector_search/test_vector_search.py @@ -0,0 +1,52 @@ +import sys +import unittest + +sys.path.append(".") + +from app.vector_search import SearchResult, filter_used_results + + +class TestVectorSearch(unittest.TestCase): + def test_filter_used_results(self): + search_results = [ + SearchResult(bot_id="1", content="content1", source="source1", rank=1), + SearchResult(bot_id="2", content="content2", source="source2", rank=2), + SearchResult(bot_id="3", content="content3", source="source3", rank=3), + ] + + generated_text = "This is a test [^1] [^3]" + + used_results = filter_used_results(generated_text, search_results) + self.assertEqual(len(used_results), 2) + self.assertEqual(used_results[0].rank, 1) + self.assertEqual(used_results[1].rank, 3) + + def test_no_reference_filter_used_results(self): + search_results = [ + SearchResult(bot_id="1", content="content1", source="source1", rank=1), + SearchResult(bot_id="2", content="content2", source="source2", rank=2), + SearchResult(bot_id="3", content="content3", source="source3", rank=3), + ] + + # 4 is not in the search results + generated_text = "This is a test [^4]" + + used_results = filter_used_results(generated_text, search_results) + self.assertEqual(len(used_results), 0) + + def test_format_not_match_filter_used_results(self): + search_results = [ + SearchResult(bot_id="1", content="content1", source="source1", rank=1), + SearchResult(bot_id="2", content="content2", source="source2", rank=2), + SearchResult(bot_id="3", content="content3", source="source3", rank=3), + ] + + # format not match + generated_text = "This is a test 1 3" + + used_results = filter_used_results(generated_text, search_results) + self.assertEqual(len(used_results), 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/package.json b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/package.json index 273362396..f986d5b85 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/package.json +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/package.json @@ -22,6 +22,7 @@ "i18next-browser-languagedetector": "^7.1.0", "immer": "^10.0.2", "react": "^18.2.0", + "react-children-utilities": "^2.10.0", "react-device-detect": "^2.2.3", "react-dom": "^18.2.0", "react-i18next": "^13.3.1", @@ -29,8 +30,11 @@ "react-markdown": "^8.0.7", "react-router-dom": "^6.14.2", "react-syntax-highlighter": "^15.5.0", + "rehype-external-links": "^3.0.0", + "rehype-katex": "^6.0.0", "remark-breaks": "^3.0.3", "remark-gfm": "^3.0.1", + "remark-math": "^6.0.0", "swr": "^2.2.0", "tailwind-merge": "^2.2.0", "ulid": "^2.3.0", diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/@types/conversation.d.ts b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/@types/conversation.d.ts index b66fef344..9eea2817e 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/@types/conversation.d.ts +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/@types/conversation.d.ts @@ -4,7 +4,10 @@ export type Model = | 'claude-v2' | 'claude-v3-opus' | 'claude-v3-sonnet' - | 'claude-v3-haiku'; + | 'claude-v3-haiku' + | 'mistral-7b-instruct' + | 'mixtral-8x7b-instruct' + | 'mistral-large'; export type Content = { contentType: 'text' | 'image'; mediaType?: string; @@ -15,6 +18,7 @@ export type MessageContent = { role: Role; content: Content[]; model: Model; + feedback: null | Feedback; }; export type RelatedDocument = { @@ -74,3 +78,15 @@ export type MessageMap = { export type Conversation = ConversationMeta & { messageMap: MessageMap; }; + +export type PutFeedbackRequest = { + thumbsUp: boolean; + category: null | string; + comment: null | string; +}; + +export type Feedback = { + thumbsUp: boolean; + category: string; + comment: string; +}; diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/components/ChatMessage.tsx b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/components/ChatMessage.tsx index 504c04648..168adcd94 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/components/ChatMessage.tsx +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/components/ChatMessage.tsx @@ -1,15 +1,26 @@ import React, { useCallback, useEffect, useMemo, useState } from 'react'; import ChatMessageMarkdown from './ChatMessageMarkdown'; import ButtonCopy from './ButtonCopy'; -import { PiCaretLeftBold, PiNotePencil, PiUserFill } from 'react-icons/pi'; +import { + PiCaretLeftBold, + PiNotePencil, + PiUserFill, + PiThumbsDown, + PiThumbsDownFill, +} from 'react-icons/pi'; import { BaseProps } from '../@types/common'; -import { DisplayMessageContent, RelatedDocument } from '../@types/conversation'; +import { + DisplayMessageContent, + RelatedDocument, + PutFeedbackRequest, +} from '../@types/conversation'; import ButtonIcon from './ButtonIcon'; import Textarea from './Textarea'; import Button from './Button'; import ModalDialog from './ModalDialog'; import { useTranslation } from 'react-i18next'; import useChat from '../hooks/useChat'; +import DialogFeedback from './DialogFeedback'; type Props = BaseProps & { chatContent?: DisplayMessageContent; @@ -21,8 +32,9 @@ const ChatMessage: React.FC = (props) => { const { t } = useTranslation(); const [isEdit, setIsEdit] = useState(false); const [changedContent, setChangedContent] = useState(''); + const [isFeedbackOpen, setIsFeedbackOpen] = useState(false); - const { getRelatedDocuments } = useChat(); + const { getRelatedDocuments, conversationId, giveFeedback } = useChat(); const [relatedDocuments, setRelatedDocuments] = useState( [] ); @@ -61,6 +73,16 @@ const ChatMessage: React.FC = (props) => { setIsEdit(false); }, [changedContent, chatContent?.sibling, props]); + const handleFeedbackSubmit = useCallback( + (messageId: string, feedback: PutFeedbackRequest) => { + if (chatContent && conversationId) { + giveFeedback(messageId, feedback); + } + setIsFeedbackOpen(false); + }, + [chatContent, conversationId, giveFeedback] + ); + return (
@@ -174,10 +196,10 @@ const ChatMessage: React.FC = (props) => {
-
+
{chatContent?.role === 'user' && !isEdit && ( { setChangedContent(chatContent.content[0].body); setIsEdit(true); @@ -186,15 +208,35 @@ const ChatMessage: React.FC = (props) => { )} {chatContent?.role === 'assistant' && ( - <> +
+ setIsFeedbackOpen(true)}> + {chatContent.feedback && !chatContent.feedback.thumbsUp ? ( + + ) : ( + + )} + - +
)}
+ setIsFeedbackOpen(false)} + onSubmit={(feedback) => { + if (chatContent) { + handleFeedbackSubmit(chatContent.id, feedback); + } + }} + />
); }; diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/components/ChatMessageMarkdown.tsx b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/components/ChatMessageMarkdown.tsx index 85b848b76..f2efc3af5 100644 --- a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/components/ChatMessageMarkdown.tsx +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/components/ChatMessageMarkdown.tsx @@ -11,6 +11,11 @@ import { twMerge } from 'tailwind-merge'; import { useTranslation } from 'react-i18next'; import { create } from 'zustand'; import { produce } from 'immer'; +import rehypeExternalLinks, { Options } from 'rehype-external-links'; +import rehypeKatex from 'rehype-katex'; +import remarkMath from 'remark-math'; +import "katex/dist/katex.min.css" +import { onlyText } from 'react-children-utilities'; type Props = BaseProps & { children: string; @@ -118,16 +123,32 @@ const ChatMessageMarkdown: React.FC = ({ : children; }, [children]); + const remarkPlugins = useMemo(() => { + return [remarkGfm, remarkBreaks, remarkMath] + }, []) + const rehypePlugins = useMemo(() => { + const rehypeExternalLinksOptions: Options = { + target: '_blank', + properties: { style: "word-break: break-all;", } + } + return [rehypeKatex, [rehypeExternalLinks, rehypeExternalLinksOptions]] + }, []) + return ( @@ -147,44 +168,50 @@ const ChatMessageMarkdown: React.FC = ({ ); }, - // + // eslint-disable-next-line @typescript-eslint/ban-ts-comment + // @ts-ignore sup({ className, children }) { // Footnote's Link is replaced with a component that displays the Reference document return ( - {children.map((child, idx) => { + { // eslint-disable-next-line @typescript-eslint/ban-ts-comment // @ts-ignore - if (child?.props['data-footnote-ref']) { + children.map((child, idx) => { // eslint-disable-next-line @typescript-eslint/ban-ts-comment // @ts-ignore - const href: string = child.props.href ?? ''; - if (/#user-content-fn-[\d]+/.test(href ?? '')) { - const docNo = Number.parseInt( - href.replace('#user-content-fn-', '') - ); - const doc = relatedDocuments?.filter( - (doc) => doc.rank === docNo - )[0]; - + if (child?.props['data-footnote-ref']) { // eslint-disable-next-line @typescript-eslint/ban-ts-comment // @ts-ignore - const refNo = child.props.children[0]; - return ( - - [{refNo}] - - ); + const href: string = child.props.href ?? ''; + if (/#user-content-fn-[\d]+/.test(href ?? '')) { + const docNo = Number.parseInt( + href.replace('#user-content-fn-', '') + ); + const doc = relatedDocuments?.filter( + (doc) => doc.rank === docNo + )[0]; + + // eslint-disable-next-line @typescript-eslint/ban-ts-comment + // @ts-ignore + const refNo = child.props.children[0]; + return ( + + [{refNo}] + + ); + } } - } - return child; - })} + return child; + })} ); }, + // eslint-disable-next-line @typescript-eslint/ban-ts-comment + // @ts-ignore section({ className, children, ...props }) { // Normal Footnote not shown for RAG reference documents // eslint-disable-next-line @typescript-eslint/ban-ts-comment diff --git a/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/components/DialogFeedback.tsx b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/components/DialogFeedback.tsx new file mode 100644 index 000000000..e8ae23cd4 --- /dev/null +++ b/packages/blueprints/gen-ai-chatbot/static-assets/chatbot-genai-components/frontend/src/components/DialogFeedback.tsx @@ -0,0 +1,66 @@ +import React, { useState } from 'react'; +import { BaseProps } from '../@types/common'; +import Button from './Button'; +import ModalDialog from './ModalDialog'; +import { useTranslation } from 'react-i18next'; +import Textarea from './Textarea'; +import Select from './Select'; +import { Feedback } from '../@types/conversation'; + +type Props = BaseProps & { + isOpen: boolean; + thumbsUp: boolean; + feedback?: Feedback; + onSubmit: (feedback: Feedback) => void; + onClose: () => void; +}; + +const DialogFeedback: React.FC = (props) => { + const { t } = useTranslation(); + const categoryOptions = t('feedbackDialog.categories', { + returnObjects: true, + }); + const [category, setCategory] = useState( + props.feedback?.category || categoryOptions[0].value + ); + const [comment, setComment] = useState(props.feedback?.comment || ''); + + const handleSubmit = () => { + props.onSubmit({ thumbsUp: props.thumbsUp, category, comment }); + }; + + return ( + +
+
{t('feedbackDialog.content')}
+ +