Harrison/cohere experimental (#638)

Co-authored-by: inyourhead <[email protected]>

Author: hwchase17

langchain/llms/ai21.py

@@ -64,6 +64,9 @@ class AI21(LLM, BaseModel):
 
     ai21_api_key: Optional[str] = None
 
+    base_url: Optional[str] = None
+    """Base url to use, if None decides based on model name."""
+
     class Config:
         """Configuration for this pydantic object."""
 
@@ -118,8 +121,15 @@ def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
         """
         if stop is None:
             stop = []
+        if self.base_url is not None:
+            base_url = self.base_url
+        else:
+            if self.model in ("j1-grande-instruct",):
+                base_url = "https://api.ai21.com/studio/v1/experimental"
+            else:
+                base_url = "https://api.ai21.com/studio/v1"
         response = requests.post(
-            url=f"https://api.ai21.com/studio/v1/{self.model}/complete",
+            url=f"{base_url}/{self.model}/complete",
             headers={"Authorization": f"Bearer {self.ai21_api_key}"},
             json={"prompt": prompt, "stopSequences": stop, **self._default_params},
         )

tests/integration_tests/llms/test_ai21.py

@@ -13,6 +13,13 @@ def test_ai21_call() -> None:
     assert isinstance(output, str)
 
 
+def test_ai21_call_experimental() -> None:
+    """Test valid call to ai21 with an experimental model."""
+    llm = AI21(maxTokens=10, model="j1-grande-instruct")
+    output = llm("Say foo:")
+    assert isinstance(output, str)
+
+
 def test_saving_loading_llm(tmp_path: Path) -> None:
     """Test saving/loading an AI21 LLM."""
     llm = AI21(maxTokens=10)