From 8c371e12eb6929d817a28f37efcab1559672005e Mon Sep 17 00:00:00 2001 From: zhujiangwei Date: Thu, 6 Jul 2023 23:56:30 +0800 Subject: [PATCH] refactor BedrockEmbeddings class (#7266) #### Description refactor BedrockEmbeddings class to clean code as below: 1. inline content type and accept 2. rewrite input_body as a dictionary literal 3. no need to declare embeddings variable, so remove it --- langchain/embeddings/bedrock.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/langchain/embeddings/bedrock.py b/langchain/embeddings/bedrock.py index 56089ed659..35a7eab5cc 100644 --- a/langchain/embeddings/bedrock.py +++ b/langchain/embeddings/bedrock.py @@ -107,27 +107,21 @@ class BedrockEmbeddings(BaseModel, Embeddings): text = text.replace(os.linesep, " ") _model_kwargs = self.model_kwargs or {} - input_body = {**_model_kwargs} - input_body["inputText"] = text + input_body = {**_model_kwargs, "inputText": text} body = json.dumps(input_body) - content_type = "application/json" - accepts = "application/json" - embeddings = [] try: response = self.client.invoke_model( body=body, modelId=self.model_id, - accept=accepts, - contentType=content_type, + accept="application/json", + contentType="application/json", ) response_body = json.loads(response.get("body").read()) - embeddings = response_body.get("embedding") + return response_body.get("embedding") except Exception as e: raise ValueError(f"Error raised by inference endpoint: {e}") - return embeddings - def embed_documents( self, texts: List[str], chunk_size: int = 1 ) -> List[List[float]]: