From b4dc7f0f17ea2d3f0cc6f8f1d1153f0cb4347ee1 Mon Sep 17 00:00:00 2001
From: Tim Xia <tim.xia@sironamedical.com>
Date: Fri, 1 Mar 2024 23:14:00 -0500
Subject: [PATCH] Add AmazonMistralConfig

---
 litellm/__init__.py     |  1 +
 litellm/llms/bedrock.py | 51 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/litellm/__init__.py b/litellm/__init__.py
index cd639ddb9b..65807a8a08 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -591,6 +591,7 @@ from .llms.bedrock import (
     AmazonCohereConfig,
     AmazonLlamaConfig,
     AmazonStabilityConfig,
+    AmazonMistralConfig
 )
 from .llms.openai import OpenAIConfig, OpenAITextCompletionConfig
 from .llms.azure import AzureOpenAIConfig, AzureOpenAIError
diff --git a/litellm/llms/bedrock.py b/litellm/llms/bedrock.py
index 4806a57e2a..3f14ac9e41 100644
--- a/litellm/llms/bedrock.py
+++ b/litellm/llms/bedrock.py
@@ -282,6 +282,55 @@ class AmazonLlamaConfig:
         }
 
 
+class AmazonMistralConfig:
+    """
+    Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mistral.html
+    Supported Params for the Amazon / Mistral models:
+
+    - `max_tokens` (integer) max tokens,
+    - `temperature` (float) temperature for model,
+    - `top_p` (float) top p for model
+    - `stop` [string] A list of stop sequences that if generated by the model, stops the model from generating further output.
+    - `top_k` (float) top k for model
+    """
+
+    max_tokens: Optional[int] = None
+    temperature: Optional[float] = None
+    topP: Optional[float] = None
+    topK: Optional[float] = None
+    stop: Optional[list[str]] = None
+
+    def __init__(
+        self,
+        max_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+        topP: Optional[int] = None,
+        topK: Optional[float] = None,
+        stop: Optional[list[str]] = None,
+    ) -> None:
+        locals_ = locals()
+        for key, value in locals_.items():
+            if key != "self" and value is not None:
+                setattr(self.__class__, key, value)
+
+    @classmethod
+    def get_config(cls):
+        return {
+            k: v
+            for k, v in cls.__dict__.items()
+            if not k.startswith("__")
+            and not isinstance(
+                v,
+                (
+                    types.FunctionType,
+                    types.BuiltinFunctionType,
+                    classmethod,
+                    staticmethod,
+                ),
+            )
+            and v is not None
+        }
+                
 class AmazonStabilityConfig:
     """
     Reference: https://us-west-2.console.aws.amazon.com/bedrock/home?region=us-west-2#/providers?model=stability.stable-diffusion-xl-v0
@@ -627,7 +676,7 @@ def completion(
             )
         elif provider == "mistral":  
             ## LOAD CONFIG
-            config = litellm.AmazonLlamaConfig.get_config()
+            config = litellm.AmazonMistralConfig.get_config()
             for k, v in config.items():
                 if (
                     k not in inference_params