| import boto3 | |||||
| from botocore.config import Config | |||||
| def get_bedrock_client(service_name, credentials=None): | |||||
| client_config = Config(region_name=credentials["aws_region"]) | |||||
| aws_access_key_id = credentials["aws_access_key_id"] | |||||
| aws_secret_access_key = credentials["aws_secret_access_key"] | |||||
| if aws_access_key_id and aws_secret_access_key: | |||||
| # use aksk to call bedrock | |||||
| client = boto3.client( | |||||
| service_name=service_name, | |||||
| config=client_config, | |||||
| aws_access_key_id=aws_access_key_id, | |||||
| aws_secret_access_key=aws_secret_access_key, | |||||
| ) | |||||
| else: | |||||
| # use iam without aksk to call | |||||
| client = boto3.client(service_name=service_name, config=client_config) | |||||
| return client |
| ) | ) | ||||
| from core.model_runtime.errors.validate import CredentialsValidateFailedError | from core.model_runtime.errors.validate import CredentialsValidateFailedError | ||||
| from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel | from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel | ||||
| from core.model_runtime.model_providers.bedrock.get_bedrock_client import get_bedrock_client | |||||
| logger = logging.getLogger(__name__) | logger = logging.getLogger(__name__) | ||||
| ANTHROPIC_BLOCK_MODE_PROMPT = """You should always follow the instructions and output a valid {{block}} object. | ANTHROPIC_BLOCK_MODE_PROMPT = """You should always follow the instructions and output a valid {{block}} object. | ||||
| :param stream: is stream response | :param stream: is stream response | ||||
| :return: full response or stream response chunk generator result | :return: full response or stream response chunk generator result | ||||
| """ | """ | ||||
| bedrock_client = boto3.client( | |||||
| service_name="bedrock-runtime", | |||||
| aws_access_key_id=credentials.get("aws_access_key_id"), | |||||
| aws_secret_access_key=credentials.get("aws_secret_access_key"), | |||||
| region_name=credentials["aws_region"], | |||||
| ) | |||||
| bedrock_client = get_bedrock_client("bedrock-runtime", credentials) | |||||
| system, prompt_message_dicts = self._convert_converse_prompt_messages(prompt_messages) | system, prompt_message_dicts = self._convert_converse_prompt_messages(prompt_messages) | ||||
| inference_config, additional_model_fields = self._convert_converse_api_model_parameters(model_parameters, stop) | inference_config, additional_model_fields = self._convert_converse_api_model_parameters(model_parameters, stop) | ||||
| from typing import Optional | from typing import Optional | ||||
| import boto3 | |||||
| from botocore.config import Config | |||||
| from core.model_runtime.entities.rerank_entities import RerankDocument, RerankResult | from core.model_runtime.entities.rerank_entities import RerankDocument, RerankResult | ||||
| from core.model_runtime.errors.invoke import ( | from core.model_runtime.errors.invoke import ( | ||||
| InvokeAuthorizationError, | InvokeAuthorizationError, | ||||
| ) | ) | ||||
| from core.model_runtime.errors.validate import CredentialsValidateFailedError | from core.model_runtime.errors.validate import CredentialsValidateFailedError | ||||
| from core.model_runtime.model_providers.__base.rerank_model import RerankModel | from core.model_runtime.model_providers.__base.rerank_model import RerankModel | ||||
| from core.model_runtime.model_providers.bedrock.get_bedrock_client import get_bedrock_client | |||||
| class BedrockRerankModel(RerankModel): | class BedrockRerankModel(RerankModel): | ||||
| return RerankResult(model=model, docs=docs) | return RerankResult(model=model, docs=docs) | ||||
| # initialize client | # initialize client | ||||
| client_config = Config(region_name=credentials["aws_region"]) | |||||
| bedrock_runtime = boto3.client( | |||||
| service_name="bedrock-agent-runtime", | |||||
| config=client_config, | |||||
| aws_access_key_id=credentials.get("aws_access_key_id", ""), | |||||
| aws_secret_access_key=credentials.get("aws_secret_access_key"), | |||||
| ) | |||||
| bedrock_runtime = get_bedrock_client("bedrock-agent-runtime", credentials) | |||||
| queries = [{"type": "TEXT", "textQuery": {"text": query}}] | queries = [{"type": "TEXT", "textQuery": {"text": query}}] | ||||
| text_sources = [] | text_sources = [] | ||||
| for text in docs: | for text in docs: |
| import time | import time | ||||
| from typing import Optional | from typing import Optional | ||||
| import boto3 | |||||
| from botocore.config import Config | |||||
| from botocore.exceptions import ( | from botocore.exceptions import ( | ||||
| ClientError, | ClientError, | ||||
| EndpointConnectionError, | EndpointConnectionError, | ||||
| InvokeServerUnavailableError, | InvokeServerUnavailableError, | ||||
| ) | ) | ||||
| from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel | from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel | ||||
| from core.model_runtime.model_providers.bedrock.get_bedrock_client import get_bedrock_client | |||||
| logger = logging.getLogger(__name__) | logger = logging.getLogger(__name__) | ||||
| :param input_type: input type | :param input_type: input type | ||||
| :return: embeddings result | :return: embeddings result | ||||
| """ | """ | ||||
| client_config = Config(region_name=credentials["aws_region"]) | |||||
| bedrock_runtime = boto3.client( | |||||
| service_name="bedrock-runtime", | |||||
| config=client_config, | |||||
| aws_access_key_id=credentials.get("aws_access_key_id"), | |||||
| aws_secret_access_key=credentials.get("aws_secret_access_key"), | |||||
| ) | |||||
| bedrock_runtime = get_bedrock_client("bedrock-runtime", credentials) | |||||
| embeddings = [] | embeddings = [] | ||||
| token_usage = 0 | token_usage = 0 |