Ctrl K
ComunidadGuías…

tests/test_query_utils.py

Rama main · réplica de lectura del repositorio original.

Archivo original ↗Preguntar a KiBot
"""
Tests for `src.pipeline.query_utils`.

  - `coerce_search_terms` accepts the four LLM output forms seen in
    practice (flat list, dict per language, list of dicts, single string).
  - `detect_term_language` assigns umlaut terms to DE, others to EN.
  - `build_terms_by_lang` aggregates over several questions and prefers
    `search_terms_by_lang` over `search_terms`.
  - robust against edge cases: empty lists, None values, invalid types.
"""

import unittest
from dataclasses import dataclass, field

from src.pipeline.query_utils import (
    build_terms_by_lang,
    coerce_search_terms,
    detect_term_language,
)


# Minimal ResearchQuestion stub for the tests — so that we do not have
# to instantiate the real model with all its fields
@dataclass
class _Q:
    search_terms: list = field(default_factory=list)
    search_terms_by_lang: dict = field(default_factory=dict)
    search_langs: list = field(default_factory=lambda: ["de", "en"])


class TestDetectTermLanguage(unittest.TestCase):

    def test_umlaut_is_german(self):
        self.assertEqual(detect_term_language("Stromaufnahme"), "en")  # no umlaut
        self.assertEqual(detect_term_language("Verlässlichkeit"), "de")
        self.assertEqual(detect_term_language("Größe"), "de")
        self.assertEqual(detect_term_language("Übertragung"), "de")

    def test_function_word_is_german(self):
        self.assertEqual(detect_term_language("der schnellste Server"), "de")
        self.assertEqual(detect_term_language("für Forschung"), "de")

    def test_pure_english_is_english(self):
        self.assertEqual(detect_term_language("machine learning"), "en")
        self.assertEqual(detect_term_language("DGX B300 power"), "en")

    def test_neutral_term_defaults_to_english(self):
        # product names without language-specific markers → en default
        self.assertEqual(detect_term_language("DGX B300"), "en")
        self.assertEqual(detect_term_language("BIFOLD"), "en")
        self.assertEqual(detect_term_language("Jonas Brenner"), "en")

    def test_empty_or_invalid(self):
        self.assertEqual(detect_term_language(""), "en")
        self.assertEqual(detect_term_language("   "), "en")
        self.assertEqual(detect_term_language(None), "en")  # type: ignore


class TestCoerceSearchTerms(unittest.TestCase):

    def test_form_a_flat_list(self):
        flat, by_lang = coerce_search_terms([
            "DGX B300 power", "Stromaufnahme Server",
        ])
        self.assertEqual(flat, ["DGX B300 power", "Stromaufnahme Server"])
        # languages via the heuristic; "Stromaufnahme" without an umlaut → en default
        # (not ideal, but consistent — the question level corrects it)
        self.assertIn("en", by_lang)

    def test_form_b_list_of_dicts(self):
        raw = [
            {"term": "DGX B300 Stromaufnahme", "lang": "de"},
            {"term": "DGX B300 power consumption", "lang": "en"},
        ]
        flat, by_lang = coerce_search_terms(raw)
        self.assertEqual(len(flat), 2)
        self.assertEqual(by_lang["de"], ["DGX B300 Stromaufnahme"])
        self.assertEqual(by_lang["en"], ["DGX B300 power consumption"])

    def test_form_c_dict_per_language(self):
        raw = {
            "de": ["term1", "term2"],
            "en": ["term3"],
        }
        flat, by_lang = coerce_search_terms(raw)
        self.assertEqual(len(flat), 3)
        self.assertEqual(by_lang["de"], ["term1", "term2"])
        self.assertEqual(by_lang["en"], ["term3"])

    def test_form_d_single_string(self):
        flat, by_lang = coerce_search_terms("DGX B300")
        self.assertEqual(flat, ["DGX B300"])

    def test_form_e_empty(self):
        for empty in (None, [], {}, ""):
            flat, by_lang = coerce_search_terms(empty)
            self.assertEqual(flat, [])
            self.assertEqual(by_lang, {})

    def test_dedup_within_call(self):
        flat, _ = coerce_search_terms(["term1", "term1", "term2", "term1"])
        self.assertEqual(flat, ["term1", "term2"])

    def test_language_normalization(self):
        """LLMs sometimes answer with 'german'/'deutsch' instead of 'de' — we normalise."""
        raw = [
            {"term": "x", "lang": "german"},
            {"term": "y", "lang": "deu"},
            {"term": "z", "lang": "english"},
        ]
        _, by_lang = coerce_search_terms(raw)
        self.assertIn("de", by_lang)
        self.assertIn("en", by_lang)
        self.assertEqual(set(by_lang["de"]), {"x", "y"})
        self.assertEqual(by_lang["en"], ["z"])

    def test_invalid_items_skipped(self):
        raw = [
            "valid",
            123,                          # int → ignored
            None,                         # None → ignored
            {"no_term_field": "x"},       # dict without a term → with the heuristic
            {"term": "ok", "lang": "de"},
        ]
        flat, by_lang = coerce_search_terms(raw)
        self.assertIn("valid", flat)
        self.assertIn("ok", flat)
        self.assertIn("x", flat)  # the only string value is taken

    def test_nested_dict_with_term_key(self):
        """Some LLMs nest: {term: {value: '...'}}"""
        raw = [{"term": "echter Term", "lang": "de"}]
        flat, by_lang = coerce_search_terms(raw)
        self.assertEqual(flat, ["echter Term"])
        self.assertEqual(by_lang["de"], ["echter Term"])


class TestBuildTermsByLang(unittest.TestCase):

    def test_uses_structured_when_available(self):
        """search_terms_by_lang is set → it is preferred."""
        q = _Q(
            search_terms_by_lang={
                "de": ["term DE"],
                "en": ["term EN"],
            },
            search_terms=["fallback"],   # should be ignored
        )
        terms_by_lang, all_langs = build_terms_by_lang([q])

        self.assertEqual(terms_by_lang["de"], ["term DE"])
        self.assertEqual(terms_by_lang["en"], ["term EN"])
        self.assertNotIn("fallback", terms_by_lang.get("de", []))
        self.assertNotIn("fallback", terms_by_lang.get("en", []))

    def test_falls_back_to_search_terms(self):
        """Without search_terms_by_lang: split heuristically."""
        q = _Q(
            search_terms=["Größe", "machine learning"],
            search_langs=["de", "en"],
        )
        terms_by_lang, all_langs = build_terms_by_lang([q])

        # Größe → de (umlaut), machine learning → en
        self.assertIn("Größe", terms_by_lang["de"])
        self.assertIn("machine learning", terms_by_lang["en"])

    def test_aggregates_over_multiple_questions(self):
        q1 = _Q(search_terms_by_lang={"de": ["a"], "en": ["b"]})
        q2 = _Q(search_terms_by_lang={"de": ["c"], "en": ["d"]})

        terms_by_lang, all_langs = build_terms_by_lang([q1, q2])

        self.assertEqual(set(terms_by_lang["de"]), {"a", "c"})
        self.assertEqual(set(terms_by_lang["en"]), {"b", "d"})
        self.assertEqual(all_langs, {"de", "en"})

    def test_dedup_across_questions(self):
        q1 = _Q(search_terms_by_lang={"de": ["term"]})
        q2 = _Q(search_terms_by_lang={"de": ["term"]})

        terms_by_lang, _ = build_terms_by_lang([q1, q2])

        self.assertEqual(terms_by_lang["de"], ["term"])  # only in it once

    def test_search_langs_appears_even_without_terms(self):
        """The declared languages are reported."""
        q = _Q(
            search_terms=[],
            search_terms_by_lang={},
            search_langs=["de", "en", "fr"],
        )
        _, all_langs = build_terms_by_lang([q])
        # all three declared languages, even if empty
        self.assertEqual(all_langs, {"de", "en", "fr"})

    def test_heuristic_respects_search_langs(self):
        """If the question only declares DE, an English-looking term belongs in DE too."""
        q = _Q(
            search_terms=["machine learning"],   # heuristically: en
            search_langs=["de"],                  # only DE declared
        )
        terms_by_lang, _ = build_terms_by_lang([q])

        # the term ends up in DE (default_lang), not in EN
        self.assertIn("machine learning", terms_by_lang["de"])
        self.assertNotIn("en", terms_by_lang)

    def test_empty_questions(self):
        terms_by_lang, all_langs = build_terms_by_lang([])
        self.assertEqual(terms_by_lang, {})
        self.assertEqual(all_langs, set())


if __name__ == "__main__":
    unittest.main()