/mem0/llms/aws_bedrock.py

import json
import logging
import re
from typing import Any, Dict, List, Optional, Union

try:
    import boto3
    from botocore.exceptions import ClientError, NoCredentialsError
except ImportError:
    raise ImportError("The 'boto3' library is required. Please install it using 'pip install boto3'.")

from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.aws_bedrock import AWSBedrockConfig
from mem0.llms.base import LLMBase
from mem0.memory.utils import extract_json

logger = logging.getLogger(__name__)

PROVIDERS = [
    "ai21", "amazon", "anthropic", "cohere", "meta", "mistral", "stability", "writer",
    "deepseek", "gpt-oss", "perplexity", "snowflake", "titan", "command", "j2", "llama",
    "minimax",
]


def extract_provider(model: str) -> str:
    """Extract provider from model identifier."""
    for provider in PROVIDERS:
        if re.search(rf"\b{re.escape(provider)}\b", model):
            return provider
    raise ValueError(f"Unknown provider in model: {model}")


class AWSBedrockLLM(LLMBase):
    """
    AWS Bedrock LLM integration for Mem0.

    Supports all available Bedrock models with automatic provider detection.
    """

    def __init__(self, config: Optional[Union[AWSBedrockConfig, BaseLlmConfig, Dict]] = None):
        """
        Initialize AWS Bedrock LLM.

        Args:
            config: AWS Bedrock configuration object
        """
        # Convert to AWSBedrockConfig if needed
        if config is None:
            config = AWSBedrockConfig()
        elif isinstance(config, dict):
            config = AWSBedrockConfig(**config)
        elif isinstance(config, BaseLlmConfig) and not isinstance(config, AWSBedrockConfig):
            # Convert BaseLlmConfig to AWSBedrockConfig
            config = AWSBedrockConfig(
                model=config.model,
                temperature=config.temperature,
                max_tokens=config.max_tokens,
                top_p=config.top_p,
                top_k=config.top_k,
                enable_vision=getattr(config, "enable_vision", False),
            )

        super().__init__(config)
        self.config = config

        # Initialize AWS client
        self._initialize_aws_client()

        # Get model configuration
        self.model_config = self.config.get_model_config()
        self.provider = extract_provider(self.config.model)

        # Initialize provider-specific settings
        self._initialize_provider_settings()

    def _initialize_aws_client(self):
        """Initialize AWS Bedrock client with proper credentials."""
        try:
            aws_config = self.config.get_aws_config()

            # Create Bedrock runtime client
            self.client = boto3.client("bedrock-runtime", **aws_config)

            # Test connection
            self._test_connection()

        except NoCredentialsError:
            raise ValueError(
                "AWS credentials not found. Please set AWS_ACCESS_KEY_ID, "
                "AWS_SECRET_ACCESS_KEY, and AWS_REGION environment variables, "
                "or provide them in the config."
            )
        except ClientError as e:
            if e.response["Error"]["Code"] == "UnauthorizedOperation":
                raise ValueError(
                    f"Unauthorized access to Bedrock. Please ensure your AWS credentials "
                    f"have permission to access Bedrock in region {self.config.aws_region}."
                )
            else:
                raise ValueError(f"AWS Bedrock error: {e}")

    def _test_connection(self):
        """Test connection to AWS Bedrock service."""
        try:
            # List available models to test connection
            bedrock_client = boto3.client("bedrock", **self.config.get_aws_config())
            response = bedrock_client.list_foundation_models()
            self.available_models = [model["modelId"] for model in response["modelSummaries"]]

            # Check if our model is available
            if self.config.model not in self.available_models:
                logger.warning(f"Model {self.config.model} may not be available in region {self.config.aws_region}")
                logger.info(f"Available models: {', '.join(self.available_models[:5])}...")

        except Exception as e:
            logger.warning(f"Could not verify model availability: {e}")
            self.available_models = []

    def _initialize_provider_settings(self):
        """Initialize provider-specific settings and capabilities."""
        # Determine capabilities based on provider and model
        self.supports_tools = self.provider in ["anthropic", "cohere", "amazon"]
        # MiniMax M2.x is intentionally excluded from supports_tools: tool use for MiniMax
        # on Amazon Bedrock is only available via the bedrock-mantle (OpenAI-compatible)
        # endpoint, not via the bedrock-runtime Converse API used by this class.
        # See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-minimax-minimax-m2-5.html
        self.supports_vision = self.provider in ["anthropic", "amazon", "meta", "mistral"]
        self.supports_streaming = self.provider in ["anthropic", "cohere", "mistral", "amazon", "meta"]

        # Set message formatting method
        if self.provider == "anthropic":
            self._format_messages = self._format_messages_anthropic
        elif self.provider == "cohere":
            self._format_messages = self._format_messages_cohere
        elif self.provider == "amazon":
            self._format_messages = self._format_messages_amazon
        elif self.provider == "meta":
            self._format_messages = self._format_messages_meta
        elif self.provider == "mistral":
            self._format_messages = self._format_messages_mistral
        else:
            self._format_messages = self._format_messages_generic

    def _format_messages_anthropic(self, messages: List[Dict[str, str]]) -> tuple[List[Dict[str, Any]], Optional[str]]:
        """Format messages for Anthropic models."""
        formatted_messages = []
        system_message = None

        for message in messages:
            role = message["role"]
            content = message["content"]

            if role == "system":
                # Anthropic supports system messages as a separate parameter
                # see: https://docs.anthropic.com/en/docs/build-with-claude/prompt-engineering/system-prompts
                system_message = content
            elif role == "user":
                # Use Converse API format
                formatted_messages.append({"role": "user", "content": [{"text": content}]})
            elif role == "assistant":
                # Use Converse API format
                formatted_messages.append({"role": "assistant", "content": [{"text": content}]})

        return formatted_messages, system_message

    def _format_messages_cohere(self, messages: List[Dict[str, str]]) -> str:
        """Format messages for Cohere models."""
        formatted_messages = []

        for message in messages:
            role = message["role"].capitalize()
            content = message["content"]
            formatted_messages.append(f"{role}: {content}")

        return "\n".join(formatted_messages)

    def _format_messages_amazon(self, messages: List[Dict[str, str]]) -> List[Dict[str, Any]]:
        """Format messages for Amazon models (including Nova)."""
        formatted_messages = []
        
        for message in messages:
            role = message["role"]
            content = message["content"]
            
            if role == "system":
                # Amazon models support system messages
                formatted_messages.append({"role": "system", "content": content})
            elif role == "user":
                formatted_messages.append({"role": "user", "content": content})
            elif role == "assistant":
                formatted_messages.append({"role": "assistant", "content": content})
        
        return formatted_messages

    def _format_messages_meta(self, messages: List[Dict[str, str]]) -> str:
        """Format messages for Meta models."""
        formatted_messages = []
        
        for message in messages:
            role = message["role"].capitalize()
            content = message["content"]
            formatted_messages.append(f"{role}: {content}")
        
        return "\n".join(formatted_messages)

    def _format_messages_mistral(self, messages: List[Dict[str, str]]) -> List[Dict[str, Any]]:
        """Format messages for Mistral models."""
        formatted_messages = []
        
        for message in messages:
            role = message["role"]
            content = message["content"]
            
            if role == "system":
                # Mistral supports system messages
                formatted_messages.append({"role": "system", "content": content})
            elif role == "user":
                formatted_messages.append({"role": "user", "content": content})
            elif role == "assistant":
                formatted_messages.append({"role": "assistant", "content": content})
        
        return formatted_messages

    def _format_messages_generic(self, messages: List[Dict[str, str]]) -> str:
        """Generic message formatting for other providers."""
        formatted_messages = []

        for message in messages:
            role = message["role"].capitalize()
            content = message["content"]
            formatted_messages.append(f"\n\n{role}: {content}")

        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.

        Args:
            prompt: Text prompt to process

        Returns:
            Prepared input dictionary
        """
        # Base configuration
        input_body = {"prompt": prompt}

        # Provider-specific parameter mappings
        provider_mappings = {
            "meta": {"max_tokens": "max_gen_len"},
            "ai21": {"max_tokens": "maxTokens", "top_p": "topP"},
            "mistral": {"max_tokens": "max_tokens"},
            "cohere": {"max_tokens": "max_tokens", "top_p": "p"},
            "amazon": {"max_tokens": "maxTokenCount", "top_p": "topP"},
            "anthropic": {"max_tokens": "max_tokens", "top_p": "top_p"},
        }

        # Apply provider mappings
        if self.provider in provider_mappings:
            for old_key, new_key in provider_mappings[self.provider].items():
                if old_key in self.model_config:
                    input_body[new_key] = self.model_config[old_key]

        # Special handling for specific providers
        if self.provider == "cohere" and "cohere.command" in self.config.model:
            input_body["message"] = input_body.pop("prompt")
        elif self.provider == "amazon":
            # Amazon Nova and other Amazon models
            if "nova" in self.config.model.lower():
                # Nova models use the converse API format
                input_body = {
                    "messages": [{"role": "user", "content": prompt}],
                    "max_tokens": self.model_config.get("max_tokens", 5000),
                    "temperature": self.model_config.get("temperature", 0.1),
                }
                self._merge_optional_top_p(input_body)
            else:
                # Legacy Amazon models
                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),
                "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),
            }
            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),
            }
            self._merge_optional_top_p(input_body)
        else:
            # Generic case - add all model config parameters
            input_body.update(self.model_config)

        return input_body

    def _convert_tool_format(self, original_tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        Convert tools to Bedrock-compatible format.

        Args:
            original_tools: List of tool definitions

        Returns:
            Converted tools in Bedrock format
        """
        new_tools = []

        for tool in original_tools:
            if tool["type"] == "function":
                function = tool["function"]
                new_tool = {
                    "toolSpec": {
                        "name": function["name"],
                        "description": function.get("description", ""),
                        "inputSchema": {
                            "json": {
                                "type": "object",
                                "properties": {},
                                "required": function["parameters"].get("required", []),
                            }
                        },
                    }
                }

                # Add properties
                for prop, details in function["parameters"].get("properties", {}).items():
                    new_tool["toolSpec"]["inputSchema"]["json"]["properties"][prop] = details

                new_tools.append(new_tool)

        return new_tools

    def _parse_response(
        self, response: Dict[str, Any], tools: Optional[List[Dict]] = None
    ) -> Union[str, Dict[str, Any]]:
        """
        Parse response from Bedrock API.

        Args:
            response: Raw API response
            tools: List of tools if used

        Returns:
            Parsed response
        """
        if tools:
            # Handle tool-enabled responses
            processed_response = {"tool_calls": []}

            if response.get("output", {}).get("message", {}).get("content"):
                for item in response["output"]["message"]["content"]:
                    if "toolUse" in item:
                        processed_response["tool_calls"].append(
                            {
                                "name": item["toolUse"]["name"],
                                "arguments": json.loads(extract_json(json.dumps(item["toolUse"]["input"]))),
                            }
                        )

            return processed_response

        # Handle regular text responses
        try:
            response_body = response.get("body").read().decode()
            response_json = json.loads(response_body)

            # Provider-specific response parsing
            if self.provider == "anthropic":
                return response_json.get("content", [{"text": ""}])[0].get("text", "")
            elif self.provider == "amazon":
                # Handle both Nova and legacy Amazon models
                if "nova" in self.config.model.lower():
                    # Nova models return content in a different format
                    if "content" in response_json:
                        return response_json["content"][0]["text"]
                    elif "completion" in response_json:
                        return response_json["completion"]
                else:
                    # Legacy Amazon models
                    return response_json.get("completion", "")
            elif self.provider == "meta":
                return response_json.get("generation", "")
            elif self.provider == "mistral":
                return response_json.get("outputs", [{"text": ""}])[0].get("text", "")
            elif self.provider == "cohere":
                return response_json.get("generations", [{"text": ""}])[0].get("text", "")
            elif self.provider == "ai21":
                return response_json.get("completions", [{"data": {"text": ""}}])[0].get("data", {}).get("text", "")
            else:
                # Generic parsing - try common response fields
                for field in ["content", "text", "completion", "generation"]:
                    if field in response_json:
                        if isinstance(response_json[field], list) and response_json[field]:
                            return response_json[field][0].get("text", "")
                        elif isinstance(response_json[field], str):
                            return response_json[field]

                # Fallback
                return str(response_json)

        except Exception as e:
            logger.warning(f"Could not parse response: {e}")
            return "Error parsing response"

    def generate_response(
        self,
        messages: List[Dict[str, str]],
        response_format: Optional[str] = None,
        tools: Optional[List[Dict]] = None,
        tool_choice: str = "auto",
        stream: bool = False,
        **kwargs,
    ) -> Union[str, Dict[str, Any]]:
        """
        Generate response using AWS Bedrock.

        Args:
            messages: List of message dictionaries
            response_format: Response format specification
            tools: List of tools for function calling
            tool_choice: Tool choice method
            stream: Whether to stream the response
            **kwargs: Additional parameters

        Returns:
            Generated response
        """
        try:
            if tools and self.supports_tools:
                # Use converse method for tool-enabled models
                return self._generate_with_tools(messages, tools, stream)
            else:
                # Use standard invoke_model method
                return self._generate_standard(messages, stream)

        except Exception as e:
            logger.error(f"Failed to generate response: {e}")
            raise RuntimeError(f"Failed to generate response: {e}")

    @staticmethod
    def _convert_tools_to_converse_format(tools: List[Dict]) -> List[Dict]:
        """Convert OpenAI-style tools to Converse API format."""
        if not tools:
            return []

        converse_tools = []
        for tool in tools:
            if tool.get("type") == "function" and "function" in tool:
                func = tool["function"]
                converse_tool = {
                    "toolSpec": {
                        "name": func["name"],
                        "description": func.get("description", ""),
                        "inputSchema": {
                            "json": func.get("parameters", {})
                        }
                    }
                }
                converse_tools.append(converse_tool)

        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 and MiniMax reasoning models reject requests that include both
        ``temperature`` and ``topP`` simultaneously, so ``topP`` is omitted for
        those providers even when the user has configured it.
        """
        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 in ("anthropic", "minimax"):
                # Both Anthropic and MiniMax M2.x (reasoning models) raise a
                # ValidationException when temperature and topP are both present
                # in inferenceConfig.  Omit topP and rely on temperature only.
                logger.debug("Omitting topP for %s Converse (using temperature); top_p=%s", self.provider, 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
        system_message = None
        if self.provider == "anthropic":
            formatted_messages, system_message = self._format_messages_anthropic(messages)
        elif self.provider == "amazon":
            formatted_messages = self._format_messages_amazon(messages)
        else:
            formatted_messages = [{"role": "user", "content": [{"text": messages[-1]["content"]}]}]

        # Prepare tool configuration in Converse API format
        tool_config = None
        if tools:
            converse_tools = self._convert_tools_to_converse_format(tools)
            if converse_tools:
                tool_config = {"tools": converse_tools}

        # Prepare converse parameters
        converse_params = {
            "modelId": self.config.model,
            "messages": formatted_messages,
            "inferenceConfig": self._build_inference_config(),
        }

        # Add system message if present (for Anthropic)
        if system_message:
            converse_params["system"] = [{"text": system_message}]

        # Add tool config if present
        if tool_config:
            converse_params["toolConfig"] = tool_config

        # Make API call
        response = self.client.converse(**converse_params)

        return self._parse_response(response, tools)

    def _generate_standard(self, messages: List[Dict[str, str]], stream: bool = False) -> str:
        """Generate standard text response using Converse API for Anthropic models."""
        # For Anthropic models, always use Converse API
        if self.provider == "anthropic":
            formatted_messages, system_message = self._format_messages_anthropic(messages)

            # Prepare converse parameters
            converse_params = {
                "modelId": self.config.model,
                "messages": formatted_messages,
                "inferenceConfig": self._build_inference_config(),
            }

            # Add system message if present
            if system_message:
                converse_params["system"] = [{"text": system_message}]

            # Use converse API for Anthropic models
            response = self.client.converse(**converse_params)

            # Parse Converse API response
            if hasattr(response, 'output') and hasattr(response.output, 'message'):
                return response.output.message.content[0].text
            elif 'output' in response and 'message' in response['output']:
                return response['output']['message']['content'][0]['text']
            else:
                return str(response)

        elif self.provider == "minimax":
            # MiniMax models (e.g. minimax.minimax-m2.5) use the Bedrock Converse API.
            # M2.5 is a reasoning model whose response content array may include a
            # `reasoningContent` block before the actual `text` block, so we iterate
            # to find the first block that contains a "text" key.
            # System messages must be passed via the top-level `system` parameter
            # (not as a message with role="system") per the Converse API spec.
            system_parts = []
            converse_messages = []
            for msg in messages:
                role = msg.get("role", "user")
                content = msg.get("content", "")
                if not isinstance(content, str):
                    content = str(content)
                if role == "system":
                    system_parts.append(content)
                else:
                    converse_messages.append({"role": role, "content": [{"text": content}]})
            if not converse_messages:
                converse_messages = [{"role": "user", "content": [{"text": ""}]}]
            converse_params = {
                "modelId": self.config.model,
                "messages": converse_messages,
                "inferenceConfig": self._build_inference_config(),
            }
            if system_parts:
                converse_params["system"] = [{"text": "\n".join(system_parts)}]
            response = self.client.converse(**converse_params)
            for block in response["output"]["message"]["content"]:
                if "text" in block:
                    return block["text"]
            return ""

        elif self.provider == "amazon" and "nova" in self.config.model.lower():
            # Nova models use the Converse API even without tools
            formatted_messages = self._format_messages_amazon(messages)
            response = self.client.converse(
                modelId=self.config.model,
                messages=formatted_messages,
                inferenceConfig=self._build_inference_config(),
            )
            return self._parse_response(response)
        else:
            # For other providers and legacy Amazon models (like Titan)
            if self.provider == "amazon":
                # Legacy Amazon models need string formatting, not array formatting
                prompt = self._format_messages_generic(messages)
            else:
                prompt = self._format_messages(messages)
            input_body = self._prepare_input(prompt)

            # Convert to JSON
            body = json.dumps(input_body)

            # Make API call
            response = self.client.invoke_model(
                body=body,
                modelId=self.config.model,
                accept="application/json",
                contentType="application/json",
            )

            return self._parse_response(response)

    def list_available_models(self) -> List[Dict[str, Any]]:
        """List all available models in the current region."""
        try:
            bedrock_client = boto3.client("bedrock", **self.config.get_aws_config())
            response = bedrock_client.list_foundation_models()

            models = []
            for model in response["modelSummaries"]:
                provider = extract_provider(model["modelId"])
                models.append(
                    {
                        "model_id": model["modelId"],
                        "provider": provider,
                        "model_name": model["modelId"].split(".", 1)[1]
                        if "." in model["modelId"]
                        else model["modelId"],
                        "modelArn": model.get("modelArn", ""),
                        "providerName": model.get("providerName", ""),
                        "inputModalities": model.get("inputModalities", []),
                        "outputModalities": model.get("outputModalities", []),
                        "responseStreamingSupported": model.get("responseStreamingSupported", False),
                    }
                )

            return models

        except Exception as e:
            logger.warning(f"Could not list models: {e}")
            return []

    def get_model_capabilities(self) -> Dict[str, Any]:
        """Get capabilities of the current model."""
        return {
            "model_id": self.config.model,
            "provider": self.provider,
            "model_name": self.config.model_name,
            "supports_tools": self.supports_tools,
            "supports_vision": self.supports_vision,
            "supports_streaming": self.supports_streaming,
            "max_tokens": self.model_config.get("max_tokens", 2000),
        }

    def validate_model_access(self) -> bool:
        """Validate if the model is accessible."""
        try:
            # Try to invoke the model with a minimal request
            if self.provider == "amazon" and "nova" in self.config.model.lower():
                # Test Nova model with converse API
                test_messages = [{"role": "user", "content": "test"}]
                self.client.converse(
                    modelId=self.config.model,
                    messages=test_messages,
                    inferenceConfig={"maxTokens": 10}
                )
            else:
                # Test other models with invoke_model
                test_body = json.dumps({"prompt": "test"})
                self.client.invoke_model(
                    body=test_body,
                    modelId=self.config.model,
                    accept="application/json",
                    contentType="application/json",
                )
            return True
        except Exception:
            return False