From b3c0728de2893b924f34b75f8107430024154669 Mon Sep 17 00:00:00 2001 From: olgavrou Date: Tue, 29 Aug 2023 05:28:43 -0400 Subject: [PATCH] fix mypy errors in tests --- .../rl_chain/test_pick_best_chain_call.py | 34 ++++++++------- .../rl_chain/test_pick_best_text_embedder.py | 34 ++++++++------- .../rl_chain/test_rl_chain_base_embedder.py | 42 +++++++++---------- 3 files changed, 58 insertions(+), 52 deletions(-) diff --git a/libs/langchain/tests/unit_tests/chains/rl_chain/test_pick_best_chain_call.py b/libs/langchain/tests/unit_tests/chains/rl_chain/test_pick_best_chain_call.py index 3fad1667d91..7bca6b470d8 100644 --- a/libs/langchain/tests/unit_tests/chains/rl_chain/test_pick_best_chain_call.py +++ b/libs/langchain/tests/unit_tests/chains/rl_chain/test_pick_best_chain_call.py @@ -1,3 +1,5 @@ +from typing import Any, Dict + import pytest from test_utils import MockEncoder @@ -10,7 +12,7 @@ encoded_text = "[ e n c o d e d ] " @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def setup(): +def setup() -> tuple: _PROMPT_TEMPLATE = """This is a dummy prompt that will be ignored by the fake llm""" PROMPT = PromptTemplate(input_variables=[], template=_PROMPT_TEMPLATE) @@ -19,7 +21,7 @@ def setup(): @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def test_multiple_ToSelectFrom_throws(): +def test_multiple_ToSelectFrom_throws() -> None: llm, PROMPT = setup() chain = pick_best_chain.PickBest.from_llm(llm=llm, prompt=PROMPT) actions = ["0", "1", "2"] @@ -32,7 +34,7 @@ def test_multiple_ToSelectFrom_throws(): @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def test_missing_basedOn_from_throws(): +def test_missing_basedOn_from_throws() -> None: llm, PROMPT = setup() chain = pick_best_chain.PickBest.from_llm(llm=llm, prompt=PROMPT) actions = ["0", "1", "2"] @@ -41,7 +43,7 @@ def test_missing_basedOn_from_throws(): @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def test_ToSelectFrom_not_a_list_throws(): +def test_ToSelectFrom_not_a_list_throws() -> None: llm, PROMPT = setup() chain = pick_best_chain.PickBest.from_llm(llm=llm, prompt=PROMPT) actions = {"actions": ["0", "1", "2"]} @@ -53,7 +55,7 @@ def test_ToSelectFrom_not_a_list_throws(): @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def test_update_with_delayed_score_with_auto_validator_throws(): +def test_update_with_delayed_score_with_auto_validator_throws() -> None: llm, PROMPT = setup() # this LLM returns a number so that the auto validator will return that auto_val_llm = FakeListChatModel(responses=["3"]) @@ -75,7 +77,7 @@ def test_update_with_delayed_score_with_auto_validator_throws(): @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def test_update_with_delayed_score_force(): +def test_update_with_delayed_score_force() -> None: llm, PROMPT = setup() # this LLM returns a number so that the auto validator will return that auto_val_llm = FakeListChatModel(responses=["3"]) @@ -99,7 +101,7 @@ def test_update_with_delayed_score_force(): @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def test_update_with_delayed_score(): +def test_update_with_delayed_score() -> None: llm, PROMPT = setup() chain = pick_best_chain.PickBest.from_llm( llm=llm, prompt=PROMPT, selection_scorer=None @@ -117,11 +119,11 @@ def test_update_with_delayed_score(): @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def test_user_defined_scorer(): +def test_user_defined_scorer() -> None: llm, PROMPT = setup() class CustomSelectionScorer(rl_chain.SelectionScorer): - def score_response(self, inputs, llm_response: str) -> float: + def score_response(self, inputs: Dict[str, Any], llm_response: str) -> float: score = 200 return score @@ -139,7 +141,7 @@ def test_user_defined_scorer(): @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def test_default_embeddings(): +def test_default_embeddings() -> None: llm, PROMPT = setup() feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) chain = pick_best_chain.PickBest.from_llm( @@ -173,7 +175,7 @@ def test_default_embeddings(): @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def test_default_embeddings_off(): +def test_default_embeddings_off() -> None: llm, PROMPT = setup() feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) chain = pick_best_chain.PickBest.from_llm( @@ -199,7 +201,7 @@ def test_default_embeddings_off(): @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def test_default_embeddings_mixed_w_explicit_user_embeddings(): +def test_default_embeddings_mixed_w_explicit_user_embeddings() -> None: llm, PROMPT = setup() feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) chain = pick_best_chain.PickBest.from_llm( @@ -234,7 +236,7 @@ def test_default_embeddings_mixed_w_explicit_user_embeddings(): @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def test_default_no_scorer_specified(): +def test_default_no_scorer_specified() -> None: _, PROMPT = setup() chain_llm = FakeListChatModel(responses=[100]) chain = pick_best_chain.PickBest.from_llm(llm=chain_llm, prompt=PROMPT) @@ -249,7 +251,7 @@ def test_default_no_scorer_specified(): @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def test_explicitly_no_scorer(): +def test_explicitly_no_scorer() -> None: llm, PROMPT = setup() chain = pick_best_chain.PickBest.from_llm( llm=llm, prompt=PROMPT, selection_scorer=None @@ -265,7 +267,7 @@ def test_explicitly_no_scorer(): @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def test_auto_scorer_with_user_defined_llm(): +def test_auto_scorer_with_user_defined_llm() -> None: llm, PROMPT = setup() scorer_llm = FakeListChatModel(responses=[300]) chain = pick_best_chain.PickBest.from_llm( @@ -284,7 +286,7 @@ def test_auto_scorer_with_user_defined_llm(): @pytest.mark.requires("vowpal_wabbit_next", "sentence_transformers") -def test_calling_chain_w_reserved_inputs_throws(): +def test_calling_chain_w_reserved_inputs_throws() -> None: llm, PROMPT = setup() chain = pick_best_chain.PickBest.from_llm(llm=llm, prompt=PROMPT) with pytest.raises(ValueError): diff --git a/libs/langchain/tests/unit_tests/chains/rl_chain/test_pick_best_text_embedder.py b/libs/langchain/tests/unit_tests/chains/rl_chain/test_pick_best_text_embedder.py index c299b187203..acc7491c402 100644 --- a/libs/langchain/tests/unit_tests/chains/rl_chain/test_pick_best_text_embedder.py +++ b/libs/langchain/tests/unit_tests/chains/rl_chain/test_pick_best_text_embedder.py @@ -8,7 +8,7 @@ encoded_text = "[ e n c o d e d ] " @pytest.mark.requires("vowpal_wabbit_next") -def test_pickbest_textembedder_missing_context_throws(): +def test_pickbest_textembedder_missing_context_throws() -> None: feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) named_action = {"action": ["0", "1", "2"]} event = pick_best_chain.PickBestEvent( @@ -19,7 +19,7 @@ def test_pickbest_textembedder_missing_context_throws(): @pytest.mark.requires("vowpal_wabbit_next") -def test_pickbest_textembedder_missing_actions_throws(): +def test_pickbest_textembedder_missing_actions_throws() -> None: feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) event = pick_best_chain.PickBestEvent( inputs={}, to_select_from={}, based_on={"context": "context"} @@ -29,7 +29,7 @@ def test_pickbest_textembedder_missing_actions_throws(): @pytest.mark.requires("vowpal_wabbit_next") -def test_pickbest_textembedder_no_label_no_emb(): +def test_pickbest_textembedder_no_label_no_emb() -> None: feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) named_actions = {"action1": ["0", "1", "2"]} expected = """shared |context context \n|action1 0 \n|action1 1 \n|action1 2 """ @@ -41,7 +41,7 @@ def test_pickbest_textembedder_no_label_no_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_pickbest_textembedder_w_label_no_score_no_emb(): +def test_pickbest_textembedder_w_label_no_score_no_emb() -> None: feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) named_actions = {"action1": ["0", "1", "2"]} expected = """shared |context context \n|action1 0 \n|action1 1 \n|action1 2 """ @@ -57,7 +57,7 @@ def test_pickbest_textembedder_w_label_no_score_no_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_pickbest_textembedder_w_full_label_no_emb(): +def test_pickbest_textembedder_w_full_label_no_emb() -> None: feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) named_actions = {"action1": ["0", "1", "2"]} expected = ( @@ -75,7 +75,7 @@ def test_pickbest_textembedder_w_full_label_no_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_pickbest_textembedder_w_full_label_w_emb(): +def test_pickbest_textembedder_w_full_label_w_emb() -> None: feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) str1 = "0" str2 = "1" @@ -99,7 +99,7 @@ def test_pickbest_textembedder_w_full_label_w_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_pickbest_textembedder_w_full_label_w_embed_and_keep(): +def test_pickbest_textembedder_w_full_label_w_embed_and_keep() -> None: feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) str1 = "0" str2 = "1" @@ -123,7 +123,7 @@ def test_pickbest_textembedder_w_full_label_w_embed_and_keep(): @pytest.mark.requires("vowpal_wabbit_next") -def test_pickbest_textembedder_more_namespaces_no_label_no_emb(): +def test_pickbest_textembedder_more_namespaces_no_label_no_emb() -> None: feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) named_actions = {"action1": [{"a": "0", "b": "0"}, "1", "2"]} context = {"context1": "context1", "context2": "context2"} @@ -136,7 +136,7 @@ def test_pickbest_textembedder_more_namespaces_no_label_no_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_pickbest_textembedder_more_namespaces_w_label_no_emb(): +def test_pickbest_textembedder_more_namespaces_w_label_no_emb() -> None: feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) named_actions = {"action1": [{"a": "0", "b": "0"}, "1", "2"]} context = {"context1": "context1", "context2": "context2"} @@ -150,7 +150,7 @@ def test_pickbest_textembedder_more_namespaces_w_label_no_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_pickbest_textembedder_more_namespaces_w_full_label_no_emb(): +def test_pickbest_textembedder_more_namespaces_w_full_label_no_emb() -> None: feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) named_actions = {"action1": [{"a": "0", "b": "0"}, "1", "2"]} context = {"context1": "context1", "context2": "context2"} @@ -164,7 +164,7 @@ def test_pickbest_textembedder_more_namespaces_w_full_label_no_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_pickbest_textembedder_more_namespaces_w_full_label_w_full_emb(): +def test_pickbest_textembedder_more_namespaces_w_full_label_w_full_emb() -> None: feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) str1 = "0" @@ -195,7 +195,9 @@ def test_pickbest_textembedder_more_namespaces_w_full_label_w_full_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_pickbest_textembedder_more_namespaces_w_full_label_w_full_embed_and_keep(): +def test_pickbest_textembedder_more_namespaces_w_full_label_w_full_embed_and_keep() -> ( + None +): feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) str1 = "0" @@ -228,7 +230,7 @@ def test_pickbest_textembedder_more_namespaces_w_full_label_w_full_embed_and_kee @pytest.mark.requires("vowpal_wabbit_next") -def test_pickbest_textembedder_more_namespaces_w_full_label_w_partial_emb(): +def test_pickbest_textembedder_more_namespaces_w_full_label_w_partial_emb() -> None: feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) str1 = "0" @@ -262,7 +264,9 @@ def test_pickbest_textembedder_more_namespaces_w_full_label_w_partial_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_pickbest_textembedder_more_namespaces_w_full_label_w_partial_embed_and_keep(): +def test_pickbest_textembedder_more_namespaces_w_full_label_w_partial_embed_and_keep() -> ( + None +): feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) str1 = "0" @@ -299,7 +303,7 @@ def test_pickbest_textembedder_more_namespaces_w_full_label_w_partial_embed_and_ @pytest.mark.requires("vowpal_wabbit_next") -def test_raw_features_underscored(): +def test_raw_features_underscored() -> None: feature_embedder = pick_best_chain.PickBestFeatureEmbedder(model=MockEncoder()) str1 = "this is a long string" str1_underscored = str1.replace(" ", "_") diff --git a/libs/langchain/tests/unit_tests/chains/rl_chain/test_rl_chain_base_embedder.py b/libs/langchain/tests/unit_tests/chains/rl_chain/test_rl_chain_base_embedder.py index 895fa8ebb60..c9f8416ceb7 100644 --- a/libs/langchain/tests/unit_tests/chains/rl_chain/test_rl_chain_base_embedder.py +++ b/libs/langchain/tests/unit_tests/chains/rl_chain/test_rl_chain_base_embedder.py @@ -7,13 +7,13 @@ encoded_text = "[ e n c o d e d ] " @pytest.mark.requires("vowpal_wabbit_next") -def test_simple_context_str_no_emb(): +def test_simple_context_str_no_emb() -> None: expected = [{"a_namespace": "test"}] assert base.embed("test", MockEncoder(), "a_namespace") == expected @pytest.mark.requires("vowpal_wabbit_next") -def test_simple_context_str_w_emb(): +def test_simple_context_str_w_emb() -> None: str1 = "test" encoded_str1 = " ".join(char for char in str1) expected = [{"a_namespace": encoded_text + encoded_str1}] @@ -28,7 +28,7 @@ def test_simple_context_str_w_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_simple_context_str_w_nested_emb(): +def test_simple_context_str_w_nested_emb() -> None: # nested embeddings, innermost wins str1 = "test" encoded_str1 = " ".join(char for char in str1) @@ -46,13 +46,13 @@ def test_simple_context_str_w_nested_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_context_w_namespace_no_emb(): +def test_context_w_namespace_no_emb() -> None: expected = [{"test_namespace": "test"}] assert base.embed({"test_namespace": "test"}, MockEncoder()) == expected @pytest.mark.requires("vowpal_wabbit_next") -def test_context_w_namespace_w_emb(): +def test_context_w_namespace_w_emb() -> None: str1 = "test" encoded_str1 = " ".join(char for char in str1) expected = [{"test_namespace": encoded_text + encoded_str1}] @@ -67,7 +67,7 @@ def test_context_w_namespace_w_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_context_w_namespace_w_emb2(): +def test_context_w_namespace_w_emb2() -> None: str1 = "test" encoded_str1 = " ".join(char for char in str1) expected = [{"test_namespace": encoded_text + encoded_str1}] @@ -82,7 +82,7 @@ def test_context_w_namespace_w_emb2(): @pytest.mark.requires("vowpal_wabbit_next") -def test_context_w_namespace_w_some_emb(): +def test_context_w_namespace_w_some_emb() -> None: str1 = "test1" str2 = "test2" encoded_str2 = " ".join(char for char in str2) @@ -111,7 +111,7 @@ def test_context_w_namespace_w_some_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_simple_action_strlist_no_emb(): +def test_simple_action_strlist_no_emb() -> None: str1 = "test1" str2 = "test2" str3 = "test3" @@ -120,7 +120,7 @@ def test_simple_action_strlist_no_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_simple_action_strlist_w_emb(): +def test_simple_action_strlist_w_emb() -> None: str1 = "test1" str2 = "test2" str3 = "test3" @@ -148,7 +148,7 @@ def test_simple_action_strlist_w_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_simple_action_strlist_w_some_emb(): +def test_simple_action_strlist_w_some_emb() -> None: str1 = "test1" str2 = "test2" str3 = "test3" @@ -181,7 +181,7 @@ def test_simple_action_strlist_w_some_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_action_w_namespace_no_emb(): +def test_action_w_namespace_no_emb() -> None: str1 = "test1" str2 = "test2" str3 = "test3" @@ -204,7 +204,7 @@ def test_action_w_namespace_no_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_action_w_namespace_w_emb(): +def test_action_w_namespace_w_emb() -> None: str1 = "test1" str2 = "test2" str3 = "test3" @@ -246,7 +246,7 @@ def test_action_w_namespace_w_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_action_w_namespace_w_emb2(): +def test_action_w_namespace_w_emb2() -> None: str1 = "test1" str2 = "test2" str3 = "test3" @@ -292,7 +292,7 @@ def test_action_w_namespace_w_emb2(): @pytest.mark.requires("vowpal_wabbit_next") -def test_action_w_namespace_w_some_emb(): +def test_action_w_namespace_w_some_emb() -> None: str1 = "test1" str2 = "test2" str3 = "test3" @@ -333,7 +333,7 @@ def test_action_w_namespace_w_some_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_action_w_namespace_w_emb_w_more_than_one_item_in_first_dict(): +def test_action_w_namespace_w_emb_w_more_than_one_item_in_first_dict() -> None: str1 = "test1" str2 = "test2" str3 = "test3" @@ -384,7 +384,7 @@ def test_action_w_namespace_w_emb_w_more_than_one_item_in_first_dict(): @pytest.mark.requires("vowpal_wabbit_next") -def test_one_namespace_w_list_of_features_no_emb(): +def test_one_namespace_w_list_of_features_no_emb() -> None: str1 = "test1" str2 = "test2" expected = [{"test_namespace": [str1, str2]}] @@ -392,7 +392,7 @@ def test_one_namespace_w_list_of_features_no_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_one_namespace_w_list_of_features_w_some_emb(): +def test_one_namespace_w_list_of_features_w_some_emb() -> None: str1 = "test1" str2 = "test2" encoded_str2 = " ".join(char for char in str2) @@ -404,24 +404,24 @@ def test_one_namespace_w_list_of_features_w_some_emb(): @pytest.mark.requires("vowpal_wabbit_next") -def test_nested_list_features_throws(): +def test_nested_list_features_throws() -> None: with pytest.raises(ValueError): base.embed({"test_namespace": [[1, 2], [3, 4]]}, MockEncoder()) @pytest.mark.requires("vowpal_wabbit_next") -def test_dict_in_list_throws(): +def test_dict_in_list_throws() -> None: with pytest.raises(ValueError): base.embed({"test_namespace": [{"a": 1}, {"b": 2}]}, MockEncoder()) @pytest.mark.requires("vowpal_wabbit_next") -def test_nested_dict_throws(): +def test_nested_dict_throws() -> None: with pytest.raises(ValueError): base.embed({"test_namespace": {"a": {"b": 1}}}, MockEncoder()) @pytest.mark.requires("vowpal_wabbit_next") -def test_list_of_tuples_throws(): +def test_list_of_tuples_throws() -> None: with pytest.raises(ValueError): base.embed({"test_namespace": [("a", 1), ("b", 2)]}, MockEncoder())