diff --git a/Containerfile b/Containerfile index 38b60b6a..1c22ebf5 100644 --- a/Containerfile +++ b/Containerfile @@ -13,7 +13,7 @@ RUN dnf install -y python3 python3-pip python3-wheel python3-aiohttp \ RUN pip3 install torch --index-url https://download.pytorch.org/whl/cpu -RUN pip3 install anthropic cohere google-cloud-aiplatform langchain langchain-core \ +RUN pip3 install anthropic boto3 cohere google-cloud-aiplatform langchain langchain-core \ langchain-huggingface langchain-text-splitters langchain-community \ pymilvus sentence-transformers transformers huggingface-hub \ pulsar-client && \ diff --git a/docker-compose-bedrock.yaml b/docker-compose-bedrock.yaml index 518c6577..b748b6e4 100644 --- a/docker-compose-bedrock.yaml +++ b/docker-compose-bedrock.yaml @@ -209,8 +209,10 @@ services: - "text-completion-bedrock" - "-p" - "pulsar://pulsar:6650" - # - "-m" - # - "mistral.mistral-large-2407-v1:0" + - "-z" + - "{AWS_ID_KEY}" + - "-k" + - "{AWS_SECRET_KEY}" restart: on-failure:100 text-completion-rag: @@ -221,6 +223,10 @@ services: - "pulsar://pulsar:6650" # - "-m" # - "mistral.mistral-large-2407-v1:0" + - "-z" + - "{AWS_ID_KEY}" + - "-k" + - "{AWS_SECRET_KEY}" - "-i" - "non-persistent://tg/request/text-completion-rag" - "-o" diff --git a/requirements.txt b/requirements.txt index a0c184ca..9a49a5aa 100644 --- a/requirements.txt +++ b/requirements.txt @@ -19,3 +19,4 @@ google-cloud-aiplatform pyyaml prometheus-client pyarrow +boto3 diff --git a/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph/model/text_completion/bedrock/llm.py index eab9c9f4..04c4c9fd 100755 --- a/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph/model/text_completion/bedrock/llm.py @@ -28,7 +28,8 @@ class Processor(ConsumerProducer): output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) model = params.get("model", default_model) - api_key = params.get("api_key") + aws_id = params.get("aws_id_key") + aws_secret = params.get("aws_secret") super(Processor, self).__init__( **params | { @@ -43,7 +44,13 @@ class Processor(ConsumerProducer): self.model = model - self.bedrock = boto3.client(service_name='bedrock-runtime', region_name="us-west-2") + 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' + ) + + self.bedrock = self.session.client(service_name='bedrock-runtime') print("Initialised", flush=True) @@ -115,6 +122,16 @@ class Processor(ConsumerProducer): help=f'Bedrock model (default: Mistral-Large-2407)' ) + parser.add_argument( + '-z', '--aws-id-key', + help=f'AWS ID Key' + ) + + parser.add_argument( + '-k', '--aws-secret', + help=f'AWS Secret Key' + ) + def run(): Processor.start(module, __doc__)