Harrison/cohere experimental (#638)

Co-authored-by: inyourhead <44607279+xettrisomeman@users.noreply.github.com>

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)