lint

aurelio-labs · Jan 13, 2024 · 2285771 · 2285771
1 parent f16529f
commit 2285771
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/tests/unit/llms/test_llm_base.py b/tests/unit/llms/test_llm_base.py
@@ -41,7 +41,7 @@ def test_base_llm_is_valid_inputs_valid_input_fail(self, base_llm):
     def test_base_llm_is_valid_inputs_invalid_false(self, base_llm):
         test_schema = {
             "name": "get_time",
-            "description": 'Finds the current time in a specific timezone.\n\n:param timezone: The timezone to find the current time in, should\n    be a valid timezone from the IANA Time Zone Database like\n    "America/New_York" or "Europe/London". Do NOT put the place\n    name itself like "rome", or "new york", you must provide\n    the IANA format.\n:type timezone: str\n:return: The current time in the specified timezone.'
+            "description": 'Finds the current time in a specific timezone.\n\n:param timezone: The timezone to find the current time in, should\n    be a valid timezone from the IANA Time Zone Database like\n    "America/New_York" or "Europe/London". Do NOT put the place\n    name itself like "rome", or "new york", you must provide\n    the IANA format.\n:type timezone: str\n:return: The current time in the specified timezone.',
         }
         test_inputs = {"timezone": "America/New_York"}
 

diff --git a/tests/unit/llms/test_llm_llamacpp.py b/tests/unit/llms/test_llm_llamacpp.py
@@ -33,7 +33,11 @@ def test_llamacpp_llm_grammar(self, llamacpp_llm):
 
     def test_llamacpp_extract_function_inputs(self, llamacpp_llm, mocker):
         llamacpp_llm.llm.create_chat_completion = mocker.Mock(
-            return_value={"choices": [{"message": {"content": "{'timezone': 'America/New_York'}"}}]}
+            return_value={
+                "choices": [
+                    {"message": {"content": "{'timezone': 'America/New_York'}"}}
+                ]
+            }
         )
         test_schema = {
             "name": "get_time",
@@ -50,7 +54,11 @@ def test_llamacpp_extract_function_inputs(self, llamacpp_llm, mocker):
     def test_llamacpp_extract_function_inputs_invalid(self, llamacpp_llm, mocker):
         with pytest.raises(ValueError):
             llamacpp_llm.llm.create_chat_completion = mocker.Mock(
-                return_value={"choices": [{"message": {"content": "{'time': 'America/New_York'}"}}]}
+                return_value={
+                    "choices": [
+                        {"message": {"content": "{'time': 'America/New_York'}"}}
+                    ]
+                }
             )
             test_schema = {
                 "name": "get_time",
@@ -62,4 +70,4 @@ def test_llamacpp_extract_function_inputs_invalid(self, llamacpp_llm, mocker):
 
             llamacpp_llm.extract_function_inputs(
                 query=test_query, function_schema=test_schema
-            )
+            )