Add support for string message in Bedrock textgen (#1291)

* Add support for string message in bedrock, update README
* Add test for string message in test script

Signed-off-by: Jonathan Minkin <[email protected]>

Author: jonminkin97

comps/llms/src/text-generation/README_bedrock.md

@@ -9,6 +9,7 @@
 In order to start Bedrock service, you need to setup the following environment variables first.
 
 ```bash
+export AWS_REGION=${aws_region}
 export AWS_ACCESS_KEY_ID=${aws_access_key_id}
 export AWS_SECRET_ACCESS_KEY=${aws_secret_access_key}
 ```
@@ -23,13 +24,13 @@ export AWS_SESSION_TOKEN=${aws_session_token}
 
 ```bash
 cd GenAIComps/
-docker build --no-cache -t opea/bedrock:latest --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy -f comps/llms/src/text-generation/Dockerfile .
+docker build --no-cache -t opea/llm-textgen:latest --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy -f comps/llms/src/text-generation/Dockerfile .
 ```
 
 ## Run the Bedrock Microservice
 
 ```bash
-docker run -d --name bedrock -p  9009:9000 --ipc=host -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e LLM_COMPONENT_NAME="OpeaTextGenBedrock" -e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID -e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY -e AWS_SESSION_TOKEN=$AWS_SESSION_TOKEN opea/bedrock:latest
+docker run -d --name bedrock -p  9009:9000 --ipc=host -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e LLM_COMPONENT_NAME="OpeaTextGenBedrock" -e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID -e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY -e AWS_SESSION_TOKEN=$AWS_SESSION_TOKEN -e BEDROCK_REGION=$AWS_REGION opea/llm-textgen:latest
 ```
 
 (You can remove `-e AWS_SESSION_TOKEN=$AWS_SESSION_TOKEN` if you are not using an IAM Role)
@@ -42,6 +43,7 @@ curl http://${host_ip}:9009/v1/chat/completions \
   -d '{"model": "us.anthropic.claude-3-5-haiku-20241022-v1:0", "messages": [{"role": "user", "content": "What is Deep Learning?"}], "max_tokens":17}' \
   -H 'Content-Type: application/json'
 
+# stream mode
 curl http://${host_ip}:9009/v1/chat/completions \
  -X POST \
  -d '{"model": "us.anthropic.claude-3-5-haiku-20241022-v1:0", "messages": [{"role": "user", "content": "What is Deep Learning?"}], "max_tokens":17, "stream": "true"}' \

comps/llms/src/text-generation/integrations/bedrock.py

@@ -87,13 +87,21 @@ async def invoke(self, input: ChatCompletionRequest):
         if logflag and len(inference_config) > 0:
             logger.info(f"[llm - chat] inference_config: {inference_config}")
 
-        # Parse messages from HuggingFace TGI format to bedrock messages format
-        # tgi: [{role: "system" | "user", content: "text"}]
+        # Parse messages to Bedrock format
+        # tgi: "prompt" or [{role: "system" | "user", content: "text"}]
         # bedrock: [role: "assistant" | "user", content: {text: "content"}]
-        messages = [
-            {"role": "assistant" if i.get("role") == "system" else "user", "content": [{"text": i.get("content", "")}]}
-            for i in input.messages
-        ]
+        messages = None
+        if isinstance(input.messages, str):
+            messages = [{"role": "user", "content": [{"text": input.messages}]}]
+        else:
+            # Convert from list of HuggingFace TGI message objects
+            messages = [
+                {
+                    "role": "assistant" if i.get("role") == "system" else "user",
+                    "content": [{"text": i.get("content", "")}],
+                }
+                for i in input.messages
+            ]
 
         # Bedrock requires that conversations start with a user prompt
         # TGI allows the first message to be an assistant prompt, defining assistant behavior

tests/llms/test_llms_text-generation_bedrock.sh

@@ -68,10 +68,25 @@ function validate_microservice() {
     -H 'Content-Type: application/json')
 
   if [[ $result == *"data: [DONE]"* ]]; then
-    echo "Result correct."
+    echo "Result correct using chat messages."
     echo "$result" >> ${LOG_PATH}/bedrock.log
   else
-    echo "Result wrong. Received was $result"
+    echo "Result wrong for chat messages. Received was $result"
+    docker logs bedrock-test >> ${LOG_PATH}/bedrock.log
+    exit 1
+  fi
+
+  # Test string message
+  result=$(http_proxy="" curl http://${ip_address}:${bedrock_port}/v1/chat/completions \
+    -X POST \
+    -d '{"model": "us.anthropic.claude-3-haiku-20240307-v1:0", "messages": "What is Deep Learning?", "max_tokens":17, "stream": "true"}' \
+    -H 'Content-Type: application/json')
+
+  if [[ $result == *"data: [DONE]"* ]]; then
+    echo "Result correct using string message."
+    echo "$result" >> ${LOG_PATH}/bedrock.log
+  else
+    echo "Result wrong for string message. Received was $result"
     docker logs bedrock-test >> ${LOG_PATH}/bedrock.log
     exit 1
   fi