From 0874885f0447c65695c31589d18ee85ba9d3d0dd Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 2 Sep 2024 11:54:40 -0700
Subject: [PATCH] add all ai21 params

---
 litellm/llms/AI21/chat.py        | 48 +++++++++++++++++---------------
 litellm/tests/test_completion.py |  3 ++
 2 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/litellm/llms/AI21/chat.py b/litellm/llms/AI21/chat.py
index 78bb8f48aed..4eabaaa8756 100644
--- a/litellm/llms/AI21/chat.py
+++ b/litellm/llms/AI21/chat.py
@@ -10,33 +10,36 @@ from typing import Optional, Union
 
 class AI21ChatConfig:
     """
-    Reference: https://inference-docs.cerebras.ai/api-reference/chat-completions
+    Reference: https://docs.ai21.com/reference/jamba-15-api-ref#request-parameters
 
     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
+    response_format: Optional[dict] = None
+    documents: Optional[list] = None
+    max_tokens: Optional[int] = None
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    stop: Optional[Union[str, list]] = None
+    n: Optional[int] = None
+    stream: Optional[bool] = None
+    seed: Optional[int] = None
+    tool_choice: Optional[str] = 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,
+        response_format: Optional[dict] = None,
+        max_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        stop: Optional[Union[str, list]] = None,
+        n: Optional[int] = None,
+        stream: Optional[bool] = None,
+        seed: Optional[int] = None,
+        tool_choice: Optional[str] = None,
         user: Optional[str] = None,
     ) -> None:
         locals_ = locals().copy()
@@ -69,15 +72,16 @@ class AI21ChatConfig:
         """
 
         return [
-            "max_tokens",
+            "tools",
             "response_format",
-            "seed",
-            "stop",
-            "stream",
+            "max_tokens",
             "temperature",
             "top_p",
+            "stop",
+            "n",
+            "stream",
+            "seed",
             "tool_choice",
-            "tools",
             "user",
         ]
 
diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index da6e33a96e6..7aff45ff6be 100644
--- a/litellm/tests/test_completion.py
+++ b/litellm/tests/test_completion.py
@@ -4481,6 +4481,9 @@ async def test_completion_ai21():
     litellm.set_verbose = True
     response = await litellm.acompletion(
         model="ai21_chat/jamba-1.5-large",
+        user="ishaan",
+        tool_choice="auto",
+        seed=123,
         messages=[{"role": "user", "content": "hi my name is ishaan"}],
     )
     pass