JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(bedrock): omit topP for Anthropic Converse; use AWSBedrockConfig in LlmFactory (#4469)

Browse Source
This commit is contained in:
Himanshu
2026-03-25 11:22:52 +05:30
committed by GitHub
parent d1b4b304c7
commit 2e0f91e70d
4 changed files with 383 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mem0/llms/aws_bedrock.py
+43 -41
View File
@@ -228,6 +228,12 @@ class AWSBedrockLLM(LLMBase):
return "\n\nHuman: " + "".join(formatted_messages) + "\n\nAssistant:"
def _merge_optional_top_p(self, target: Dict[str, Any], *, key: str = "top_p") -> None:
"""Add nucleus sampling to ``target`` only when ``model_config`` has ``top_p`` set."""
top_p = self.model_config.get("top_p")
if top_p is not None:
target[key] = top_p
def _prepare_input(self, prompt: str) -> Dict[str, Any]:
"""
Prepare input for the current provider's model.
@@ -268,44 +274,38 @@ class AWSBedrockLLM(LLMBase):
"messages": [{"role": "user", "content": prompt}],
"max_tokens": self.model_config.get("max_tokens", 5000),
"temperature": self.model_config.get("temperature", 0.1),
"top_p": self.model_config.get("top_p", 0.9),
}
self._merge_optional_top_p(input_body)
else:
# Legacy Amazon models
input_body = {
"inputText": prompt,
"textGenerationConfig": {
"maxTokenCount": self.model_config.get("max_tokens", 5000),
"topP": self.model_config.get("top_p", 0.9),
"temperature": self.model_config.get("temperature", 0.1),
},
}
# Remove None values
input_body["textGenerationConfig"] = {
k: v for k, v in input_body["textGenerationConfig"].items() if v is not None
text_gen_config: Dict[str, Any] = {
"maxTokenCount": self.model_config.get("max_tokens", 5000),
"temperature": self.model_config.get("temperature", 0.1),
}
self._merge_optional_top_p(text_gen_config, key="topP")
input_body = {"inputText": prompt, "textGenerationConfig": text_gen_config}
elif self.provider == "anthropic":
input_body = {
"messages": [{"role": "user", "content": [{"type": "text", "text": prompt}]}],
"max_tokens": self.model_config.get("max_tokens", 2000),
"temperature": self.model_config.get("temperature", 0.1),
"top_p": self.model_config.get("top_p", 0.9),
"anthropic_version": "bedrock-2023-05-31",
}
self._merge_optional_top_p(input_body)
elif self.provider == "meta":
input_body = {
"prompt": prompt,
"max_gen_len": self.model_config.get("max_tokens", 5000),
"temperature": self.model_config.get("temperature", 0.1),
"top_p": self.model_config.get("top_p", 0.9),
}
self._merge_optional_top_p(input_body)
elif self.provider == "mistral":
input_body = {
"prompt": prompt,
"max_tokens": self.model_config.get("max_tokens", 5000),
"temperature": self.model_config.get("temperature", 0.1),
"top_p": self.model_config.get("top_p", 0.9),
}
self._merge_optional_top_p(input_body)
else:
# Generic case - add all model config parameters
input_body.update(self.model_config)
@@ -479,6 +479,29 @@ class AWSBedrockLLM(LLMBase):
return converse_tools
def _default_max_tokens_for_converse(self) -> int:
"""Default maxTokens if ``max_tokens`` is missing (Nova: 5000, else 2000)."""
model_id = (self.config.model or "").lower()
if self.provider == "amazon" and "nova" in model_id:
return 5000
return 2000
def _build_inference_config(self) -> Dict[str, Any]:
"""Build Converse ``inferenceConfig``. Anthropic allows only one of temperature or topP; we keep temperature and omit topP."""
inference_config: Dict[str, Any] = {
"maxTokens": self.model_config.get("max_tokens", self._default_max_tokens_for_converse()),
"temperature": self.model_config.get("temperature", 0.1),
}
top_p = self.model_config.get("top_p")
if top_p is not None:
if self.provider == "anthropic":
logger.debug("Omitting topP for Anthropic Converse (using temperature); top_p=%s", top_p)
else:
inference_config["topP"] = top_p
return inference_config
def _generate_with_tools(self, messages: List[Dict[str, str]], tools: List[Dict], stream: bool = False) -> Dict[str, Any]:
"""Generate response with tool calling support using correct message format."""
# Format messages for tool-enabled models
@@ -501,11 +524,7 @@ class AWSBedrockLLM(LLMBase):
converse_params = {
"modelId": self.config.model,
"messages": formatted_messages,
"inferenceConfig": {
"maxTokens": self.model_config.get("max_tokens", 2000),
"temperature": self.model_config.get("temperature", 0.1),
"topP": self.model_config.get("top_p", 0.9),
}
"inferenceConfig": self._build_inference_config(),
}
# Add system message if present (for Anthropic)
@@ -531,11 +550,7 @@ class AWSBedrockLLM(LLMBase):
converse_params = {
"modelId": self.config.model,
"messages": formatted_messages,
"inferenceConfig": {
"maxTokens": self.model_config.get("max_tokens", 2000),
"temperature": self.model_config.get("temperature", 0.1),
"topP": self.model_config.get("top_p", 0.9),
}
"inferenceConfig": self._build_inference_config(),
}
# Add system message if present
@@ -554,26 +569,13 @@ class AWSBedrockLLM(LLMBase):
return str(response)
elif self.provider == "amazon" and "nova" in self.config.model.lower():
# Nova models use converse API even without tools
# Nova models use the Converse API even without tools
formatted_messages = self._format_messages_amazon(messages)
input_body = {
"messages": formatted_messages,
"max_tokens": self.model_config.get("max_tokens", 5000),
"temperature": self.model_config.get("temperature", 0.1),
"top_p": self.model_config.get("top_p", 0.9),
}
# Use converse API for Nova models
response = self.client.converse(
modelId=self.config.model,
messages=input_body["messages"],
inferenceConfig={
"maxTokens": input_body["max_tokens"],
"temperature": input_body["temperature"],
"topP": input_body["top_p"],
}
messages=formatted_messages,
inferenceConfig=self._build_inference_config(),
)
return self._parse_response(response)
else:
# For other providers and legacy Amazon models (like Titan)