Merge pull request #132 from aws-samples/dev

chore: merge commits in dev into cn-marketing branch

docs/draw/online_function_modules.drawio

@@ -0,0 +1,61 @@
+<mxfile host="65bd71144e">
+    <diagram id="Jpx9hanNh7IF5f1KVsD0" name="Page-1">
+        <mxGraphModel dx="881" dy="481" grid="0" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="850" pageHeight="1100" background="#ffffff" math="0" shadow="0">
+            <root>
+                <mxCell id="0"/>
+                <mxCell id="1" parent="0"/>
+                <mxCell id="2" value="Intention module" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#60a917;fontColor=#ffffff;strokeColor=#2D7600;" vertex="1" parent="1">
+                    <mxGeometry x="160" y="220" width="120" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="3" value="knowledge retrieval" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#d80073;fontColor=#ffffff;strokeColor=#A50040;" vertex="1" parent="1">
+                    <mxGeometry x="320" y="220" width="120" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="4" value="LLM generation" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#6a00ff;strokeColor=#3700CC;fontColor=#ffffff;" vertex="1" parent="1">
+                    <mxGeometry x="480" y="220" width="120" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="5" value="Content guardrails" style="rounded=1;whiteSpace=wrap;html=1;fillColor=#a20025;strokeColor=#6F0000;fontColor=#ffffff;" vertex="1" parent="1">
+                    <mxGeometry x="640" y="220" width="120" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="9" value="" style="endArrow=classic;html=1;entryX=0.061;entryY=0.538;entryDx=0;entryDy=0;entryPerimeter=0;strokeColor=#000000;" edge="1" parent="1">
+                    <mxGeometry width="50" height="50" relative="1" as="geometry">
+                        <mxPoint x="60" y="250" as="sourcePoint"/>
+                        <mxPoint x="160.26" y="249.94000000000005" as="targetPoint"/>
+                    </mxGeometry>
+                </mxCell>
+                <mxCell id="10" value="" style="endArrow=classic;html=1;entryX=0.303;entryY=0.538;entryDx=0;entryDy=0;entryPerimeter=0;exitX=1;exitY=0.5;exitDx=0;exitDy=0;strokeColor=#000000;" edge="1" parent="1" source="2">
+                    <mxGeometry width="50" height="50" relative="1" as="geometry">
+                        <mxPoint x="110" y="360" as="sourcePoint"/>
+                        <mxPoint x="319.98" y="249.94000000000005" as="targetPoint"/>
+                    </mxGeometry>
+                </mxCell>
+                <mxCell id="11" value="" style="endArrow=classic;html=1;entryX=0;entryY=0.5;entryDx=0;entryDy=0;exitX=1;exitY=0.5;exitDx=0;exitDy=0;strokeColor=#000000;" edge="1" parent="1" source="3" target="4">
+                    <mxGeometry width="50" height="50" relative="1" as="geometry">
+                        <mxPoint x="290" y="260" as="sourcePoint"/>
+                        <mxPoint x="329.98" y="259.94000000000005" as="targetPoint"/>
+                    </mxGeometry>
+                </mxCell>
+                <mxCell id="12" value="" style="endArrow=classic;html=1;exitX=1;exitY=0.5;exitDx=0;exitDy=0;strokeColor=#000000;" edge="1" parent="1" source="4" target="5">
+                    <mxGeometry width="50" height="50" relative="1" as="geometry">
+                        <mxPoint x="450" y="260" as="sourcePoint"/>
+                        <mxPoint x="490" y="260" as="targetPoint"/>
+                    </mxGeometry>
+                </mxCell>
+                <mxCell id="13" value="" style="endArrow=classic;html=1;exitX=1;exitY=0.5;exitDx=0;exitDy=0;strokeColor=#000000;" edge="1" parent="1" source="5">
+                    <mxGeometry width="50" height="50" relative="1" as="geometry">
+                        <mxPoint x="450" y="260" as="sourcePoint"/>
+                        <mxPoint x="840" y="250" as="targetPoint"/>
+                    </mxGeometry>
+                </mxCell>
+                <mxCell id="14" value="" style="rounded=1;whiteSpace=wrap;html=1;strokeWidth=2;fillColor=none;strokeColor=#000000;" vertex="1" parent="1">
+                    <mxGeometry x="105" y="200" width="680" height="100" as="geometry"/>
+                </mxCell>
+                <mxCell id="17" value="&lt;font color=&quot;#000000&quot;&gt;query&lt;/font&gt;" style="text;html=1;align=center;verticalAlign=middle;resizable=0;points=[];autosize=1;strokeColor=none;fillColor=none;" vertex="1" parent="1">
+                    <mxGeometry x="49" y="218" width="48" height="26" as="geometry"/>
+                </mxCell>
+                <mxCell id="18" value="&lt;font color=&quot;#000000&quot;&gt;response&lt;/font&gt;" style="text;html=1;align=center;verticalAlign=middle;resizable=0;points=[];autosize=1;strokeColor=none;fillColor=none;" vertex="1" parent="1">
+                    <mxGeometry x="786" y="215" width="67" height="26" as="geometry"/>
+                </mxCell>
+            </root>
+        </mxGraphModel>
+    </diagram>
+</mxfile>

source/infrastructure/bin/main.ts

@@ -84,6 +84,7 @@ export class RootStack extends Stack {
       _embeddingEndPoints:_LLMStack._embeddingEndPoints || '',
       _instructEndPoint:_LLMStack._instructEndPoint || '',
       _chatSessionTable: _DynamoDBStack._chatSessionTable,
+      _workspaceTable: _EtlStack._workspaceTableName,
       _sfnOutput: _EtlStack._sfnOutput,
       _OpenSearchIndex: _CdkParameters._OpenSearchIndex.valueAsString,
       _OpenSearchIndexDict: _CdkParameters._OpenSearchIndexDict.valueAsString,
@@ -113,7 +114,7 @@ export class RootStack extends Stack {
     new CfnOutput(this, 'Cross Model Endpoint', {value:_LLMStack._rerankEndPoint || 'No Cross Endpoint Created'});
     new CfnOutput(this, 'Embedding Model Endpoint', {value:_LLMStack._embeddingEndPoints[0] || 'No Embedding Endpoint Created'});
     new CfnOutput(this, 'Instruct Model Endpoint', {value:_LLMStack._instructEndPoint || 'No Instruct Endpoint Created'});
-    new CfnOutput(this, 'Processed Object Table', {value:_EtlStack._processedObjectsTable});
+    new CfnOutput(this, 'Processed Object Table', {value:_EtlStack._processedObjectsTableName});
     new CfnOutput(this, 'Chunk Bucket', {value:_EtlStack._resBucketName});
     new CfnOutput(this, '_aosIndexDict', {value:_CdkParameters._OpenSearchIndexDict.valueAsString});
   }

source/infrastructure/lib/api/api-stack.ts

@@ -26,6 +26,7 @@ interface apiStackProps extends StackProps {
     _embeddingEndPoints: string[];
     _instructEndPoint: string;
     _chatSessionTable: string;
+    _workspaceTable: string;
     // type of StepFunctions
     _sfnOutput: sfn.StateMachine;
     _OpenSearchIndex: string;
@@ -51,6 +52,7 @@ export class LLMApiStack extends NestedStack {
         const _aosIndex = props._OpenSearchIndex
         const _aosIndexDict = props._OpenSearchIndexDict
         const _chatSessionTable = props._chatSessionTable
+        const _workspaceTable = props._workspaceTable
         const _jobQueueArn = props._jobQueueArn
         const _jobDefinitionArn = props._jobDefinitionArn
         const _etlEndpoint = props._etlEndpoint
@@ -291,7 +293,7 @@ export class LLMApiStack extends NestedStack {
         const lambdaDdbIntegration = new apigw.LambdaIntegration(lambdaDdb, { proxy: true, });
 
         // All AOS wrapper should be within such lambda
-        const apiResourceDdb = api.root.addResource('ddb');
+        const apiResourceDdb = api.root.addResource('feedback');
         apiResourceDdb.addMethod('POST', lambdaDdbIntegration);
 
         const apiResourceStepFunction = api.root.addResource('etl');
@@ -323,14 +325,16 @@ export class LLMApiStack extends NestedStack {
                 architecture: Architecture.X86_64,
                 environment: {
                     aos_endpoint: _domainEndpoint,
-                    llm_endpoint: props._instructEndPoint,
+                    llm_model_endpoint_name: props._instructEndPoint,
+                    llm_model_id: "internlm2-chat-7b",
                     embedding_endpoint: props._embeddingEndPoints[0],
                     zh_embedding_endpoint: props._embeddingEndPoints[0],
                     en_embedding_endpoint: props._embeddingEndPoints[1],
                     rerank_endpoint: props._rerankEndPoint,
                     aos_index: _aosIndex,
                     aos_index_dict: _aosIndexDict,
                     chat_session_table: _chatSessionTable,
+                    workspace_table: _workspaceTable,
                 },
                 layers: [_ApiLambdaExecutorLayer]
             });

source/infrastructure/lib/etl/etl-stack.ts

@@ -39,7 +39,8 @@ export class EtlStack extends NestedStack {
     _sfnOutput;
     _jobName;
     _jobArn;
-    _processedObjectsTable;
+    _processedObjectsTableName;
+    _workspaceTableName;
     _etlEndpoint: string;
     _resBucketName: string;
 
@@ -121,6 +122,33 @@ export class EtlStack extends NestedStack {
             // No sort key for this index
         });
 
+        const workspaceTable = new dynamodb.Table(this, "WorkspaceTable", {
+            partitionKey: {
+              name: "workspace_id",
+              type: dynamodb.AttributeType.STRING,
+            },
+            sortKey: {
+              name: "object_type",
+              type: dynamodb.AttributeType.STRING,
+            },
+            billingMode: dynamodb.BillingMode.PAY_PER_REQUEST,
+            encryption: dynamodb.TableEncryption.AWS_MANAGED,
+            pointInTimeRecovery: true,
+            removalPolicy: RemovalPolicy.DESTROY,
+          });
+
+        workspaceTable.addGlobalSecondaryIndex({
+            indexName: "by_object_type_idx",
+            partitionKey: {
+                name: "object_type",
+                type: dynamodb.AttributeType.STRING,
+            },
+            sortKey: {
+                name: "created_at",
+                type: dynamodb.AttributeType.STRING,
+            },
+        });
+
         const _S3Bucket = new s3.Bucket(this, 'llm-bot-glue-res-bucket', {
             // bucketName: `llm-bot-glue-lib-${Aws.ACCOUNT_ID}-${Aws.REGION}`,
             blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL,
@@ -181,19 +209,20 @@ export class EtlStack extends NestedStack {
                 '--QA_ENHANCEMENT.$': sfn.JsonPath.stringAt('$.qaEnhance'),
                 '--AOS_ENDPOINT': props._domainEndpoint,
                 '--REGION': props._region,
-                '--EMBEDDING_MODEL_ENDPOINT': props._embeddingEndpoint.join(','),
+                '--EMBEDDING_MODEL_ENDPOINT': props._embeddingEndpoint[0],
                 '--ETL_MODEL_ENDPOINT': this._etlEndpoint,
                 '--DOC_INDEX_TABLE': props._OpenSearchIndex,
                 '--RES_BUCKET': _S3Bucket.bucketName,
                 '--ProcessedObjectsTable': table.tableName,
-                '--additional-python-modules': 'langchain==0.0.312,beautifulsoup4==4.12.2,requests-aws4auth==1.2.3,boto3==1.28.84,openai==0.28.1,pyOpenSSL==23.3.0,tenacity==8.2.3,markdownify==0.11.6,mammoth==1.6.0,chardet==5.2.0,python-docx==1.1.0,nltk==3.8.1,pdfminer.six==20221105',
+                '--WORKSPACE_TABLE': workspaceTable.tableName,
+                '--WORKSPACE_ID.$': sfn.JsonPath.stringAt('$.workspaceId'),
+                '--additional-python-modules': 'langchain==0.1.0,beautifulsoup4==4.12.2,requests-aws4auth==1.2.3,boto3==1.28.84,openai==0.28.1,pyOpenSSL==23.3.0,tenacity==8.2.3,markdownify==0.11.6,mammoth==1.6.0,chardet==5.2.0,python-docx==1.1.0,nltk==3.8.1,pdfminer.six==20221105',
                 '--python-modules-installer-option': BuildConfig.JOB_PIP_OPTION,
                 // add multiple extra python files
                 '--extra-py-files': extraPythonFilesList,
                 '--CONTENT_TYPE': 'ug',
                 '--EMBEDDING_LANG': 'zh,zh,en,en',
                 '--EMBEDDING_TYPE': 'similarity,relevance,similarity,relevance',
-                '--AOS_INDEX.$': sfn.JsonPath.stringAt('$.aosIndex'),
             }
         });
 
@@ -238,7 +267,7 @@ export class EtlStack extends NestedStack {
                     's3Prefix.$': '$.Payload.s3Prefix',
                     'qaEnhance.$': '$.Payload.qaEnhance',
                     'offline.$': '$.Payload.offline',
-                    'aosIndex.$': '$.Payload.aosIndex',
+                    'workspaceId.$': '$.Payload.workspaceId',
                 }
             },
             // we need the original input
@@ -259,16 +288,16 @@ export class EtlStack extends NestedStack {
                 '--S3_BUCKET.$': '$.s3Bucket',
                 '--S3_PREFIX.$': '$.s3Prefix',
                 '--AOS_ENDPOINT': props._domainEndpoint,
-                '--EMBEDDING_MODEL_ENDPOINT': props._embeddingEndpoint.join(','),
+                '--EMBEDDING_MODEL_ENDPOINT': props._embeddingEndpoint[0],
                 '--ETL_MODEL_ENDPOINT': this._etlEndpoint,
                 '--DOC_INDEX_TABLE': props._OpenSearchIndex,
                 '--REGION': props._region,
                 '--RES_BUCKET': _S3Bucket.bucketName,
                 '--OFFLINE': 'true',
                 '--QA_ENHANCEMENT.$': '$.qaEnhance',
+                '--WORKSPACE_ID.$': '$.workspaceId',
                 // Convert the numeric index to a string
                 '--BATCH_INDICE.$': 'States.Format(\'{}\', $.batchIndices)',
-                '--AOS_INDEX.$': '$.aosIndex',
                 '--ProcessedObjectsTable': table.tableName,
                 '--CONTENT_TYPE': 'ug',
                 '--EMBEDDING_LANG': 'zh,zh,en,en',
@@ -289,7 +318,7 @@ export class EtlStack extends NestedStack {
                 's3Bucket.$': '$.s3Bucket',
                 's3Prefix.$': '$.s3Prefix',
                 'qaEnhance.$': '$.qaEnhance',
-                'aosIndex.$': '$.aosIndex',
+                'workspaceId.$': '$.workspaceId',
                 // 'index' is a special variable within the Map state that represents the current index
                 'batchIndices.$': '$$.Map.Item.Index' // Add this if you need to know the index of the current item in the map state
             },
@@ -308,17 +337,17 @@ export class EtlStack extends NestedStack {
                 '--S3_BUCKET.$': '$.s3Bucket',
                 '--S3_PREFIX.$': '$.s3Prefix',
                 '--AOS_ENDPOINT': props._domainEndpoint,
-                '--EMBEDDING_MODEL_ENDPOINT': props._embeddingEndpoint.join(','),
+                '--EMBEDDING_MODEL_ENDPOINT': props._embeddingEndpoint[0],
                 '--ETL_MODEL_ENDPOINT': this._etlEndpoint,
                 '--DOC_INDEX_TABLE': props._OpenSearchIndex,
                 '--REGION': props._region,
                 '--RES_BUCKET': _S3Bucket.bucketName,
                 '--OFFLINE': 'false',
                 '--QA_ENHANCEMENT.$': '$.qaEnhance',
+                '--WORKSPACE_ID.$': '$.workspaceId',
                 // set the batch indice to 0 since we are running online
                 '--BATCH_INDICE': '0',
                 '--ProcessedObjectsTable': table.tableName,
-                '--AOS_INDEX.$': '$.aosIndex',
             }),
         });
 
@@ -348,7 +377,8 @@ export class EtlStack extends NestedStack {
         this._sfnOutput = sfnStateMachine;
         this._jobName = glueJob.jobName;
         this._jobArn = glueJob.jobArn;
-        this._processedObjectsTable = table.tableName
+        this._processedObjectsTableName = table.tableName;
+        this._workspaceTableName = workspaceTable.tableName;
         this._resBucketName = _S3Bucket.bucketName
     }
 }

source/infrastructure/lib/model/llm-stack.ts

@@ -200,7 +200,7 @@ export class LLMStack extends NestedStack {
                 tags: instruct_tag_array,
             });
 
-            this._instructEndPoint = InstructEndpoint.endpointName as string;
+            this._instructEndPoint = InstructEndpointName;
         }
 
 

source/lambda/ddb/rating.py

@@ -74,6 +74,34 @@ def get_session(table, session_id, user_id):
     return response.get("Item", {})
 
 
+def get_feedback(table, session_id, user_id, message_id):
+    session = get_session(table, session_id, user_id)
+    messages = session.get("History", [])
+    feedback = None
+
+    if not messages:
+        return {
+            "added": False,
+            "error": "Failed to add feedback. No messages found in session.",
+        }
+    elif not message_id:
+        return {
+            "added": False,
+            "error": "Failed to add feedback. Please specify the message_id in the request to add feedback.",
+        }
+
+    for message in messages:
+
+        ddb_message_id = (
+            message.get("data", {}).get("additional_kwargs", {}).get("message_id", "")
+        )
+        if message_id == ddb_message_id:
+            feedback = message["data"]["additional_kwargs"].get("feedback", {})
+            return feedback
+
+    return feedback
+
+
 # SESSIONS_BY_USER_ID_INDEX_NAME = "byUserId"
 def list_sessions_by_user_id(table, user_id, SESSIONS_BY_USER_ID_INDEX_NAME):
     response = {}
@@ -196,13 +224,14 @@ def add_feedback(table, session_id, user_id, message_id, feedback) -> None:
             break
 
     try:
-        table.put_item(
-            Item={
+        table.update_item(
+            Key={
                 "SessionId": session_id,
                 "UserId": user_id,
-                "StartTime": start_time,
-                "History": messages,
-            }
+            },
+            UpdateExpression="SET History = :msg",
+            ExpressionAttributeValues={":msg": messages},
+            ReturnValues="ALL_NEW",
         )
         response = {"added": True}
 
@@ -270,6 +299,9 @@ def lambda_handler(event, context):
     operations_mapping = {
         "POST": {
             "get_session": lambda: get_session(session_table, session_id, user_id),
+            "get_feedback": lambda: get_feedback(
+                session_table, session_id, user_id, message_id
+            ),
             "list_sessions_by_user_id": lambda: list_sessions_by_user_id(
                 session_table, user_id, SESSIONS_BY_USER_ID_INDEX_NAME
             ),