|
@@ -59,15 +59,15 @@ class BedrockTextEmbeddingModel(TextEmbeddingModel):
|
|
|
model_prefix = model.split('.')[0]
|
|
|
|
|
|
if model_prefix == "amazon" :
|
|
|
- for text in texts:
|
|
|
- body = {
|
|
|
+ for text in texts:
|
|
|
+ body = {
|
|
|
"inputText": text,
|
|
|
- }
|
|
|
- response_body = self._invoke_bedrock_embedding(model, bedrock_runtime, body)
|
|
|
- embeddings.extend([response_body.get('embedding')])
|
|
|
- token_usage += response_body.get('inputTextTokenCount')
|
|
|
- logger.warning(f'Total Tokens: {token_usage}')
|
|
|
- result = TextEmbeddingResult(
|
|
|
+ }
|
|
|
+ response_body = self._invoke_bedrock_embedding(model, bedrock_runtime, body)
|
|
|
+ embeddings.extend([response_body.get('embedding')])
|
|
|
+ token_usage += response_body.get('inputTextTokenCount')
|
|
|
+ logger.warning(f'Total Tokens: {token_usage}')
|
|
|
+ result = TextEmbeddingResult(
|
|
|
model=model,
|
|
|
embeddings=embeddings,
|
|
|
usage=self._calc_response_usage(
|
|
@@ -75,20 +75,20 @@ class BedrockTextEmbeddingModel(TextEmbeddingModel):
|
|
|
credentials=credentials,
|
|
|
tokens=token_usage
|
|
|
)
|
|
|
- )
|
|
|
- return result
|
|
|
-
|
|
|
+ )
|
|
|
+ return result
|
|
|
+
|
|
|
if model_prefix == "cohere" :
|
|
|
- input_type = 'search_document' if len(texts) > 1 else 'search_query'
|
|
|
- for text in texts:
|
|
|
- body = {
|
|
|
+ input_type = 'search_document' if len(texts) > 1 else 'search_query'
|
|
|
+ for text in texts:
|
|
|
+ body = {
|
|
|
"texts": [text],
|
|
|
"input_type": input_type,
|
|
|
- }
|
|
|
- response_body = self._invoke_bedrock_embedding(model, bedrock_runtime, body)
|
|
|
- embeddings.extend(response_body.get('embeddings'))
|
|
|
- token_usage += len(text)
|
|
|
- result = TextEmbeddingResult(
|
|
|
+ }
|
|
|
+ response_body = self._invoke_bedrock_embedding(model, bedrock_runtime, body)
|
|
|
+ embeddings.extend(response_body.get('embeddings'))
|
|
|
+ token_usage += len(text)
|
|
|
+ result = TextEmbeddingResult(
|
|
|
model=model,
|
|
|
embeddings=embeddings,
|
|
|
usage=self._calc_response_usage(
|
|
@@ -96,9 +96,9 @@ class BedrockTextEmbeddingModel(TextEmbeddingModel):
|
|
|
credentials=credentials,
|
|
|
tokens=token_usage
|
|
|
)
|
|
|
- )
|
|
|
- return result
|
|
|
-
|
|
|
+ )
|
|
|
+ return result
|
|
|
+
|
|
|
#others
|
|
|
raise ValueError(f"Got unknown model prefix {model_prefix} when handling block response")
|
|
|
|
|
@@ -183,7 +183,7 @@ class BedrockTextEmbeddingModel(TextEmbeddingModel):
|
|
|
)
|
|
|
|
|
|
return usage
|
|
|
-
|
|
|
+
|
|
|
def _map_client_to_invoke_error(self, error_code: str, error_msg: str) -> type[InvokeError]:
|
|
|
"""
|
|
|
Map client error to invoke error
|
|
@@ -212,9 +212,9 @@ class BedrockTextEmbeddingModel(TextEmbeddingModel):
|
|
|
content_type = 'application/json'
|
|
|
try:
|
|
|
response = bedrock_runtime.invoke_model(
|
|
|
- body=json.dumps(body),
|
|
|
- modelId=model,
|
|
|
- accept=accept,
|
|
|
+ body=json.dumps(body),
|
|
|
+ modelId=model,
|
|
|
+ accept=accept,
|
|
|
contentType=content_type
|
|
|
)
|
|
|
response_body = json.loads(response.get('body').read().decode('utf-8'))
|