edsl 0.1.49__py3-none-any.whl → 0.1.51__py3-none-any.whl

edsl/questions/question_multiple_choice.py

@@ -4,12 +4,34 @@ from typing import Union, Literal, Optional, List, Any
 from jinja2 import Template
 from pydantic import BaseModel, Field
 
-from ..scenarios import Scenario
 from .question_base import QuestionBase
 from .descriptors import QuestionOptionsDescriptor
 from .decorators import inject_exception
 from .response_validator_abc import ResponseValidatorABC
 
+class BaseMultipleChoiceResponse(BaseModel):
+    """
+    Base model for multiple choice responses.
+    
+    Attributes:
+        answer: The selected choice
+        comment: Optional comment field
+        generated_tokens: Optional raw tokens generated by the model
+    """
+    answer: Any = Field(..., description="Selected choice")
+    comment: Optional[str] = Field(None, description="Optional comment field")
+    generated_tokens: Optional[Any] = Field(None, description="Generated tokens")
+    
+    """
+    Examples:
+        >>> model = BaseMultipleChoiceResponse(answer="Option A", comment="My reasoning")
+        >>> model.answer
+        'Option A'
+        >>> model.comment
+        'My reasoning'
+    """
+
+
 def create_response_model(choices: List[str], permissive: bool = False):
     """
     Create a ChoiceResponse model class with a predefined list of choices.
@@ -17,95 +39,238 @@ def create_response_model(choices: List[str], permissive: bool = False):
     :param choices: A list of allowed values for the answer field.
     :param permissive: If True, any value will be accepted as an answer.
     :return: A new Pydantic model class.
+    
+    Examples:
+        >>> model = create_response_model(["Red", "Green", "Blue"], permissive=False)
+        >>> response = model(answer="Red")
+        >>> response.answer
+        'Red'
+        
+        >>> try:
+        ...     model(answer="Purple")
+        ... except Exception:
+        ...     print("Invalid value")
+        Invalid value
+        
+        >>> permissive_model = create_response_model(["Red", "Green", "Blue"], permissive=True)
+        >>> response = permissive_model(answer="Purple")
+        >>> response.answer
+        'Purple'
     """
     choice_tuple = tuple(choices)
 
     if not permissive:
-
-        class ChoiceResponse(BaseModel):
+        class ChoiceResponse(BaseMultipleChoiceResponse):
+            """
+            A model for multiple choice responses with strict validation.
+            
+            Attributes:
+                answer: Must be one of the predefined choices
+                
+            Examples:
+                >>> choices = ["Option A", "Option B", "Option C"]
+                >>> model = create_response_model(choices, permissive=False)
+                >>> response = model(answer="Option A")
+                >>> response.answer
+                'Option A'
+            """
             answer: Literal[choice_tuple] = Field(description="Selected choice")
-            comment: Optional[str] = Field(None, description="Optional comment field")
-            generated_tokens: Optional[Any] = Field(
-                None, description="Generated tokens"
-            )
-
-            class Config:
-                @staticmethod
-                def json_schema_extra(schema: dict, model: BaseModel) -> None:
-                    for prop in schema.get("properties", {}).values():
-                        if prop.get("title") == "answer":
-                            prop["enum"] = choices
 
+            model_config = {
+                "json_schema_extra": {
+                    "properties": {
+                        "answer": {
+                            "enum": choices
+                        }
+                    }
+                }
+            }
     else:
+        class ChoiceResponse(BaseMultipleChoiceResponse):
+            """
+            A model for multiple choice responses with permissive validation.
+            
+            Attributes:
+                answer: Can be any value, with suggested choices provided
+                
+            Examples:
+                >>> choices = ["Option A", "Option B", "Option C"]
+                >>> model = create_response_model(choices, permissive=True)
+                >>> response = model(answer="Something Else")
+                >>> response.answer
+                'Something Else'
+            """
+            answer: Any = Field(description=f"Selected choice (can be any value). Suggested choices are: {choices}")
 
-        class ChoiceResponse(BaseModel):
-            answer: Any = Field(description="Selected choice (can be any value)")
-            comment: Optional[str] = Field(None, description="Optional comment field")
-            generated_tokens: Optional[Any] = Field(
-                None, description="Generated tokens"
-            )
-
-            class Config:
-                @staticmethod
-                def json_schema_extra(schema: dict, model: BaseModel) -> None:
-                    for prop in schema.get("properties", {}).values():
-                        if prop.get("title") == "answer":
-                            prop["description"] += f". Suggested choices are: {choices}"
-                    schema["title"] += " (Permissive)"
+            model_config = {
+                "title": "PermissiveChoiceResponse"
+            }
 
     return ChoiceResponse
 
 
 class MultipleChoiceResponseValidator(ResponseValidatorABC):
+    """
+    Validator for multiple choice responses.
+    
+    This validator ensures that the answer is one of the allowed options.
+    In permissive mode, any answer is accepted.
+    
+    Examples:
+        >>> from edsl.questions import QuestionMultipleChoice
+        >>> q = QuestionMultipleChoice(
+        ...     question_name="feeling",
+        ...     question_text="How are you feeling?",
+        ...     question_options=["Good", "Bad", "Neutral"]
+        ... )
+        >>> validator = q.response_validator
+        >>> result = validator.validate({"answer": "Good"})
+        >>> sorted(result.keys())
+        ['answer', 'comment', 'generated_tokens']
+        >>> result["answer"]
+        'Good'
+    """
     required_params = ["question_options", "use_code"]
 
     def fix(self, response, verbose=False):
-        response_text = str(response.get("answer"))
-        if response_text is None:
-            response_text = response.get("generated_tokens", "")
+        """
+        Attempt to fix an invalid multiple choice response.
+        
+        Strategies:
+        1. Extract an option mentioned in the generated text
+        2. Check for exact matches in the text
+        3. Look for substring matches
+        4. Normalize whitespace and check for matches
+        5. Check if the answer is a prefix of any option (ignoring trailing spaces/punctuation)
+        
+        Parameters:
+            response: The invalid response to fix
+            verbose: Whether to print debug information
+            
+        Returns:
+            A fixed response dict if possible, otherwise the original response
+            
+        Examples:
+            >>> from edsl.questions import QuestionMultipleChoice
+            >>> q = QuestionMultipleChoice.example()
+            >>> validator = q.response_validator
+            >>> result = validator.fix({"answer": "I'm feeling Good today"})
+            >>> sorted(result.keys())
+            ['answer', 'comment', 'generated_tokens']
+            >>> result["answer"]
+            'Good'
+        """
+        # Don't attempt to fix None values - they should be properly rejected
+        if response.get("answer") is None:
+            if verbose:
+                print("Not attempting to fix None answer value")
+            return response
+            
+        # Get the raw text to analyze
+        response_text = str(response.get("answer", ""))
+        if not response_text:
+            response_text = str(response.get("generated_tokens", ""))
 
         if verbose:
-            print(f"Invalid generated tokens was: {response_text}")
+            print(f"Invalid response text: {response_text}")
+            print(f"Looking for options among: {self.question_options}")
 
+        # Strategy 1: Look for exact options in the text
         matches = []
-        for idx, option in enumerate(self.question_options):
-            if verbose:
-                print("The options are: ", self.question_options)
-            if str(option) in response_text:
+        for option in self.question_options:
+            option_str = str(option)
+            if option_str in response_text:
                 if verbose:
-                    print("Match found with option ", option)
+                    print(f"Match found with option: {option_str}")
                 if option not in matches:
                     matches.append(option)
 
-        if verbose:
-            print("The matches are: ", matches)
+        # If we have exactly one match, use it
         if len(matches) == 1:
+            fixed_answer = matches[0]
             proposed_data = {
-                "answer": matches[0],
-                "generated_tokens": response.get("generated_tokens", None),
+                "answer": fixed_answer,
+                "comment": response.get("comment"),
+                "generated_tokens": response.get("generated_tokens"),
             }
+            
             try:
-                self.response_model(**proposed_data)
+                # Validate the fixed answer
+                self.response_model.model_validate(proposed_data)
+                if verbose:
+                    print(f"Fixed answer: {fixed_answer}")
                 return proposed_data
             except Exception as e:
                 if verbose:
-                    print(f"Proposed solution {proposed_data} is invalid. Error: {e}")
-            return response
+                    print(f"Validation failed for fixed answer: {e}")
+        
+        # Strategy 2: Check if the answer is a match when normalized (strip whitespace)
+        response_text_normalized = response_text.strip()
+        for option in self.question_options:
+            option_str = str(option).strip()
+            if option_str == response_text_normalized:
+                if verbose:
+                    print(f"Normalized match found with option: {option}")
+                proposed_data = {
+                    "answer": option,  # Use the exact option from the list
+                    "comment": response.get("comment"),
+                    "generated_tokens": response.get("generated_tokens"),
+                }
+                try:
+                    self.response_model.model_validate(proposed_data)
+                    if verbose:
+                        print(f"Fixed answer with normalization: {option}")
+                    return proposed_data
+                except Exception as e:
+                    if verbose:
+                        print(f"Validation failed for normalized answer: {e}")
+        
+        # Strategy 3: Check if the answer is a prefix of any option
+        # This handles cases where the model returns a partial answer
+        # Only apply this strategy if we have a meaningful response text
+        if response_text_normalized and not response_text_normalized.lower() == "none":
+            for option in self.question_options:
+                option_str = str(option).strip()
+                if option_str.startswith(response_text_normalized) or response_text_normalized.startswith(option_str):
+                    if verbose:
+                        print(f"Prefix match found with option: {option}")
+                    proposed_data = {
+                        "answer": option,  # Use the exact option from the list
+                        "comment": response.get("comment"),
+                        "generated_tokens": response.get("generated_tokens"),
+                    }
+                    try:
+                        self.response_model.model_validate(proposed_data)
+                        if verbose:
+                            print(f"Fixed answer with prefix matching: {option}")
+                        return proposed_data
+                    except Exception as e:
+                        if verbose:
+                            print(f"Validation failed for prefix answer: {e}")
+        
+        # If multiple or no matches, return original response
+        if verbose:
+            if len(matches) > 1:
+                print(f"Multiple matches found: {matches}, cannot determine correct option")
+            else:
+                print("No matches found in response text")
+        
+        return response
 
     valid_examples = [
-        ({"answer": 1}, {"question_options": ["Good", "Great", "OK", "Bad"]})
+        ({"answer": "Good"}, {"question_options": ["Good", "Great", "OK", "Bad"]})
     ]
 
     invalid_examples = [
         (
-            {"answer": -1},
+            {"answer": "Terrible"},
             {"question_options": ["Good", "Great", "OK", "Bad"]},
-            "Answer code must be a non-negative integer",
+            "Value error, Permitted values are 'Good', 'Great', 'OK', 'Bad'",
         ),
         (
             {"answer": None},
             {"question_options": ["Good", "Great", "OK", "Bad"]},
-            "Answer code must not be missing.",
+            "Answer must not be null",
         ),
     ]
 
@@ -314,7 +479,8 @@ class QuestionMultipleChoice(QuestionBase):
 
             if potential_replacement is None:
                 # Nope - maybe it's in the substition dict?
-                raise ValueError(
+                from .exceptions import QuestionValueError
+                raise QuestionValueError(
                     f"Could not find the key '{question_option_key}' in the scenario."
                     f"The substition dict was: '{substitution_dict}.'"
                     f"The question options were: '{question_options}'."
@@ -353,11 +519,13 @@ class QuestionMultipleChoice(QuestionBase):
             try:
                 return translated_options[int(answer_code)]
             except IndexError:
-                raise ValueError(
+                from .exceptions import QuestionValueError
+                raise QuestionValueError(
                     f"Answer code is out of range. The answer code index was: {int(answer_code)}. The options were: {translated_options}."
                 )
             except TypeError:
-                raise ValueError(
+                from .exceptions import QuestionValueError
+                raise QuestionValueError(
                     f"The answer code was: '{answer_code}.'",
                     f"The options were: '{translated_options}'.",
                 )