chore: chain tests (#62)

Author: Efreet408

tests/chain_tests/test_decontextualization_chain.py

@@ -0,0 +1,100 @@
+from unittest.mock import patch
+
+from langchain_core.language_models.fake_chat_models import FakeListChatModel
+
+from aidial_rag_eval.generation.models.converters.llm_decontextualization_converter import (
+    LLMNoPronounsConverter,
+)
+from aidial_rag_eval.generation.utils.segmented_text import SegmentedText
+
+
+def test_valid_json_response():
+    fake_llm = FakeListChatModel(
+        responses=['{"segments": ["John went to the store.", "John bought milk."]}']
+    )
+    converter = LLMNoPronounsConverter(model=fake_llm, max_concurrency=1)
+
+    segmented_text = SegmentedText(
+        segments=["John went to the store.", "He bought milk."], delimiters=[" "]
+    )
+
+    converter.transform_texts([segmented_text], show_progress_bar=False)
+
+    assert segmented_text.segments == ["John went to the store.", "John bought milk."]
+
+
+def test_invalid_json_response():
+    fake_llm = FakeListChatModel(responses=["not a valid json at all"])
+    converter = LLMNoPronounsConverter(model=fake_llm, max_concurrency=1)
+
+    segmented_text = SegmentedText(
+        segments=["John went to the store.", "He bought milk."], delimiters=[" "]
+    )
+    original_segments = segmented_text.segments.copy()
+
+    converter.transform_texts([segmented_text], show_progress_bar=False)
+
+    assert segmented_text.segments == original_segments
+
+
+def test_json_missing_segments_key():
+    fake_llm = FakeListChatModel(
+        responses=['{"wrong_key": ["John went to the store.", "John bought milk."]}']
+    )
+    converter = LLMNoPronounsConverter(model=fake_llm, max_concurrency=1)
+
+    segmented_text = SegmentedText(
+        segments=["John went to the store.", "He bought milk."], delimiters=[" "]
+    )
+    original_segments = segmented_text.segments.copy()
+
+    converter.transform_texts([segmented_text], show_progress_bar=False)
+
+    assert segmented_text.segments == original_segments
+
+
+def test_segment_count_mismatch():
+    fake_llm = FakeListChatModel(responses=['{"segments": ["only one segment"]}'])
+    converter = LLMNoPronounsConverter(model=fake_llm, max_concurrency=1)
+
+    segmented_text = SegmentedText(
+        segments=["John went to the store.", "He bought milk."], delimiters=[" "]
+    )
+    original_segments = segmented_text.segments.copy()
+
+    converter.transform_texts([segmented_text], show_progress_bar=False)
+
+    assert segmented_text.segments == original_segments
+
+
+def test_empty_response():
+    fake_llm = FakeListChatModel(responses=[""])
+    converter = LLMNoPronounsConverter(model=fake_llm, max_concurrency=1)
+
+    segmented_text = SegmentedText(
+        segments=["John went to the store.", "He bought milk."], delimiters=[" "]
+    )
+    original_segments = segmented_text.segments.copy()
+
+    converter.transform_texts([segmented_text], show_progress_bar=False)
+
+    assert segmented_text.segments == original_segments
+
+
+def test_invoke_raises_exception():
+    fake_llm = FakeListChatModel(responses=[""])
+
+    converter = LLMNoPronounsConverter(model=fake_llm, max_concurrency=1)
+
+    segmented_text = SegmentedText(
+        segments=["John went to the store.", "He bought milk."], delimiters=[" "]
+    )
+
+    with patch.object(
+        FakeListChatModel, "invoke", side_effect=Exception("LLM invoke failed")
+    ):
+        try:
+            converter.transform_texts([segmented_text], show_progress_bar=False)
+            raise AssertionError("Expected exception was not raised")
+        except Exception as e:
+            assert str(e) == "LLM invoke failed"

tests/chain_tests/test_inference_scorer_chain.py

@@ -0,0 +1,129 @@
+from unittest.mock import patch
+
+import pytest
+from langchain_core.language_models.fake_chat_models import FakeListChatModel
+
+from aidial_rag_eval.generation.models.inference_scorers.llm_inference_scorer import (
+    LLMInferenceScorer,
+)
+from aidial_rag_eval.generation.types import InferenceInputs
+
+
+def _create_inference_input(
+    statements: list[str], premise: str = "Water is wet."
+) -> InferenceInputs:
+    return InferenceInputs(
+        hypothesis_id=0,
+        premise=premise,
+        statements=statements,
+        document_name="test_doc",
+    )
+
+
+def test_valid_json_response():
+    fake_llm = FakeListChatModel(
+        responses=['{"results": [{"tag": "ENT", "explanation": "test"}]}']
+    )
+    scorer = LLMInferenceScorer(model=fake_llm, max_concurrency=1)
+
+    inputs = [_create_inference_input(["Water is wet."])]
+
+    results = scorer.get_inference(inputs, show_progress_bar=False)
+
+    assert len(results) == 1
+    assert results[0].inference == 1.0
+    assert (
+        results[0].explanation
+        == '[{"tag": "ENT", "explanation": "test", "statement": "Water is wet."}]'
+    )
+
+
+def test_invalid_json_response():
+    fake_llm = FakeListChatModel(responses=["not valid json"])
+    scorer = LLMInferenceScorer(model=fake_llm, max_concurrency=1)
+
+    inputs = [_create_inference_input(["Statement1"])]
+
+    results = scorer.get_inference(inputs, show_progress_bar=False)
+
+    assert len(results) == 1
+    assert results[0].inference == 0.0
+    assert results[0].explanation == ""
+
+
+def test_json_missing_tag_key():
+    fake_llm = FakeListChatModel(responses=['{"results": [{"explanation": "value"}]}'])
+    scorer = LLMInferenceScorer(model=fake_llm, max_concurrency=1)
+
+    inputs = [_create_inference_input(["Statement1"])]
+
+    results = scorer.get_inference(inputs, show_progress_bar=False)
+
+    assert results[0].inference == 0.0
+    assert results[0].explanation == ""
+
+
+@pytest.mark.skip(reason="explanation key check is not implemented")
+def test_json_missing_explanation_key():
+    fake_llm = FakeListChatModel(responses=['{"results": [{"tag": "ENT"}]}'])
+    scorer = LLMInferenceScorer(model=fake_llm, max_concurrency=1)
+
+    inputs = [_create_inference_input(["Statement1"])]
+
+    results = scorer.get_inference(inputs, show_progress_bar=False)
+
+    assert results[0].inference == 0.0
+    assert results[0].explanation == ""
+
+
+def test_output_count_mismatch():
+    fake_llm = FakeListChatModel(
+        responses=['{"results": [{"tag": "ENT", "explanation": "test"}]}']
+    )
+    scorer = LLMInferenceScorer(model=fake_llm, max_concurrency=1)
+
+    inputs = [_create_inference_input(["Statement1", "Statement2"])]
+
+    results = scorer.get_inference(inputs, show_progress_bar=False)
+
+    assert results[0].inference == 0.0
+
+
+def test_empty_response():
+    fake_llm = FakeListChatModel(responses=[""])
+    scorer = LLMInferenceScorer(model=fake_llm, max_concurrency=1)
+
+    inputs = [_create_inference_input(["Statement1"])]
+
+    results = scorer.get_inference(inputs, show_progress_bar=False)
+
+    assert results[0].inference == 0.0
+    assert results[0].explanation == ""
+
+
+def test_empty_statements():
+    fake_llm = FakeListChatModel(responses=["should not be called"])
+    scorer = LLMInferenceScorer(model=fake_llm, max_concurrency=1)
+
+    inputs = [_create_inference_input([])]
+
+    results = scorer.get_inference(inputs, show_progress_bar=False)
+
+    assert results[0].inference == 0.0
+    assert results[0].explanation == ""
+
+
+def test_invoke_raises_exception():
+    fake_llm = FakeListChatModel(responses=[""])
+    scorer = LLMInferenceScorer(model=fake_llm, max_concurrency=1)
+
+    inputs = [_create_inference_input(["Statement1"])]
+
+    with patch.object(
+        FakeListChatModel, "invoke", side_effect=Exception("LLM invoke failed")
+    ):
+        try:
+            scorer.get_inference(inputs, show_progress_bar=False)
+            raise AssertionError("Expected exception was not raised")
+        except Exception as e:
+            assert str(e) == "LLM invoke failed"

tests/chain_tests/test_statement_extractor_chain.py

@@ -0,0 +1,103 @@
+from unittest.mock import patch
+
+from langchain_core.language_models.fake_chat_models import FakeListChatModel
+
+from aidial_rag_eval.generation.models.statement_extractor.llm_statement_extractor import (
+    LLMStatementExtractor,
+)
+
+
+def test_valid_json_response():
+    fake_llm = FakeListChatModel(
+        responses=[
+            """
+            {
+                "hypothesis_statements":
+                    [
+                        {
+                            "statements": ["statement11"]
+                        },
+                        {
+                            "statements": ["statement21"]
+                        }
+                    ]
+            }"""
+        ]
+    )
+    extractor = LLMStatementExtractor(model=fake_llm, max_concurrency=1)
+
+    hypothesis_segments = ["hypothesis_segment1", "hypothesis_segment1"]
+
+    result = extractor.extract([hypothesis_segments], show_progress_bar=False)
+
+    assert result == [[["statement11"], ["statement21"]]]
+
+
+def test_invalid_json_response():
+    fake_llm = FakeListChatModel(responses=["not valid json"])
+    extractor = LLMStatementExtractor(model=fake_llm, max_concurrency=1)
+
+    hypothesis_segments = ["hypothesis_segment1", "hypothesis_segment2"]
+
+    result = extractor.extract([hypothesis_segments], show_progress_bar=False)
+
+    assert result == [
+        [[hypothesis_segment] for hypothesis_segment in hypothesis_segments]
+    ]
+
+
+def test_json_wrong_structure():
+    fake_llm = FakeListChatModel(responses=['{"wrong_key": "not a list"}'])
+    extractor = LLMStatementExtractor(model=fake_llm, max_concurrency=1)
+
+    hypothesis_segments = ["hypothesis_segment1", "hypothesis_segment2"]
+
+    result = extractor.extract([hypothesis_segments], show_progress_bar=False)
+
+    assert result == [
+        [[hypothesis_segment] for hypothesis_segment in hypothesis_segments]
+    ]
+
+
+def test_statement_count_mismatch():
+    fake_llm = FakeListChatModel(
+        responses=['{"hypothesis_statements": [{"statements": ["statement1"]}]}']
+    )
+    extractor = LLMStatementExtractor(model=fake_llm, max_concurrency=1)
+
+    hypothesis_segments = ["hypothesis_segment1", "hypothesis_segment1"]
+
+    result = extractor.extract([hypothesis_segments], show_progress_bar=False)
+
+    assert result == [
+        [[hypothesis_segment] for hypothesis_segment in hypothesis_segments]
+    ]
+
+
+def test_empty_response():
+    fake_llm = FakeListChatModel(responses=[""])
+    extractor = LLMStatementExtractor(model=fake_llm, max_concurrency=1)
+
+    hypothesis_segments = ["hypothesis_segment1", "hypothesis_segment1"]
+
+    result = extractor.extract([hypothesis_segments], show_progress_bar=False)
+
+    assert result == [
+        [[hypothesis_segment] for hypothesis_segment in hypothesis_segments]
+    ]
+
+
+def test_invoke_raises_exception():
+    fake_llm = FakeListChatModel(responses=[""])
+    extractor = LLMStatementExtractor(model=fake_llm, max_concurrency=1)
+
+    hypothesis_segments = ["hypothesis_segment1", "hypothesis_segment1"]
+
+    with patch.object(
+        FakeListChatModel, "invoke", side_effect=Exception("LLM invoke failed")
+    ):
+        try:
+            extractor.extract([hypothesis_segments], show_progress_bar=False)
+            raise AssertionError("Expected exception was not raised")
+        except Exception as e:
+            assert str(e) == "LLM invoke failed"