Dage 16: response validation Schemas (#171)

* Minor fixes / improvements
* Used simple data structures to perform validation of LLM responses in the LLM Service
* Explored Pydantic, decided simple dict / lists / custom objetcts

Author: Intrinsical-AI

rre-dataset-generator/src/llm/llm_service.py

@@ -1,13 +1,13 @@
+import json
 from json import JSONDecodeError
-from typing import List, Union
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import HumanMessage, SystemMessage
-import json
+from src.model.query_response import LLMQueryResponse
+from src.model.score_response import LLMScoreResponse
 import logging
 
 from src.model.document import Document
-from typing import Union, List
 
 log = logging.getLogger(__name__)
 
@@ -16,7 +16,7 @@ class LLMService:
     def __init__(self, chat_model: BaseChatModel):
         self.chat_model = chat_model
 
-    def generate_queries(self, document: Document, num_queries_generate_per_doc: int) -> Union[List[str], str]:
+    def generate_queries(self, document: Document, num_queries_generate_per_doc: int) -> LLMQueryResponse:
         """
         Generate queries based on the given document and num_queries_generate_per_doc and
         Returns a list of generated queries or just a generated string in case of LLM hallucination
@@ -35,32 +35,35 @@ def generate_queries(self, document: Document, num_queries_generate_per_doc: int
             HumanMessage(content=f"Document:\n{doc_json}")
         ]
 
-        raw = self.chat_model.invoke(messages).content.strip()
+        # The response from invoke is an AIMessage object which contains all the needed info
+        response = self.chat_model.invoke(messages)
 
-        # in case LLM hallucinates
         try:
-            return json.loads(raw)
-        except JSONDecodeError:
-            log.warning("LLM hallucinated and its response: %s", raw)
-            return raw
+           output = LLMQueryResponse(response_content=response.content)
+        except (KeyError, JSONDecodeError, ValueError) as e:
+            log.warning(f"LLM unexpected response. Raw output: {response.content}")
+            raise ValueError(f"Invalid LLM response: {e}")
 
-    def generate_score(self, document: Document, query: str, relevance_scale: str) -> int:
+        return output
+    
+
+    def generate_score(self, document: Document, query: str, relevance_scale: str) -> LLMScoreResponse:
         """
         Generates a relevance score for a given document-query pair using a specified relevance scale.
         """
         if relevance_scale == "binary":
-            scale = {0, 1}
-            description = (" - 0: the query is NOT relevant to the given document"
-                           " - 1: the query is relevant to the given document")
+            allowed = {0, 1}
+            description = (" - 0: the query is NOT relevant to the given document\n"
+                        " - 1: the query is relevant to the given document")
         elif relevance_scale == "graded":
-            scale = {0, 1, 2}
-            description = (" - 0: the query is NOT relevant to the given document"
-                           " - 1: the query may be relevant to the given document"
-                           " - 2: the document proposed is the answer to the query")
+            allowed = {0, 1, 2}
+            description = (" - 0: the query is NOT relevant to the given document\n"
+                        " - 1: the query may be relevant to the given document\n"
+                        " - 2: the document proposed is the answer to the query")
         else:
-            error_msg = "The relevance scale must be either 'binary' or 'graded'"
-            log.error(error_msg)
-            raise ValueError(error_msg)
+            msg = f"Invalid relevance scale: {relevance_scale}"
+            log.error(msg)
+            raise ValueError(msg)
 
         messages = [
             SystemMessage(
@@ -76,16 +79,19 @@ def generate_score(self, document: Document, query: str, relevance_scale: str) -
             )
         ]
 
-        #response = self.chat_model.with_structured_output(method="json_mode").invoke(messages)
         raw = self.chat_model.invoke(messages).content.strip()
+
         try:
-            response =  int(json.loads(raw)['score'])
-            if response not in scale:
-                error_msg = f"LLM hallucinated the value of the scale. Returned: {response}"
-                log.warning(error_msg)
-                raise ValueError(error_msg)
-            return response
-        except JSONDecodeError:
-            error_msg = f"LLM hallucinated and its response: {raw}"
-            log.warning(error_msg)
-            raise ValueError(error_msg)
+            score = json.loads(raw)['score']
+        except (JSONDecodeError, KeyError) as e:
+            log.debug(f"LLM unexpected response. Raw output: {raw}")
+            raise ValueError(f"Invalid LLM response: {e}")
+
+        try:
+            parsed = LLMScoreResponse(score=score, scale=relevance_scale)
+            return parsed
+        except ValueError as e:
+            log.warning(f"Validation error for score '{score}' on scale '{relevance_scale}': {e}")
+            raise e
+
+

rre-dataset-generator/src/model/query_response.py

@@ -0,0 +1,40 @@
+import json
+from typing import List
+
+class LLMQueryResponse:
+    """
+    Parses and validates a JSON response containing a list of queries.
+    """
+    _queries: List[str]
+
+    def __init__(self, response_content: str):
+        """
+        Initializes the object by parsing and validating the JSON string.
+
+        Args:
+            response_content: A string containing a JSON-encoded list of strings.
+
+        Raises:
+            ValueError: If the content is not a valid JSON list of non-empty strings.
+        """
+        try:
+            parsed_content = json.loads(response_content)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON in `response_content`: {e}")
+
+        if not isinstance(parsed_content, list):
+            raise ValueError("`response_content` must be a JSON list.")
+
+        if not all(isinstance(item, str) for item in parsed_content):
+            raise ValueError("All items in `response_content` must be strings.")
+
+        if any(not item.strip() for item in parsed_content):
+            raise ValueError("Items in `response_content` must not be empty or only whitespace.")
+
+        self._queries = parsed_content
+
+    def get_queries(self) -> List[str]:
+        """
+        Returns the validated list of queries.
+        """
+        return self._queries

rre-dataset-generator/src/model/score_response.py

@@ -0,0 +1,33 @@
+class LLMScoreResponse:
+    """
+    Parses and validates an LLM score response.
+    """
+    def __init__(self, score: int, scale: str = "graded"):
+        """
+        Initializes the object by validating the score.
+
+        Args:
+            score: The relevance score.
+            scale: The relevance scale, either 'binary' (0-1) or 'graded' (0-2).
+
+        Raises:
+            ValueError: If the score is not valid for the given scale.
+        """
+        if scale not in ["binary", "graded"]:
+            raise ValueError(f"Invalid scale: {scale}. Must be 'binary' or 'graded'.")
+            
+        if scale == "binary" and score not in {0, 1}:
+            raise ValueError(f"Score must be 0 or 1 for binary scale, got {score}")
+        elif scale == "graded" and score not in {0, 1, 2}:
+            raise ValueError(f"Score must be 0, 1, or 2 for graded scale, got {score}")
+            
+        self.score = score
+
+    def get_score(self) -> int:
+        """
+        Returns the validated score.
+
+        Returns:
+            The validated score.
+        """
+        return self.score

rre-dataset-generator/tests/unit/llm/__init__.py

@@ -0,0 +1,60 @@
+import pytest
+from langchain_core.language_models.fake_chat_models import FakeListChatModel
+from src.llm.llm_service import LLMService
+from src.model.document import Document
+from src.model.query_response import LLMQueryResponse
+
+
+@pytest.fixture
+def example_doc():
+    return Document(
+        id="doc1",
+        fields={
+            "title": "Car of the Year",
+            "description": "The Toyota Camry, the nation's most popular car has now been rated as its best new model."
+        }
+    )
+
+
+def test_llm_service_generate_queries_EXPECTED_valid(example_doc):
+    fake_llm = FakeListChatModel(responses=['["Toyota", "Best Car"]'])
+    service = LLMService(chat_model=fake_llm)
+    response = service.generate_queries(example_doc, 2)
+
+    assert isinstance(response, LLMQueryResponse)
+    assert response.get_queries() == ["Toyota", "Best Car"]
+
+
+def test_llm_service_generate_queries_EXPECTED_empty_list(example_doc):
+    fake_llm = FakeListChatModel(responses=['[]'])
+    service = LLMService(chat_model=fake_llm)
+    response = service.generate_queries(example_doc, 0)
+    assert response.get_queries() == []
+
+
+@pytest.mark.parametrize("invalid_response, expected_error", [
+    ('not a json', "Invalid JSON in `response_content`"),
+    ('["", " ", "Valid"]', "must not be empty or only whitespace"),
+    ('["Good", 123, null]', "must be strings"),
+])
+def test_llm_service_generate_queries_with_invalid_responses_EXPECTED_error(invalid_response, expected_error, example_doc):
+    fake_llm = FakeListChatModel(responses=[invalid_response])
+    service = LLMService(chat_model=fake_llm)
+    with pytest.raises(ValueError, match=expected_error):
+        service.generate_queries(example_doc, 3)
+
+
+def test_generate_queries_with_unicode_strings_EXPECTED_list_of_unicode_strings(example_doc):
+    unicode_list = '["こんにちは", "你好", "¡Hola!"]'
+    fake_llm = FakeListChatModel(responses=[unicode_list])
+    service = LLMService(chat_model=fake_llm)
+    response = service.generate_queries(example_doc, 3)
+    assert response.get_queries() == ["こんにちは", "你好", "¡Hola!"]
+
+
+def test_generate_queries_with_leading_trailing_whitespace_EXPECTED_strings_preserved(example_doc):
+    list_with_whitespace = '["  hello  ", " world "]'
+    fake_llm = FakeListChatModel(responses=[list_with_whitespace])
+    service = LLMService(chat_model=fake_llm)
+    response = service.generate_queries(example_doc, 2)
+    assert response.get_queries() == ["  hello  ", " world "]

rre-dataset-generator/tests/unit/llm/test_llm_service_queries.py

@@ -0,0 +1,61 @@
+import pytest
+import json
+from langchain_core.language_models.fake_chat_models import FakeListChatModel
+from src.llm.llm_service import LLMService
+from src.model.document import Document
+from src.model.score_response import LLMScoreResponse
+from src.model.query_response import LLMQueryResponse
+
+
+@pytest.fixture
+def example_doc():
+    return Document(
+        id="doc1",
+        fields={
+            "title": "Car of the Year",
+            "description": "The Toyota Camry, the nation's most popular car has now been rated as its best new model."
+        }
+    )
+
+
+@pytest.mark.parametrize("scale, valid_score", [
+    ("binary", 0),
+    ("binary", 1),
+    ("graded", 0),
+    ("graded", 1),
+    ("graded", 2),
+])
+def test_generate_score_with_valid_scale_EXPECTED_integer_score(scale, valid_score, example_doc):
+    fake_llm = FakeListChatModel(responses=[f'{{"score": {valid_score}}}'])
+    service = LLMService(chat_model=fake_llm)
+    query = "Is a Toyota the car of the year?"
+    response = service.generate_score(example_doc, query, relevance_scale=scale)
+    assert isinstance(response, LLMScoreResponse)
+    assert response.get_score() == valid_score
+
+
+@pytest.mark.parametrize("scale, response_json, expected_error", [
+    # Binary scale errors
+    ('binary', 'not a json', 'Invalid LLM response'),
+    ('binary', '{"not_score": 1}', 'Invalid LLM response'),
+    ('binary', '{"score": "one"}', 'Score must be 0 or 1 for binary scale, got one'),
+    ('binary', '{"score": 3}', 'Score must be 0 or 1 for binary scale, got 3'),
+    # Graded scale errors
+    ('graded', '{"score": -1}', 'Score must be 0, 1, or 2 for graded scale, got -1'),
+    ('graded', '{"score": 1.5}', 'Score must be 0, 1, or 2 for graded scale, got 1.5'),
+    ('graded', '{"score": null}', 'Score must be 0, 1, or 2 for graded scale, got None'),
+])
+def test_generate_score_with_invalid_llm_responses_EXPECTED_value_error(scale, response_json, expected_error, example_doc):
+    fake_llm = FakeListChatModel(responses=[response_json])
+    service = LLMService(chat_model=fake_llm)
+    query = "Is a Toyota the car of the year?"
+    with pytest.raises(ValueError, match=expected_error):
+        service.generate_score(example_doc, query, relevance_scale=scale)
+
+
+def test_generate_score_with_invalid_relevance_scale_EXPECTED_value_error(example_doc):
+    fake_llm = FakeListChatModel(responses=['{"score": 1}'])
+    service = LLMService(chat_model=fake_llm)
+    query = "What car won?"
+    with pytest.raises(ValueError, match="Invalid relevance scale"):
+        service.generate_score(example_doc, query, relevance_scale='fuzzy')