diff --git a/ai4rag/search_space/prepare/language_detection.py b/ai4rag/search_space/prepare/language_detection.py index 56655d6f..d0c333b7 100644 --- a/ai4rag/search_space/prepare/language_detection.py +++ b/ai4rag/search_space/prepare/language_detection.py @@ -2,6 +2,7 @@ # Copyright IBM Corp. 2026 # SPDX-License-Identifier: Apache-2.0 # ----------------------------------------------------------------------------- +import json import re from ai4rag import logger @@ -202,9 +203,8 @@ def detect_language_with_llm( ) -> dict[str, str] | None: """Detect the dominant language from sample questions using an LLM. - Sends a small sample of questions to a generation model registered in OGX - and asks it to return the ISO 639-1 code. Models listed in - *allowed_generation_models* are preferred when available. + Sends a small sample of questions to a generation model and uses + JSON-schema structured output to obtain a single ISO 639-1 code. Parameters ---------- @@ -217,8 +217,8 @@ def detect_language_with_llm( Returns ------- dict[str, str] | None - A dictionary with ``code`` and ``name`` keys when a non-English - language is detected, or ``None`` for English / on failure. + ``{"code": "", "name": ""}`` on success, + or ``None`` on failure. """ sample_text = "\n".join(f"- {q}" for q in questions[:5]) @@ -229,8 +229,8 @@ def detect_language_with_llm( "role": "system", "content": ( "You are a language detection assistant. " - "Given text samples, respond with ONLY the ISO 639-1 language code. " - "Nothing else — just the code." + "Given text samples, respond with the ISO 639-1 language code " + "of the dominant language." ), }, { @@ -238,34 +238,42 @@ def detect_language_with_llm( "content": f"What language are these questions written in?\n{sample_text}", }, ], - max_completion_tokens=10, + max_completion_tokens=15, temperature=0.0, + response_format={ + "type": "json_schema", + "json_schema": { + "name": "language_detection", + "schema": { + "type": "object", + "properties": { + "code": { + "type": "string", + "description": "ISO 639-1 language code", + }, + }, + "required": ["code"], + "additionalProperties": False, + }, + "strict": True, + }, + }, ) raw_content = response[0].message.content if not raw_content or not isinstance(raw_content, str): raise ValueError(f"Invalid response content: {type(raw_content)}") - cleaned = raw_content.strip().lower().replace('"', "").replace("'", "") - if not cleaned: - raise ValueError("Empty response after cleanup") - - code_pattern = r"[a-z]{2}(?:-[a-z]{2,4})?" - # Try targeted patterns to avoid matching English stop words (e.g. "is", "it", "no") - match = ( - re.match(rf"^({code_pattern})\s*$", cleaned) # code only - or re.match(rf"^({code_pattern})\s", cleaned) # code at start, then more text - or re.search(rf"\(({code_pattern})\)", cleaned) # code in parentheses - ) - if not match: - raise ValueError(f"No ISO 639-1 code found in response: {cleaned[:50]}") + detected_code = json.loads(raw_content)["code"].strip().lower() + if not re.fullmatch(r"[a-z]{2}(?:-[a-z]{2,4})?", detected_code): + raise ValueError(f"Malformed language code: {detected_code!r}") - detected_code = match.group(1).split("-")[0] - name = LANGUAGE_MAP.get(match.group(1)) or LANGUAGE_MAP.get(detected_code) + base_code = detected_code.split("-")[0] + name = LANGUAGE_MAP.get(detected_code) or LANGUAGE_MAP.get(base_code) if not name: - raise ValueError(f"Unsupported language code '{detected_code}' from response: {cleaned[:50]}") + raise ValueError(f"Unsupported language code: {detected_code!r}") - logger.info("Language detected via LLM: %s (%s)", detected_code, name) - return {"code": detected_code, "name": name} + logger.info("Language detected via LLM: %s (%s)", base_code, name) + return {"code": base_code, "name": name} except Exception as exc: logger.warning("LLM language detection failed: %s", exc) diff --git a/tests/unit/ai4rag/search_space/prepare/test_language_detection.py b/tests/unit/ai4rag/search_space/prepare/test_language_detection.py index 87574e8b..82671dc0 100644 --- a/tests/unit/ai4rag/search_space/prepare/test_language_detection.py +++ b/tests/unit/ai4rag/search_space/prepare/test_language_detection.py @@ -4,6 +4,7 @@ # ----------------------------------------------------------------------------- from __future__ import annotations +import json from unittest.mock import MagicMock import pytest @@ -18,7 +19,7 @@ def mock_generation_model() -> MagicMock: """Return a MagicMock that behaves like an OGXFoundationModel.""" mock_choice = MagicMock() - mock_choice.message.content = "ja" + mock_choice.message.content = json.dumps({"code": "ja"}) model = MagicMock() model.chat.return_value = [mock_choice] @@ -65,7 +66,7 @@ def test_detects_japanese(self, mock_generation_model, sample_questions): mock_generation_model.chat.assert_called_once() def test_detects_english(self, mock_generation_model, sample_questions): - mock_generation_model.chat.return_value[0].message.content = "en" + mock_generation_model.chat.return_value[0].message.content = json.dumps({"code": "en"}) result = detect_language_with_llm(sample_questions, mock_generation_model) @@ -79,7 +80,7 @@ def test_api_failure_returns_none(self, mock_generation_model, sample_questions) assert result is None def test_unsupported_language_code_returns_none(self, mock_generation_model, sample_questions): - mock_generation_model.chat.return_value[0].message.content = "xx" + mock_generation_model.chat.return_value[0].message.content = json.dumps({"code": "xx"}) result = detect_language_with_llm(sample_questions, mock_generation_model) @@ -94,55 +95,56 @@ def test_samples_at_most_five_questions(self, mock_generation_model): user_content = call_kwargs.kwargs["messages"][1]["content"] assert user_content.count("- Question") == 5 - def test_empty_llm_response_returns_none(self, mock_generation_model, sample_questions): - mock_generation_model.chat.return_value[0].message.content = " " + def test_malformed_json_returns_none(self, mock_generation_model, sample_questions): + mock_generation_model.chat.return_value[0].message.content = "not json" result = detect_language_with_llm(sample_questions, mock_generation_model) assert result is None - def test_passes_overridden_chat_params(self, mock_generation_model, sample_questions): - """Verify that chat is called with overridden max_completion_tokens and temperature.""" + def test_passes_expected_chat_params(self, mock_generation_model, sample_questions): detect_language_with_llm(sample_questions, mock_generation_model) call_kwargs = mock_generation_model.chat.call_args.kwargs - assert call_kwargs["max_completion_tokens"] == 10 + assert call_kwargs["max_completion_tokens"] == 15 assert call_kwargs["temperature"] == 0.0 + assert call_kwargs["response_format"]["type"] == "json_schema" + assert call_kwargs["response_format"]["json_schema"]["strict"] is True - def test_strips_quotes_from_response(self, mock_generation_model, sample_questions): - mock_generation_model.chat.return_value[0].message.content = '"fr"' + def test_malformed_code_format_returns_none(self, mock_generation_model, sample_questions): + mock_generation_model.chat.return_value[0].message.content = json.dumps({"code": "123"}) result = detect_language_with_llm(sample_questions, mock_generation_model) - assert result == {"code": "fr", "name": "French"} + assert result is None - def test_extracts_code_at_start_of_verbose_response(self, mock_generation_model, sample_questions): - mock_generation_model.chat.return_value[0].message.content = "de (German)" + def test_missing_code_key_returns_none(self, mock_generation_model, sample_questions): + mock_generation_model.chat.return_value[0].message.content = json.dumps({"language": "en"}) result = detect_language_with_llm(sample_questions, mock_generation_model) - assert result == {"code": "de", "name": "German"} + assert result is None - def test_extracts_code_in_parentheses(self, mock_generation_model, sample_questions): - mock_generation_model.chat.return_value[0].message.content = "German (de)" + def test_extracts_code_with_region_suffix(self, mock_generation_model, sample_questions): + mock_generation_model.chat.return_value[0].message.content = json.dumps({"code": "zh-cn"}) result = detect_language_with_llm(sample_questions, mock_generation_model) - assert result == {"code": "de", "name": "German"} + assert result == {"code": "zh", "name": "Chinese"} - def test_verbose_response_without_extractable_code_returns_none(self, mock_generation_model, sample_questions): - mock_generation_model.chat.return_value[0].message.content = "The language is German" + def test_normalises_uppercase_code(self, mock_generation_model, sample_questions): + mock_generation_model.chat.return_value[0].message.content = json.dumps({"code": "FR"}) result = detect_language_with_llm(sample_questions, mock_generation_model) - assert result is None + assert result == {"code": "fr", "name": "French"} - def test_extracts_code_with_region_suffix(self, mock_generation_model, sample_questions): - mock_generation_model.chat.return_value[0].message.content = "zh-cn" + def test_strips_whitespace_from_code(self, mock_generation_model, sample_questions): + mock_generation_model.chat.return_value[0].message.content = json.dumps({"code": " de "}) result = detect_language_with_llm(sample_questions, mock_generation_model) - assert result == {"code": "zh", "name": "Chinese"} + assert result == {"code": "de", "name": "German"} def test_none_content_returns_none(self, mock_generation_model, sample_questions): mock_generation_model.chat.return_value[0].message.content = None