edsl 0.1.57__py3-none-any.whl → 0.1.59__py3-none-any.whl

edsl/dataset/dataset.py

@@ -93,6 +93,38 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
         """
         _, values = list(self.data[0].items())[0]
         return len(values)
+    
+    def drop(self, field_name):
+        """
+        Returns a new Dataset with the specified field removed.
+        
+        Args:
+            field_name (str): The name of the field to remove.
+            
+        Returns:
+            Dataset: A new Dataset instance without the specified field.
+            
+        Raises:
+            KeyError: If the field_name doesn't exist in the dataset.
+            
+        Examples:
+            >>> from .dataset import Dataset
+            >>> d = Dataset([{'a': [1, 2, 3]}, {'b': [4, 5, 6]}])
+            >>> d.drop('a')
+            Dataset([{'b': [4, 5, 6]}])
+            
+            >>> # Testing drop with nonexistent field raises DatasetKeyError - tested in unit tests
+        """
+        from .dataset import Dataset
+        
+        # Check if field exists in the dataset
+        if field_name not in self.relevant_columns():
+            raise DatasetKeyError(f"Field '{field_name}' not found in dataset")
+        
+        # Create a new dataset without the specified field
+        new_data = [entry for entry in self.data if field_name not in entry]
+        return Dataset(new_data)
+
 
     def tail(self, n: int = 5) -> Dataset:
         """Return the last n observations in the dataset.
@@ -1054,6 +1086,48 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
         
         return Dataset(new_data)
 
+    def unique(self) -> "Dataset":
+        """Return a new dataset with only unique observations.
+        
+        Examples:
+            >>> d = Dataset([{'a': [1, 2, 2, 3]}, {'b': [4, 5, 5, 6]}])
+            >>> d.unique().data
+            [{'a': [1, 2, 3]}, {'b': [4, 5, 6]}]
+            
+            >>> d = Dataset([{'x': ['a', 'a', 'b']}, {'y': [1, 1, 2]}])
+            >>> d.unique().data
+            [{'x': ['a', 'b']}, {'y': [1, 2]}]
+        """
+        # Get all column names and values
+        headers, data = self._tabular()
+        
+        # Create a list of unique rows
+        unique_rows = []
+        seen = set()
+        
+        for row in data:
+            # Convert the row to a hashable representation for comparison
+            # We need to handle potential unhashable types
+            try:
+                row_key = tuple(map(lambda x: str(x) if isinstance(x, (list, dict)) else x, row))
+                if row_key not in seen:
+                    seen.add(row_key)
+                    unique_rows.append(row)
+            except:
+                # Fallback for complex objects: compare based on string representation
+                row_str = str(row)
+                if row_str not in seen:
+                    seen.add(row_str)
+                    unique_rows.append(row)
+        
+        # Create a new dataset with unique combinations
+        new_data = []
+        for i, header in enumerate(headers):
+            values = [row[i] for row in unique_rows]
+            new_data.append({header: values})
+        
+        return Dataset(new_data)
+
 
 if __name__ == "__main__":
     import doctest

edsl/dataset/dataset_operations_mixin.py

@@ -1070,7 +1070,6 @@ class DataOperationsBase:
             - All dictionaries in the field must have compatible structures
             - If a dictionary is missing a key, the corresponding value will be None
             - Non-dictionary values in the field will cause a warning
-
         Examples:
             >>> from edsl.dataset import Dataset
 
@@ -1086,48 +1085,85 @@ class DataOperationsBase:
             >>> d = Dataset([{'a': [{'a': 1, 'b': 2}]}, {'c': [5]}])
             >>> d.flatten('a', keep_original=True)
             Dataset([{'a': [{'a': 1, 'b': 2}]}, {'c': [5]}, {'a.a': [1]}, {'a.b': [2]}])
+            
+            # Can also use unambiguous unprefixed field name
+            >>> result = Dataset([{'answer.pros_cons': [{'pros': ['Safety'], 'cons': ['Cost']}]}]).flatten('pros_cons')
+            >>> sorted(result.keys()) == ['answer.pros_cons.cons', 'answer.pros_cons.pros']
+            True
+            >>> sorted(result.to_dicts()[0].items()) == sorted({'cons': ['Cost'], 'pros': ['Safety']}.items())
+            True
         """
         from ..dataset import Dataset
 
         # Ensure the dataset isn't empty
         if not self.data:
             return self.copy()
-
-        # Find all columns that contain the field
-        matching_entries = []
-        for entry in self.data:
-            col_name = next(iter(entry.keys()))
-            if field == col_name or (
-                "." in col_name
-                and (col_name.endswith("." + field) or col_name.startswith(field + "."))
-            ):
-                matching_entries.append(entry)
-
-        # Check if the field is ambiguous
-        if len(matching_entries) > 1:
-            matching_cols = [next(iter(entry.keys())) for entry in matching_entries]
-            from .exceptions import DatasetValueError
-
-            raise DatasetValueError(
-                f"Ambiguous field name '{field}'. It matches multiple columns: {matching_cols}. "
-                f"Please specify the full column name to flatten."
-            )
-
-        # Get the number of observations
-        num_observations = self.num_observations()
-
-        # Find the column to flatten
+        
+        # First try direct match with the exact field name
         field_entry = None
         for entry in self.data:
-            if field in entry:
+            col_name = next(iter(entry.keys()))
+            if field == col_name:
                 field_entry = entry
                 break
+        
+        # If not found, try to match by unprefixed name
+        if field_entry is None:
+            # Find any columns that have field as their unprefixed name
+            candidates = []
+            for entry in self.data:
+                col_name = next(iter(entry.keys()))
+                if '.' in col_name:
+                    prefix, col_field = col_name.split('.', 1)
+                    if col_field == field:
+                        candidates.append(entry)
+            
+            # If we found exactly one match by unprefixed name, use it
+            if len(candidates) == 1:
+                field_entry = candidates[0]
+            # If we found multiple matches, it's ambiguous
+            elif len(candidates) > 1:
+                matching_cols = [next(iter(entry.keys())) for entry in candidates]
+                from .exceptions import DatasetValueError
+                raise DatasetValueError(
+                    f"Ambiguous field name '{field}'. It matches multiple columns: {matching_cols}. "
+                    f"Please specify the full column name to flatten."
+                )
+            # If no candidates by unprefixed name, check partial matches
+            else:
+                partial_matches = []
+                for entry in self.data:
+                    col_name = next(iter(entry.keys()))
+                    if '.' in col_name and (
+                        col_name.endswith('.' + field) or 
+                        col_name.startswith(field + '.')
+                    ):
+                        partial_matches.append(entry)
+                
+                # If we found exactly one partial match, use it
+                if len(partial_matches) == 1:
+                    field_entry = partial_matches[0]
+                # If we found multiple partial matches, it's ambiguous
+                elif len(partial_matches) > 1:
+                    matching_cols = [next(iter(entry.keys())) for entry in partial_matches]
+                    from .exceptions import DatasetValueError
+                    raise DatasetValueError(
+                        f"Ambiguous field name '{field}'. It matches multiple columns: {matching_cols}. "
+                        f"Please specify the full column name to flatten."
+                    )
+        
+        # Get the number of observations
+        num_observations = self.num_observations()
 
+        # If we still haven't found the field, it's not in the dataset
         if field_entry is None:
             warnings.warn(
                 f"Field '{field}' not found in dataset, returning original dataset"
             )
             return self.copy()
+        
+        # Get the actual field name as it appears in the data
+        actual_field = next(iter(field_entry.keys()))
 
         # Create new dictionary for flattened data
         flattened_data = []
@@ -1135,14 +1171,14 @@ class DataOperationsBase:
         # Copy all existing columns except the one we're flattening (if keep_original is False)
         for entry in self.data:
             col_name = next(iter(entry.keys()))
-            if col_name != field or keep_original:
+            if col_name != actual_field or keep_original:
                 flattened_data.append(entry.copy())
 
         # Get field data and make sure it's valid
-        field_values = field_entry[field]
+        field_values = field_entry[actual_field]
         if not all(isinstance(item, dict) for item in field_values if item is not None):
             warnings.warn(
-                f"Field '{field}' contains non-dictionary values that cannot be flattened"
+                f"Field '{actual_field}' contains non-dictionary values that cannot be flattened"
             )
             return self.copy()
 
@@ -1162,7 +1198,7 @@ class DataOperationsBase:
                 new_values.append(value)
 
             # Add this as a new column
-            flattened_data.append({f"{field}.{key}": new_values})
+            flattened_data.append({f"{actual_field}.{key}": new_values})
 
         # Return a new Dataset with the flattened data
         return Dataset(flattened_data)
@@ -1244,37 +1280,6 @@ class DataOperationsBase:
 
         return result
 
-    def drop(self, field_name):
-        """
-        Returns a new Dataset with the specified field removed.
-
-        Args:
-            field_name (str): The name of the field to remove.
-
-        Returns:
-            Dataset: A new Dataset instance without the specified field.
-
-        Raises:
-            KeyError: If the field_name doesn't exist in the dataset.
-
-        Examples:
-            >>> from .dataset import Dataset
-            >>> d = Dataset([{'a': [1, 2, 3]}, {'b': [4, 5, 6]}])
-            >>> d.drop('a')
-            Dataset([{'b': [4, 5, 6]}])
-
-            >>> # Testing drop with nonexistent field raises DatasetKeyError - tested in unit tests
-        """
-        from .dataset import Dataset
-
-        # Check if field exists in the dataset
-        if field_name not in self.relevant_columns():
-            raise DatasetKeyError(f"Field '{field_name}' not found in dataset")
-
-        # Create a new dataset without the specified field
-        new_data = [entry for entry in self.data if field_name not in entry]
-        return Dataset(new_data)
-
     def remove_prefix(self):
         """Returns a new Dataset with the prefix removed from all column names.
 

edsl/inference_services/services/test_service.py

@@ -54,7 +54,7 @@ class TestService(InferenceServiceABC):
             input_token_name = cls.input_token_name
             output_token_name = cls.output_token_name
             _rpm = 1000
-            _tpm = 100000
+            _tpm = 8000000
 
             @property
             def _canned_response(self):

edsl/interviews/exception_tracking.py

@@ -16,13 +16,19 @@ class InterviewExceptionEntry:
         invigilator: "InvigilatorBase",
         traceback_format="text",
         answers=None,
+        time=None,  # Added time parameter for deserialization
     ):
-        self.time = datetime.datetime.now().isoformat()
+        self.time = time or datetime.datetime.now().isoformat()
         self.exception = exception
         self.invigilator = invigilator
         self.traceback_format = traceback_format
         self.answers = answers
 
+    @property
+    def exception_type(self) -> str:
+        """Return the type of the exception."""
+        return type(self.exception).__name__
+
     @property
     def question_type(self) -> str:
         """Return the type of the question that failed."""
@@ -125,7 +131,12 @@ class InterviewExceptionEntry:
         'Traceback (most recent call last):...'
         """
         e = self.exception
-        tb_str = "".join(traceback.format_exception(type(e), e, e.__traceback__))
+        # Check if the exception has a traceback attribute
+        if hasattr(e, "__traceback__") and e.__traceback__:
+            tb_str = "".join(traceback.format_exception(type(e), e, e.__traceback__))
+        else:
+            # Use the message as traceback if no traceback available
+            tb_str = f"Exception: {str(e)}"
         return tb_str
 
     @property
@@ -139,14 +150,19 @@ class InterviewExceptionEntry:
 
         console = Console(file=html_output, record=True)
 
-        tb = Traceback.from_exception(
-            type(self.exception),
-            self.exception,
-            self.exception.__traceback__,
-            show_locals=True,
-        )
-        console.print(tb)
-        return html_output.getvalue()
+        # Check if the exception has a traceback attribute
+        if hasattr(self.exception, "__traceback__") and self.exception.__traceback__:
+            tb = Traceback.from_exception(
+                type(self.exception),
+                self.exception,
+                self.exception.__traceback__,
+                show_locals=True,
+            )
+            console.print(tb)
+            return html_output.getvalue()
+        else:
+            # Return a simple string if no traceback available
+            return f"<pre>Exception: {str(self.exception)}</pre>"
 
     @staticmethod
     def serialize_exception(exception: Exception) -> dict:
@@ -155,14 +171,25 @@ class InterviewExceptionEntry:
         >>> entry = InterviewExceptionEntry.example()
         >>> _ = entry.serialize_exception(entry.exception)
         """
-        return {
-            "type": type(exception).__name__,
-            "message": str(exception),
-            "traceback": "".join(
+        # Store the original exception type for proper reconstruction
+        exception_type = type(exception).__name__
+        module_name = getattr(type(exception), "__module__", "builtins")
+
+        # Extract traceback if available
+        if hasattr(exception, "__traceback__") and exception.__traceback__:
+            tb_str = "".join(
                 traceback.format_exception(
                     type(exception), exception, exception.__traceback__
                 )
-            ),
+            )
+        else:
+            tb_str = f"Exception: {str(exception)}"
+
+        return {
+            "type": exception_type,
+            "module": module_name,
+            "message": str(exception),
+            "traceback": tb_str,
         }
 
     @staticmethod
@@ -172,11 +199,31 @@ class InterviewExceptionEntry:
         >>> entry = InterviewExceptionEntry.example()
         >>> _ = entry.deserialize_exception(entry.to_dict()["exception"])
         """
+        exception_type = data.get("type", "Exception")
+        module_name = data.get("module", "builtins")
+        message = data.get("message", "")
+
         try:
-            exception_class = globals()[data["type"]]
-        except KeyError:
-            exception_class = Exception
-        return exception_class(data["message"])
+            # Try to import the module and get the exception class
+            # if module_name != "builtins":
+            #     import importlib
+
+            #     module = importlib.import_module(module_name)
+            #     exception_class = getattr(module, exception_type, Exception)
+            # else:
+            #     # Look for exception in builtins
+            import builtins
+
+            exception_class = getattr(builtins, exception_type, Exception)
+
+        except (ImportError, AttributeError):
+            # Fall back to a generic Exception but preserve the type name
+            exception = Exception(message)
+            exception.__class__.__name__ = exception_type
+            return exception
+
+        # Create instance of the original exception type if possible
+        return exception_class(message)
 
     def to_dict(self) -> dict:
         """Return the exception as a dictionary.
@@ -216,7 +263,11 @@ class InterviewExceptionEntry:
             invigilator = None
         else:
             invigilator = InvigilatorAI.from_dict(data["invigilator"])
-        return cls(exception=exception, invigilator=invigilator)
+
+        # Use the original timestamp from serialization
+        time = data.get("time")
+
+        return cls(exception=exception, invigilator=invigilator, time=time)
 
 
 class InterviewExceptionCollection(UserDict):
@@ -239,6 +290,27 @@ class InterviewExceptionCollection(UserDict):
         """Return the number of unfixed exceptions."""
         return sum(len(v) for v in self.unfixed_exceptions().values())
 
+    def list(self) -> list[dict]:
+        """
+        Return a list of exception dicts with the following metadata:
+        - exception_type: the type of the exception
+        - inference_service: the inference service used
+        - model: the model used
+        - question_name: the name of the question that failed
+        """
+        exception_list = []
+        for question_name, exceptions in self.data.items():
+            for exception in exceptions:
+                exception_list.append(
+                    {
+                        "exception_type": exception.exception_type,
+                        "inference_service": exception.invigilator.model._inference_service_,
+                        "model": exception.invigilator.model.model,
+                        "question_name": question_name,
+                    }
+                )
+        return exception_list
+
     def num_unfixed(self) -> int:
         """Return a list of unfixed questions."""
         return len([k for k in self.data.keys() if k not in self.fixed])

edsl/invigilators/invigilators.py

@@ -105,7 +105,11 @@ class InvigilatorBase(ABC):
             value = getattr(self, attr)
             if value is None:
                 return None
-            if hasattr(value, "to_dict"):
+            if attr == "scenario" and hasattr(value, "offload"):
+                # Use the scenario's offload method to replace base64_string values
+                offloaded = value.offload()
+                return offloaded.to_dict()
+            elif hasattr(value, "to_dict"):
                 return value.to_dict()
             if isinstance(value, (int, float, str, bool, dict, list)):
                 return value