Skip to content

fix(language): use structured output for LLM response parsing #92

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jakub-walaszczyk merged 1 commit into from
Jul 3, 2026
+59 −49
Merged

fix(language): use structured output for LLM response parsing #92

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 34 additions & 26 deletions ai4rag/search_space/prepare/language_detection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
# Copyright IBM Corp. 2026
# SPDX-License-Identifier: Apache-2.0
# -----------------------------------------------------------------------------
import json
import re

from ai4rag import logger
Expand Down Expand Up @@ -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
----------
Expand All @@ -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": "<iso-639-1>", "name": "<language>"}`` on success,
or ``None`` on failure.
"""
sample_text = "\n".join(f"- {q}" for q in questions[:5])

Expand All @@ -229,43 +229,51 @@ 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."
),
},
{
"role": "user",
"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)
Expand Down
48 changes: 25 additions & 23 deletions tests/unit/ai4rag/search_space/prepare/test_language_detection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
# -----------------------------------------------------------------------------
from __future__ import annotations

import json
from unittest.mock import MagicMock

import pytest
Expand All @@ -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]
Expand Down Expand Up @@ -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)

Expand All @@ -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)

Expand All @@ -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
Expand Down
Loading