klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: bedrock llm, embeddings, tools, temporary creds (#3023)

Browse Source
This commit is contained in:
Laith Al-Saadoon
2025-06-24 10:16:06 -05:00
committed by GitHub
parent b4b27f099e
commit 8139b5887f
2 changed files with 10 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
mem0/embeddings/aws_bedrock.py
View File

@@ -11,7 +11,6 @@ import numpy as np
from mem0.configs.embeddings.base import BaseEmbedderConfig
from mem0.embeddings.base import EmbeddingBase
from mem0.memory.utils import extract_json
class AWSBedrockEmbedding(EmbeddingBase):
@@ -28,6 +27,7 @@ class AWSBedrockEmbedding(EmbeddingBase):
# Get AWS config from environment variables or use defaults
aws_access_key = os.environ.get("AWS_ACCESS_KEY_ID", "")
aws_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY", "")
aws_session_token = os.environ.get("AWS_SESSION_TOKEN", "")
aws_region = os.environ.get("AWS_REGION", "us-west-2")
# Check if AWS config is provided in the config
@@ -43,6 +43,7 @@ class AWSBedrockEmbedding(EmbeddingBase):
region_name=aws_region,
aws_access_key_id=aws_access_key if aws_access_key else None,
aws_secret_access_key=aws_secret_key if aws_secret_key else None,
aws_session_token=aws_session_token if aws_session_token else None,
)
def _normalize_vector(self, embeddings):
@@ -75,7 +76,7 @@ class AWSBedrockEmbedding(EmbeddingBase):
contentType="application/json",
)
response_body = json.loads(extract_json(response.get("body").read()))
response_body = json.loads(response.get("body").read())
if provider == "cohere":
embeddings = response_body.get("embeddings")[0]

19
mem0/llms/aws_bedrock.py
View File

@@ -10,7 +10,7 @@ except ImportError:
from mem0.configs.llms.base import BaseLlmConfig
from mem0.llms.base import LLMBase
from mem0.memory.utils import extract_json
PROVIDERS = ["ai21", "amazon", "anthropic", "cohere", "meta", "mistral", "stability", "writer"]
@@ -92,17 +92,15 @@ class AWSBedrockLLM(LLMBase):
if response["output"]["message"]["content"]:
for item in response["output"]["message"]["content"]:
if "toolUse" in item:
processed_response["tool_calls"].append(
{
"name": item["toolUse"]["name"],
"arguments": item["toolUse"]["input"],
}
)
processed_response["tool_calls"].append({
"name": item["toolUse"]["name"],
"arguments": item["toolUse"]["input"],
})
return processed_response
response_body = response.get("body").read().decode()
response_json = json.loads(extract_json(response_body))
response_json = json.loads(response_body)
return response_json.get("content", [{"text": ""}])[0].get("text", "")
def _prepare_input(
@@ -190,10 +188,7 @@ class AWSBedrockLLM(LLMBase):
}
for prop, details in function["parameters"].get("properties", {}).items():
new_tool["toolSpec"]["inputSchema"]["json"]["properties"][prop] = {
"type": details.get("type", "string"),
"description": details.get("description", ""),
}
new_tool["toolSpec"]["inputSchema"]["json"]["properties"][prop] = details
new_tools.append(new_tool)