From 838484b1f3ec5fb2faeb0547f1877cfd014a2ec8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pavel=20Kr=C3=A1l?= <pavel@pavelkral.eu>
Date: Wed, 26 Jul 2023 19:59:09 +0200
Subject: [PATCH] Set min/max constraints to float arguments

* Add .idea/ to .gitignore
* Set min and max constraints to float parameters

Closes #115
---
 .gitignore                           |  1 +
 llm/default_plugins/openai_models.py |  8 ++++++++
 tests/test_cli_openai_models.py      | 25 +++++++++++++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/.gitignore b/.gitignore
index 53605b7a..a8925435 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,3 +7,4 @@ venv
 .pytest_cache
 *.egg-info
 .DS_Store
+.idea/
diff --git a/llm/default_plugins/openai_models.py b/llm/default_plugins/openai_models.py
index 01724132..f331c228 100644
--- a/llm/default_plugins/openai_models.py
+++ b/llm/default_plugins/openai_models.py
@@ -94,6 +94,8 @@ class Options(llm.Options):
                 "0.8 will make the output more random, while lower values like 0.2 will "
                 "make it more focused and deterministic."
             ),
+            ge=0,
+            le=2,
             default=None,
         )
         max_tokens: Optional[int] = Field(
@@ -107,6 +109,8 @@ class Options(llm.Options):
                 "10% probability mass are considered. Recommended to use top_p or "
                 "temperature but not both."
             ),
+            ge=0,
+            le=1,
             default=None,
         )
         frequency_penalty: Optional[float] = Field(
@@ -115,6 +119,8 @@ class Options(llm.Options):
                 "on their existing frequency in the text so far, decreasing the model's "
                 "likelihood to repeat the same line verbatim."
             ),
+            ge=-2,
+            le=2,
             default=None,
         )
         presence_penalty: Optional[float] = Field(
@@ -123,6 +129,8 @@ class Options(llm.Options):
                 "on whether they appear in the text so far, increasing the model's "
                 "likelihood to talk about new topics."
             ),
+            ge=-2,
+            le=2,
             default=None,
         )
         stop: Optional[str] = Field(
diff --git a/tests/test_cli_openai_models.py b/tests/test_cli_openai_models.py
index fd5ff07a..6af71073 100644
--- a/tests/test_cli_openai_models.py
+++ b/tests/test_cli_openai_models.py
@@ -37,3 +37,28 @@ def test_openai_models(mocked_models):
         "ada:2020-05-03        openai      2020-05-03T20:26:40\n"
         "babbage:2020-05-03    openai      2020-05-03T20:26:40\n"
     )
+
+
+def test_openai_options_min_max(mocked_models):
+    options = {
+        "temperature": [0, 2],
+        "top_p": [0, 1],
+        "frequency_penalty": [-2, 2],
+        "presence_penalty": [-2, 2],
+    }
+    runner = CliRunner()
+
+    for option, [min_val, max_val] in options.items():
+        result = runner.invoke(cli, ["-m", "chatgpt", "-o", option, "-10"])
+        assert result.exit_code == 1
+        assert (
+            result.output
+            == f"Error: {option}\n  Input should be greater than or equal to {min_val}\n"
+        )
+
+        result = runner.invoke(cli, ["-m", "chatgpt", "-o", option, "10"])
+        assert result.exit_code == 1
+        assert (
+            result.output
+            == f"Error: {option}\n  Input should be less than or equal to {max_val}\n"
+        )