From af0d2d1f3b9a7f8d8fdd5be2e673edb0599feeeb Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 2 Sep 2024 11:47:04 -0700
Subject: [PATCH] refactor ai21

---
 litellm/llms/AI21/chat.py                    | 91 ++++++++++++++++++++
 litellm/llms/{ai21.py => AI21/completion.py} |  0
 2 files changed, 91 insertions(+)
 create mode 100644 litellm/llms/AI21/chat.py
 rename litellm/llms/{ai21.py => AI21/completion.py} (100%)

diff --git a/litellm/llms/AI21/chat.py b/litellm/llms/AI21/chat.py
new file mode 100644
index 0000000000..78bb8f48ae
--- /dev/null
+++ b/litellm/llms/AI21/chat.py
@@ -0,0 +1,91 @@
+"""
+AI21 Chat Completions API
+
+this is OpenAI compatible - no translation needed / occurs
+"""
+
+import types
+from typing import Optional, Union
+
+
+class AI21ChatConfig:
+    """
+    Reference: https://inference-docs.cerebras.ai/api-reference/chat-completions
+
+    Below are the parameters:
+    """
+
+    max_tokens: Optional[int] = None
+    response_format: Optional[dict] = None
+    seed: Optional[int] = None
+    stop: Optional[str] = None
+    stream: Optional[bool] = None
+    temperature: Optional[float] = None
+    top_p: Optional[int] = None
+    tool_choice: Optional[str] = None
+    tools: Optional[list] = None
+    user: Optional[str] = None
+
+    def __init__(
+        self,
+        max_tokens: Optional[int] = None,
+        response_format: Optional[dict] = None,
+        seed: Optional[int] = None,
+        stop: Optional[str] = None,
+        stream: Optional[bool] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[int] = None,
+        tool_choice: Optional[str] = None,
+        tools: Optional[list] = None,
+        user: Optional[str] = None,
+    ) -> None:
+        locals_ = locals().copy()
+        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
+        }
+
+    def get_supported_openai_params(self, model: str) -> list:
+        """
+        Get the supported OpenAI params for the given model
+
+        """
+
+        return [
+            "max_tokens",
+            "response_format",
+            "seed",
+            "stop",
+            "stream",
+            "temperature",
+            "top_p",
+            "tool_choice",
+            "tools",
+            "user",
+        ]
+
+    def map_openai_params(
+        self, model: str, non_default_params: dict, optional_params: dict
+    ) -> dict:
+        supported_openai_params = self.get_supported_openai_params(model=model)
+        for param, value in non_default_params.items():
+            if param in supported_openai_params:
+                optional_params[param] = value
+        return optional_params
diff --git a/litellm/llms/ai21.py b/litellm/llms/AI21/completion.py
similarity index 100%
rename from litellm/llms/ai21.py
rename to litellm/llms/AI21/completion.py