fix(qa): read properly var to fix streaming (#341)

awslabs · Mar 26, 2024 · e58f614 · e58f614
1 parent d9394b5
commit e58f614
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 19 deletions.
diff --git a/lambda/aws-qa-appsync-opensearch/question_answering/src/adapters/base/base.py b/lambda/aws-qa-appsync-opensearch/question_answering/src/adapters/base/base.py
@@ -16,9 +16,10 @@
 
 
 class ModelAdapter:
-    def __init__(self, callback=None, modality='Text', model_kwargs={}):
+    def __init__(self, callback=None, modality='Text', model_kwargs={}, streaming=False):
         self.model_kwargs = model_kwargs
         self.modality = modality
+        self.streaming = streaming
 
         self.callback_handler = callback
 

diff --git a/lambda/aws-qa-appsync-opensearch/question_answering/src/adapters/bedrock/claude.py b/lambda/aws-qa-appsync-opensearch/question_answering/src/adapters/bedrock/claude.py
@@ -50,9 +50,9 @@ def get_llm(self, model_kwargs={}):
             "streaming": False 
         }
 
-        if self.callback_handler:
+        if self.callback_handler and self.streaming:
             kwargs["callbacks"] = self.callback_handler
-            kwargs["streaming"] = model_kwargs.get("streaming", False)
+            kwargs["streaming"] = True
 
         return Bedrock(
             **kwargs
@@ -104,9 +104,9 @@ def get_llm(self, model_kwargs={}):
             "streaming": False 
         }
 
-        if self.callback_handler:
+        if self.callback_handler and self.streaming:
             kwargs["callbacks"] = self.callback_handler
-            kwargs["streaming"] = model_kwargs.get("streaming", False)
+            kwargs["streaming"] = True
 
         return BedrockChat(
             **kwargs

diff --git a/lambda/aws-qa-appsync-opensearch/question_answering/src/adapters/bedrock/titan.py b/lambda/aws-qa-appsync-opensearch/question_answering/src/adapters/bedrock/titan.py
@@ -46,9 +46,9 @@ def get_llm(self, model_kwargs={}):
             "streaming": False 
         }
 
-        if self.callback_handler:
+        if self.callback_handler and self.streaming:
             kwargs["callbacks"] = self.callback_handler
-            kwargs["streaming"] = model_kwargs.get("streaming", False)
+            kwargs["streaming"] = True
 
         return Bedrock(
             **kwargs

diff --git a/lambda/aws-qa-appsync-opensearch/question_answering/src/qa_agent/doc_qa.py b/lambda/aws-qa-appsync-opensearch/question_answering/src/qa_agent/doc_qa.py
@@ -81,6 +81,7 @@ def run_qa_agent_rag_no_memory(input_params):
         callback=callback_manager,
         modality=qa_modality,
         model_kwargs=qa_model_args,
+        streaming=streaming
     )
 
     # get embeddings model
@@ -104,6 +105,7 @@ def run_qa_agent_rag_no_memory(input_params):
         model_id=em_model_id,
         modality=qa_modality,
         model_kwargs=em_model_args,
+        streaming=streaming
     )
 
     logger.info(decoded_question)
@@ -231,6 +233,7 @@ def run_qa_agent_from_single_document_no_memory(input_params):
         callback=callback_manager,
         modality=qa_modality,
         model_kwargs=qa_model_args,
+        streaming=streaming
     )
 
     # 1 : load the document

diff --git a/src/patterns/gen-ai/aws-qa-appsync-opensearch/README.md b/src/patterns/gen-ai/aws-qa-appsync-opensearch/README.md
@@ -189,7 +189,7 @@ Mutation call to trigger the question:
     jobid: "123",
     jobstatus: "", 
     qa_model: 
-   	 {
+      {
       provider: "Bedrock",
       modality: "Text",
       modelId: "anthropic.claude-v2:1", 
@@ -299,23 +299,23 @@ Parameters
 
 Question answering
 
-| **Provider**  | **Model id** | **Modalities** | **Notes** |
-| :----------------------------- | :---- | ----- | --- |
-| Bedrock                          |  anthropic.claude-v2 | Text |  |
-| Bedrock                          |  anthropic.claude-v2:1 | Text |  |
-| Bedrock                          |  anthropic.claude-3-haiku-20240307-v1:0 | Text, Image |  |
-| Bedrock                          |  anthropic.claude-3-sonnet-20240229-v1:0 | Text, Image |  |
-| Bedrock                          |  anthropic.claude-instant-v1 | Text |  |
-| Bedrock                          |  amazon.titan-text-lite-v1 | Text |  |
-| Bedrock                          |  amazon.titan-text-express-v1 | Text |  |
-| SageMaker                          |  idefics | Text, Image | The model is not deployed as part of the construct and requires to be provisioned separately |
+| **Provider**  | **Model id** | **Modalities** | **Streaming** | **Notes** |
+| :----------------------------- | :---- | ----- | --- | --- |
+| Bedrock                          |  anthropic.claude-v2 | Text | ✅ | |
+| Bedrock                          |  anthropic.claude-v2:1 | Text | ✅ | Default model is none selected |
+| Bedrock                          |  anthropic.claude-3-haiku-20240307-v1:0 | Text, Image | ✅ | |
+| Bedrock                          |  anthropic.claude-3-sonnet-20240229-v1:0 | Text, Image | ✅ | |
+| Bedrock                          |  anthropic.claude-instant-v1 | Text | ✅ | |
+| Bedrock                          |  amazon.titan-text-lite-v1 | Text | ✅  |
+| Bedrock                          |  amazon.titan-text-express-v1 | Text | ✅  |
+| SageMaker                          |  idefics | Text, Image | ❌ | The model is not deployed as part of the construct and requires to be provisioned separately | 
 
 Embeddings
 
 | **Provider**  | **Model id** | **Modalities** | **Notes** |
 | :----------------------------- | :---- | ----- | --- |
 | Bedrock                          |  amazon.titan-embed-image-v1 | Text |  |
-| Bedrock                          |  amazon.titan-embed-text-v1 | Text |  |
+| Bedrock                          |  amazon.titan-embed-text-v1 | Text | Default model is none selected  |
 
 ## Default properties