chat with document updated to support history and a default model

Author: ilmarmur

src/api/schema.graphql

@@ -203,7 +203,7 @@ type Query @aws_cognito_user_pools @aws_iam {
   getFileContents(s3Uri: String!): FileContentsResponse
   getConfiguration: ConfigurationResponse
   queryKnowledgeBase(input: String!, sessionId: String): String
-  chatWithDocument(s3Uri: String!, prompt: String!, modelId: String!): String
+  chatWithDocument(s3Uri: String!, prompt: String!, history: AWSJSON!, modelId: String!): String
   getStepFunctionExecution(executionArn: String!): StepFunctionExecutionResponse
   submitAnalyticsQuery(query: String!): AnalyticsJob @aws_cognito_user_pools
   getAnalyticsJobStatus(jobId: ID!): AnalyticsJob @aws_cognito_user_pools

src/lambda/chat_with_document_resolver/index.py

@@ -47,11 +47,18 @@ def handler(event, context):
 
         objectKey = event['arguments']['s3Uri']
         prompt = event['arguments']['prompt']
+        history = event['arguments']['history']
+
+        full_prompt = "You are an assistant that's responsible for getting details from document text attached here based on questions from the user.\n\n"
+        full_prompt += "If you don't know the answer, just say that you don't know. Don't try to make up an answer.\n\n"
+        full_prompt += "Additionally, use the user and assistant responses in the following JSON object to see what's been asked and what the resposes were in the past.\n\n"
+        full_prompt += "The JSON object is: " + json.dumps(history) + ".\n\n"
+        full_prompt += "The user's question is: " + prompt
 
         # this feature is not enabled until the model can be selected on the chat screen
         # selectedModelId = event['arguments']['modelId']
         selectedModelId = get_summarization_model()
-        
+
         logger.info(f"Processing S3 URI: {objectKey}")
 
         output_bucket = os.environ['OUTPUT_BUCKET']
@@ -67,6 +74,9 @@ def handler(event, context):
             # full text key
             fulltext_key = objectKey + '/summary/fulltext.txt'
 
+            logger.info(f"Output Bucket: {output_bucket}")
+            logger.info(f"Full Text Key: {fulltext_key}")
+
             # read full contents of the object as text
             s3 = boto3.client('s3')
             response = s3.get_object(Bucket=output_bucket, Key=fulltext_key)
@@ -90,7 +100,7 @@ def handler(event, context):
                     "role":"user",
                     "content": [
                         {
-                            "text": prompt
+                            "text": full_prompt
                         }
                     ]
                 }

src/ui/package-lock.json

@@ -10,13 +10,17 @@ import './ChatPanel.css';
 
 const logger = new Logger('chatWithDocument');
 
-const getChatResponse = async (s3Uri, prompt) => {
+const getChatResponse = async (s3Uri, prompt, history) => {
   logger.debug('s3URI:', s3Uri);
-
+  logger.debug('history:', history);
   // commenting this out until model selection for chat is available again on this screen
   // logger.debug('modelId:', modelId);
   const modelId = 'us.amazon.nova-pro-v1:0';
-  const response = await API.graphql({ query: chatWithDocument, variables: { s3Uri, prompt, modelId } });
+  const strHistory = JSON.stringify(history);
+  const response = await API.graphql({
+    query: chatWithDocument,
+    variables: { s3Uri, prompt, history: strHistory, modelId },
+  });
   // logger.debug('response:', response);
   return response;
 };
@@ -35,6 +39,7 @@ const ChatPanel = (item) => {
   const [error, setError] = useState(null);
   // const [modelId, setModelId] = useState(modelOptions[0].value);
   const [chatQueries, setChatQueries] = useState([]);
+  const [jsonChatHistory, setJsonChatHistory] = useState([]);
   const textareaRef = useRef(null);
   const { objectKey } = item;
   let rowId = 0;
@@ -71,8 +76,8 @@ const ChatPanel = (item) => {
     textareaRef.current.value = '';
 
     // comment out sending the model ID until model selection is available again on this screen
-    // const chatResponse = getChatResponse(objectKey, prompt, modelId);
-    const chatResponse = getChatResponse(objectKey, prompt);
+    // const chatResponse = getChatResponse(objectKey, prompt, history, modelId);
+    const chatResponse = getChatResponse(objectKey, prompt, jsonChatHistory);
 
     let chatResponseData = {};
 
@@ -87,6 +92,13 @@ const ChatPanel = (item) => {
             dt: new Date().toLocaleTimeString(),
             type: 'msg',
           };
+
+          const chatItem = {
+            ask: prompt,
+            response: cResponse.cr.content[0].text,
+          };
+
+          setJsonChatHistory((prevChatHistory) => [...prevChatHistory, chatItem]);
         }
       })
       .catch((r) => {
@@ -158,9 +170,30 @@ const ChatPanel = (item) => {
             <p>To start chatting to this document, enter your message below.</p>
           )}
 
-          <FormField label="Your message" className="chat-composer-container">
+          {/* <FormField label="Your message" className="chat-composer-container">
             <textarea name="postContent" ref={textareaRef} rows={6} className="chat-textarea" id="chatTextarea" />
-          </FormField>
+          </FormField> */}
+
+          <div style={{ gap: '8px', width: '100%' }}>
+            <FormField label="Your message" style={{ flex: 8 }}>
+              <input
+                type="text"
+                name="postContent"
+                ref={textareaRef}
+                style={{ padding: '3px', width: '100%' }}
+                id="chatTextarea"
+                onKeyDown={(e) => {
+                  if (e.key === 'Enter') {
+                    handlePromptSubmit();
+                  }
+                }}
+              />
+            </FormField>
+
+            <Button variant="primary" onClick={handlePromptSubmit}>
+              Send
+            </Button>
+          </div>
 
           {/* <FormField label="Model">
             <select name="model" id="modelSelect" onChange={handleModelIdChange}>
@@ -171,10 +204,6 @@ const ChatPanel = (item) => {
               ))}
             </select>
           </FormField> */}
-
-          <Button variant="primary" onClick={handlePromptSubmit}>
-            Send
-          </Button>
         </Container>
       </SpaceBetween>
     </div>

src/ui/src/components/chat-panel/ChatPanel.jsx

@@ -3,7 +3,7 @@
 import gql from 'graphql-tag';
 
 export default gql`
-  query Query($s3Uri: String!, $prompt: String!, $modelId: String!) {
-    chatWithDocument(s3Uri: $s3Uri, prompt: $prompt, modelId: $modelId)
+  query Query($s3Uri: String!, $prompt: String!, $history: AWSJSON!, $modelId: String!) {
+    chatWithDocument(s3Uri: $s3Uri, prompt: $prompt, history: $history, modelId: $modelId)
   }
 `;

src/ui/src/graphql/queries/chatWithDocument.js

@@ -5725,11 +5725,14 @@ Resources:
         Variables:
           LOG_LEVEL: !Ref LogLevel
           OUTPUT_BUCKET: !Ref OutputBucket
+          CONFIGURATION_TABLE_NAME: !Ref ConfigurationTable
       LoggingConfig:
         LogGroup: !Ref ChatWithDocumentResolverFunctionLogGroup
       Policies:
         - S3ReadPolicy:
             BucketName: !Ref OutputBucket
+        - DynamoDBCrudPolicy:
+            TableName: !Ref ConfigurationTable
         - Statement: 
           - !If
             - ShouldUseDocumentKnowledgeBase