diff --git a/docker-compose-bedrock.yaml b/docker-compose-bedrock.yaml index b748b6e4..00150a93 100644 --- a/docker-compose-bedrock.yaml +++ b/docker-compose-bedrock.yaml @@ -210,9 +210,11 @@ services: - "-p" - "pulsar://pulsar:6650" - "-z" - - "{AWS_ID_KEY}" + - "${AWS_ID_KEY}" - "-k" - - "{AWS_SECRET_KEY}" + - "${AWS_SECRET_KEY}" + - "-r" + - "us-west-2" restart: on-failure:100 text-completion-rag: @@ -224,9 +226,11 @@ services: # - "-m" # - "mistral.mistral-large-2407-v1:0" - "-z" - - "{AWS_ID_KEY}" + - "${AWS_ID_KEY}" - "-k" - - "{AWS_SECRET_KEY}" + - "${AWS_SECRET_KEY}" + - "-r" + - "us-west-2" - "-i" - "non-persistent://tg/request/text-completion-rag" - "-o" diff --git a/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph/model/text_completion/bedrock/llm.py index 04c4c9fd..21ff73d8 100755 --- a/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph/model/text_completion/bedrock/llm.py @@ -6,6 +6,7 @@ Input is prompt, output is response. Mistral is default. import boto3 import json +import re from .... schema import TextCompletionRequest, TextCompletionResponse from .... schema import text_completion_request_queue @@ -19,6 +20,7 @@ default_input_queue = text_completion_request_queue default_output_queue = text_completion_response_queue default_subscriber = module default_model = 'mistral.mistral-large-2407-v1:0' +default_region = 'us-west-2' class Processor(ConsumerProducer): @@ -30,6 +32,7 @@ class Processor(ConsumerProducer): model = params.get("model", default_model) aws_id = params.get("aws_id_key") aws_secret = params.get("aws_secret") + aws_region = params.get("aws_region", default_region) super(Processor, self).__init__( **params | { @@ -47,7 +50,7 @@ class Processor(ConsumerProducer): self.session = boto3.Session( aws_access_key_id=aws_id, aws_secret_access_key=aws_secret, - region_name='us-west-2' # e.g., 'us-west-2' + region_name=aws_region ) self.bedrock = self.session.client(service_name='bedrock-runtime') @@ -99,11 +102,23 @@ class Processor(ConsumerProducer): response_body = json.loads(response.get("body").read()) outputtext = response_body['outputs'][0]['text'] - resp = outputtext - print(resp, flush=True) + print(outputtext, flush=True) + + # Parse output for ```json``` delimiters + pattern = r'```json\s*([\s\S]*?)\s*```' + match = re.search(pattern, outputtext) + + if match: + # If delimiters are found, extract the JSON content + json_content = match.group(1) + json_resp = json_content.strip() + + else: + # If no delimiters are found, return the original text + json_resp = outputtext.strip() print("Send response...", flush=True) - r = TextCompletionResponse(response=resp) + r = TextCompletionResponse(response=json_resp) self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -132,6 +147,11 @@ class Processor(ConsumerProducer): help=f'AWS Secret Key' ) + parser.add_argument( + '-r', '--aws-region', + help=f'AWS Region (default: us-west-2)' + ) + def run(): Processor.start(module, __doc__) diff --git a/trustgraph/model/text_completion/cohere/llm.py b/trustgraph/model/text_completion/cohere/llm.py index c5d5ec27..5d8ba57c 100755 --- a/trustgraph/model/text_completion/cohere/llm.py +++ b/trustgraph/model/text_completion/cohere/llm.py @@ -62,7 +62,7 @@ class Processor(ConsumerProducer): output = self.cohere.chat( model=self.model, message=prompt, - preamble = "You are an AI-assistant chatbot. You are trained to read text and find entities in that text. You respond only with well-formed JSON.", + preamble = "You are a helpful AI-assistant." temperature=0.0, chat_history=[], prompt_truncation='auto',