edsl 0.1.14__py3-none-any.whl → 0.1.40__py3-none-any.whl

edsl/jobs/tasks/TaskHistory.py

@@ -0,0 +1,470 @@
+from typing import List, Optional
+from io import BytesIO
+import base64
+from edsl.jobs.tasks.task_status_enum import TaskStatus
+from edsl.Base import RepresentationMixin
+
+
+class TaskHistory(RepresentationMixin):
+    def __init__(
+        self,
+        interviews: List["Interview"] = None,
+        include_traceback: bool = False,
+        max_interviews: int = 10,
+        interviews_with_exceptions_only: bool = False,
+    ):
+        """
+        The structure of a TaskHistory exception
+
+        [Interview.exceptions, Interview.exceptions, Interview.exceptions, ...]
+
+        >>> _ = TaskHistory.example()
+        ...
+        """
+        self.interviews_with_exceptions_only = interviews_with_exceptions_only
+        self._interviews = {}
+        self.total_interviews = []
+        if interviews is not None:
+            for interview in interviews:
+                self.add_interview(interview)
+
+        self.include_traceback = include_traceback
+        self._interviews = {
+            index: interview for index, interview in enumerate(self.total_interviews)
+        }
+        self.max_interviews = max_interviews
+
+        # self.total_interviews = interviews
+        self.include_traceback = include_traceback
+
+        # self._interviews = {index: i for index, i in enumerate(self.total_interviews)}
+        self.max_interviews = max_interviews
+
+    def add_interview(self, interview: "Interview"):
+        """Add a single interview to the history"""
+        if self.interviews_with_exceptions_only and interview.exceptions == {}:
+            return
+
+        self.total_interviews.append(interview)
+        self._interviews[len(self._interviews)] = interview
+
+    @classmethod
+    def example(cls):
+        """ """
+        from edsl.jobs.interviews.Interview import Interview
+
+        from edsl.jobs.Jobs import Jobs
+
+        j = Jobs.example(throw_exception_probability=1, test_model=True)
+
+        from edsl.config import CONFIG
+
+        results = j.run(
+            print_exceptions=False,
+            skip_retry=True,
+            cache=False,
+            raise_validation_errors=True,
+            disable_remote_cache=True,
+            disable_remote_inference=True,
+        )
+
+        return cls(results.task_history.total_interviews)
+
+    @property
+    def exceptions(self):
+        """
+        >>> len(TaskHistory.example().exceptions)
+        4
+        """
+        return [i.exceptions for k, i in self._interviews.items() if i.exceptions != {}]
+
+    @property
+    def unfixed_exceptions(self):
+        """
+        >>> len(TaskHistory.example().unfixed_exceptions)
+        4
+        """
+        return [
+            i.exceptions
+            for k, i in self._interviews.items()
+            if i.exceptions.num_unfixed() > 0
+        ]
+
+    @property
+    def indices(self):
+        return [k for k, i in self._interviews.items() if i.exceptions != {}]
+
+    def __repr__(self):
+        """Return a string representation of the TaskHistory."""
+        return f"TaskHistory(interviews={self.total_interviews})."
+
+    def to_dict(self, add_edsl_version=True):
+        """Return the TaskHistory as a dictionary."""
+        d = {
+            "interviews": [
+                i.to_dict(add_edsl_version=add_edsl_version)
+                for i in self.total_interviews
+            ],
+            "include_traceback": self.include_traceback,
+        }
+        if add_edsl_version:
+            from edsl import __version__
+
+            d["edsl_version"] = __version__
+            d["edsl_class_name"] = "TaskHistory"
+        return d
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        """Create a TaskHistory from a dictionary."""
+        if data is None:
+            return cls([], include_traceback=False)
+
+        from edsl.jobs.interviews.Interview import Interview
+
+        interviews = [Interview.from_dict(i) for i in data["interviews"]]
+        return cls(interviews, include_traceback=data["include_traceback"])
+
+    @property
+    def has_exceptions(self) -> bool:
+        """Return True if there are any exceptions.
+
+        >>> TaskHistory.example().has_exceptions
+        True
+
+        """
+        return len(self.exceptions) > 0
+
+    @property
+    def has_unfixed_exceptions(self) -> bool:
+        """Return True if there are any exceptions."""
+        return len(self.unfixed_exceptions) > 0
+
+    def show_exceptions(self, tracebacks=False):
+        """Print the exceptions."""
+        for index in self.indices:
+            self.total_interviews[index].exceptions.print(tracebacks)
+
+    def get_updates(self):
+        """Return a list of all the updates."""
+        updates = []
+        for interview in self.total_interviews:
+            for question_name, logs in interview.task_status_logs.items():
+                updates.append(logs)
+        return updates
+
+    def print(self):
+        from rich import print
+
+        print(self.get_updates())
+
+    def plot_completion_times(self):
+        """Plot the completion times for each task."""
+        import matplotlib.pyplot as plt
+
+        updates = self.get_updates()
+
+        elapsed = [update.max_time - update.min_time for update in updates]
+        for i, update in enumerate(updates):
+            if update[-1]["value"] != TaskStatus.SUCCESS:
+                elapsed[i] = 0
+        x = range(len(elapsed))
+        y = elapsed
+
+        plt.bar(x, y)
+        plt.title("Per-interview completion times")
+        plt.xlabel("Task")
+        plt.ylabel("Time (seconds)")
+        plt.show()
+
+    def plotting_data(self, num_periods=100):
+        updates = self.get_updates()
+
+        min_t = min([update.min_time for update in updates])
+        max_t = max([update.max_time for update in updates])
+        delta_t = (max_t - min_t) / (num_periods * 1.0)
+        time_periods = [min_t + delta_t * i for i in range(num_periods)]
+
+        def counts(t):
+            d = {}
+            for update in updates:
+                status = update.status_at_time(t)
+                if status in d:
+                    d[status] += 1
+                else:
+                    d[status] = 1
+            return d
+
+        status_counts = [counts(t) for t in time_periods]
+
+        new_counts = []
+        for status_count in status_counts:
+            d = {task_status: 0 for task_status in TaskStatus}
+            d.update(status_count)
+            new_counts.append(d)
+
+        return new_counts
+
+    def plot(self, num_periods=100, get_embedded_html=False):
+        """Plot the number of tasks in each state over time."""
+        new_counts = self.plotting_data(num_periods)
+        max_count = max([max(entry.values()) for entry in new_counts])
+
+        rows = int(len(TaskStatus) ** 0.5) + 1
+        cols = (len(TaskStatus) + rows - 1) // rows  # Ensure all plots fit
+
+        import matplotlib.pyplot as plt
+
+        fig, axes = plt.subplots(rows, cols, figsize=(15, 10))
+        axes = axes.flatten()  # Flatten in case of a single row/column
+
+        for i, status in enumerate(TaskStatus):
+            ax = axes[i]
+            x = range(len(new_counts))
+            y = [
+                item.get(status, 0) for item in new_counts
+            ]  # Use .get() to handle missing keys safely
+            ax.plot(x, y, marker="o", linestyle="-")
+            ax.set_title(status.name)
+            ax.set_xlabel("Time Periods")
+            ax.set_ylabel("Count")
+            ax.grid(True)
+            ax.set_ylim(0, max_count)
+
+        # Hide any unused subplots
+        for ax in axes[len(TaskStatus) :]:
+            ax.axis("off")
+
+        plt.tight_layout()
+
+        if get_embedded_html:
+            buffer = BytesIO()
+            fig.savefig(buffer, format="png")
+            plt.close(fig)
+            buffer.seek(0)
+
+            # Encode plot to base64 string
+            img_data = base64.b64encode(buffer.getvalue()).decode("utf-8")
+            buffer.close()
+            return f'<img src="data:image/png;base64,{img_data}" alt="Plot">'
+        else:
+            plt.show()
+
+    def css(self):
+        from importlib import resources
+
+        env = resources.files("edsl").joinpath("templates/error_reporting")
+        css = env.joinpath("report.css").read_text()
+        return css
+
+    def javascript(self):
+        from importlib import resources
+
+        env = resources.files("edsl").joinpath("templates/error_reporting")
+        js = env.joinpath("report.js").read_text()
+        return js
+
+    @property
+    def exceptions_by_type(self) -> dict:
+        """Return a dictionary of exceptions by type."""
+        exceptions_by_type = {}
+        for interview in self.total_interviews:
+            for question_name, exceptions in interview.exceptions.items():
+                for exception in exceptions:
+                    exception_type = exception.exception.__class__.__name__
+                    if exception_type in exceptions_by_type:
+                        exceptions_by_type[exception_type] += 1
+                    else:
+                        exceptions_by_type[exception_type] = 1
+        return exceptions_by_type
+
+    @property
+    def exceptions_by_service(self) -> dict:
+        """Return a dictionary of exceptions tallied by service."""
+        exceptions_by_service = {}
+        for interview in self.total_interviews:
+            service = interview.model._inference_service_
+            if service not in exceptions_by_service:
+                exceptions_by_service[service] = 0
+            if interview.exceptions != {}:
+                exceptions_by_service[service] += len(interview.exceptions)
+        return exceptions_by_service
+
+    @property
+    def exceptions_by_question_name(self) -> dict:
+        """Return a dictionary of exceptions tallied by question name."""
+        exceptions_by_question_name = {}
+        for interview in self.total_interviews:
+            for question_name, exceptions in interview.exceptions.items():
+                question_type = interview.survey._get_question_by_name(
+                    question_name
+                ).question_type
+                if (question_name, question_type) not in exceptions_by_question_name:
+                    exceptions_by_question_name[(question_name, question_type)] = 0
+                exceptions_by_question_name[(question_name, question_type)] += len(
+                    exceptions
+                )
+
+        for question in self.total_interviews[0].survey.questions:
+            if (
+                question.question_name,
+                question.question_type,
+            ) not in exceptions_by_question_name:
+                exceptions_by_question_name[
+                    (question.question_name, question.question_type)
+                ] = 0
+
+        sorted_exceptions_by_question_name = {
+            k: v
+            for k, v in sorted(
+                exceptions_by_question_name.items(),
+                key=lambda item: item[1],
+                reverse=True,
+            )
+        }
+        return sorted_exceptions_by_question_name
+
+    @property
+    def exceptions_by_model(self) -> dict:
+        """Return a dictionary of exceptions tallied by model and question name."""
+        exceptions_by_model = {}
+        for interview in self.total_interviews:
+            model = interview.model.model
+            service = interview.model._inference_service_
+            if (service, model) not in exceptions_by_model:
+                exceptions_by_model[(service, model)] = 0
+            if interview.exceptions != {}:
+                exceptions_by_model[(service, model)] += len(interview.exceptions)
+
+        # sort the exceptions by model
+        sorted_exceptions_by_model = {
+            k: v
+            for k, v in sorted(
+                exceptions_by_model.items(), key=lambda item: item[1], reverse=True
+            )
+        }
+        return sorted_exceptions_by_model
+
+    def generate_html_report(self, css: Optional[str], include_plot=False):
+        if include_plot:
+            performance_plot_html = self.plot(num_periods=100, get_embedded_html=True)
+        else:
+            performance_plot_html = ""
+
+        if css is None:
+            css = self.css()
+
+        models_used = set([i.model.model for index, i in self._interviews.items()])
+
+        from jinja2 import Environment, FileSystemLoader
+        from edsl.TemplateLoader import TemplateLoader
+
+        env = Environment(loader=TemplateLoader("edsl", "templates/error_reporting"))
+
+        # Get current memory usage at this point
+
+        template = env.get_template("base.html")
+
+        # Render the template with data
+        output = template.render(
+            interviews=self._interviews,
+            css=css,
+            javascript=self.javascript(),
+            num_exceptions=len(self.exceptions),
+            performance_plot_html=performance_plot_html,
+            exceptions_by_type=self.exceptions_by_type,
+            exceptions_by_question_name=self.exceptions_by_question_name,
+            exceptions_by_model=self.exceptions_by_model,
+            exceptions_by_service=self.exceptions_by_service,
+            models_used=models_used,
+            max_interviews=self.max_interviews,
+        )
+        return output
+
+    def html(
+        self,
+        filename: Optional[str] = None,
+        return_link=False,
+        css=None,
+        cta="Open Report in New Tab",
+        open_in_browser=False,
+    ):
+        """Return an HTML report."""
+
+        from IPython.display import display, HTML
+        import tempfile
+        import os
+        from edsl.utilities.utilities import is_notebook
+
+        output = self.generate_html_report(css)
+
+        # Save the rendered output to a file
+        with open("output.html", "w") as f:
+            f.write(output)
+
+        if filename is None:
+            current_directory = os.getcwd()
+            filename = tempfile.NamedTemporaryFile(
+                "w", delete=False, suffix=".html", dir=current_directory
+            ).name
+
+        with open(filename, "w") as f:
+            with open(filename, "w") as f:
+                f.write(output)
+
+        if is_notebook():
+            import html
+
+            html_url = f"/files/{filename}"
+            html_link = f'<a href="{html_url}" target="_blank">{cta}</a>'
+            display(HTML(html_link))
+            escaped_output = html.escape(output)
+            iframe = f""""
+            <iframe srcdoc="{ escaped_output }" style="width: 800px; height: 600px;"></iframe>
+            """
+            display(HTML(iframe))
+        else:
+            print(f"Exception report saved to {filename}")
+
+        if open_in_browser:
+            import webbrowser
+
+            webbrowser.open(f"file://{os.path.abspath(filename)}")
+
+        if return_link:
+            return filename
+
+    def notebook(self):
+        """Create a notebook with the HTML content embedded in the first cell, then delete the cell content while keeping the output."""
+        from nbformat import v4 as nbf
+        from nbconvert.preprocessors import ExecutePreprocessor
+        import nbformat
+        import os
+
+        # Use the existing html method to generate the HTML content
+        output_html = self.generate_html_report(css=None)
+        nb = nbf.new_notebook()
+
+        # Add a code cell that renders the HTML content
+        code_cell = nbf.new_code_cell(
+            f"""
+    from IPython.display import HTML, display
+    display(HTML('''{output_html}'''))
+            """
+        )
+        nb.cells.append(code_cell)
+
+        # Execute the notebook
+        ep = ExecutePreprocessor(timeout=600, kernel_name="python3")
+        ep.preprocess(nb, {"metadata": {"path": os.getcwd()}})
+
+        # After execution, clear the cell's source code
+        nb.cells[0].source = ""
+
+        return nb
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)

edsl/jobs/tasks/TaskStatusLog.py

@@ -0,0 +1,23 @@
+from collections import UserList
+
+
+class TaskStatusLog(UserList):
+    """A list of TaskStatusEntry objects."""
+
+    @property
+    def min_time(self):
+        return self[0]["log_time"]
+
+    @property
+    def max_time(self):
+        return self[-1]["log_time"]
+
+    def status_at_time(self, t):
+        """Return the status at time t.
+
+        TODO: Could re-factor with bisect to make this faster.
+        """
+        for entry in self:
+            if entry["log_time"] > t:
+                return entry["value"]
+        return self[-1]["value"]

edsl/jobs/tasks/task_status_enum.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+from collections import UserDict
+import enum
+import time
+
+
+class TaskStatus(enum.Enum):
+    "These are the possible states a task can be in."
+    NOT_STARTED = enum.auto()
+    WAITING_FOR_DEPENDENCIES = enum.auto()
+    CANCELLED = enum.auto()
+    PARENT_FAILED = enum.auto()
+    WAITING_FOR_REQUEST_CAPACITY = enum.auto()
+    WAITING_FOR_TOKEN_CAPACITY = enum.auto()
+    API_CALL_IN_PROGRESS = enum.auto()
+    SUCCESS = enum.auto()
+    FAILED = enum.auto()
+
+
+class TaskStatusLogEntry(UserDict):
+    def __init__(self, log_time, value):
+        self.data = {"log_time": log_time, "value": value}
+        super().__init__(self.data)
+
+
+class TaskStatusDescriptor:
+    "The descriptor ensures that the task status is always an instance of the TaskStatus enum."
+
+    def __init__(self):
+        self._task_status = None
+
+    def __get__(self, instance, owner):
+        return self._task_status
+
+    def __set__(self, instance, value):
+        """Ensure that the value is an instance of TaskStatus."""
+        if not isinstance(value, TaskStatus):
+            raise ValueError("Value must be an instance of TaskStatus enum")
+        t = time.monotonic()
+        if hasattr(instance, "status_log"):
+            instance.status_log.append(TaskStatusLogEntry(t, value))
+        self._task_status = value
+
+    def __delete__(self, instance):
+        self._task_status = None
+
+
+status_colors = {
+    TaskStatus.NOT_STARTED: "grey",
+    TaskStatus.WAITING_FOR_DEPENDENCIES: "orange",
+    TaskStatus.WAITING_FOR_REQUEST_CAPACITY: "yellow",
+    TaskStatus.WAITING_FOR_TOKEN_CAPACITY: "gold",
+    TaskStatus.CANCELLED: "white",
+    TaskStatus.PARENT_FAILED: "darkred",
+    TaskStatus.FAILED: "red",
+    TaskStatus.API_CALL_IN_PROGRESS: "blue",
+    TaskStatus.SUCCESS: "green",
+}
+
+
+def get_enum_from_string(str_key):
+    """Parse the string to extract the enum member name."""
+    try:
+        _, member_name = str_key.split(".")
+        enum_member = getattr(TaskStatus, member_name)
+        return enum_member
+    except ValueError:
+        return str_key
+
+
+class InterviewTaskLogDict(UserDict):
+    """A dictionary of TaskStatusLog objects.
+
+    The key is the name of the task.
+    """
+
+    @property
+    def min_time(self):
+        return min([log.min_time for log in self.values()])
+
+    @property
+    def max_time(self):
+        return max([log.max_time for log in self.values()])
+
+    def status_matrix(self, num_periods):
+        """Return a matrix of status values."""
+        start_time = self.min_time
+        end_time = self.max_time
+        time_increment = (end_time - start_time) / num_periods
+        status_matrix = {}
+        time_periods = [start_time + i * time_increment for i in range(num_periods)]
+        for task_name, log in self.items():
+            status_matrix[task_name] = [log.status_at_time(t) for t in time_periods]
+        return status_matrix
+
+    def numerical_matrix(self, num_periods):
+        """Return a numerical matrix of status values."""
+        status_dicts = self.status_matrix(num_periods)
+
+        num_cols = num_periods
+        num_rows = len(status_dicts)
+        matrix = [[0 for _ in range(num_cols)] for _ in range(num_rows)]
+
+        for row_index, (task_name, status_list) in enumerate(status_dicts.items()):
+            matrix[row_index] = [
+                list(status_colors.keys()).index(status) for status in status_list
+            ]
+
+        index_to_names = {i: name for i, name in enumerate(status_dicts.keys())}
+        return matrix, index_to_names
+
+    def visualize(self, num_periods=10):
+        """Visualize the status matrix with outlined squares."""
+        import matplotlib.pyplot as plt
+        from matplotlib.colors import ListedColormap
+        import numpy as np
+        from matplotlib.patches import Rectangle
+
+        # Define your custom colormap
+        custom_cmap = ListedColormap(list(status_colors.values()))
+
+        # Generate the matrix
+        matrix, index_to_names = self.numerical_matrix(num_periods)
+
+        # Create the figure and axes
+        plt.figure(figsize=(10, 5))
+        ax = plt.gca()
+
+        # Display the matrix and keep a reference to the imshow object
+        im = ax.imshow(matrix, aspect="auto", cmap=custom_cmap)
+
+        # Adding color bar, now correctly associating it with 'im'
+        cbar = plt.colorbar(im, ticks=range(len(status_colors)), label="Task Status")
+
+        cbar_labels = [status.name for status in status_colors.keys()]
+        # breakpoint()
+        cbar.set_ticklabels(cbar_labels)  # Setting the custom labels for the colorbar
+
+        im.set_clim(
+            -0.5, len(status_colors) - 0.5
+        )  # Setting color limits directly on the imshow object
+
+        # Outline each cell by drawing rectangles
+        for (j, i), val in np.ndenumerate(matrix):
+            ax.add_patch(
+                Rectangle(
+                    (i - 0.5, j - 0.5), 1, 1, fill=False, edgecolor="black", lw=0.5
+                )
+            )
+
+        # Set custom y-axis ticks and labels
+        yticks = list(index_to_names.keys())
+        yticklabels = list(index_to_names.values())
+        plt.yticks(ticks=yticks, labels=yticklabels)
+
+        # Show the plot
+        plt.show()
+
+
+if __name__ == "__main__":
+    pass

edsl/jobs/tokens/InterviewTokenUsage.py

@@ -0,0 +1,27 @@
+from edsl.jobs.tokens.TokenUsage import TokenUsage
+from edsl.enums import TokenPricing
+
+
+class InterviewTokenUsage:
+    def __init__(
+        self, new_token_usage: TokenUsage = None, cached_token_usage: TokenUsage = None
+    ):
+        self.new_token_usage = new_token_usage or TokenUsage(from_cache=False)
+        self.cached_token_usage = cached_token_usage or TokenUsage(from_cache=True)
+
+    def __add__(self, other):
+        if not isinstance(other, InterviewTokenUsage):
+            raise ValueError(f"Can't add {type(other)} to InterviewTokenSummary")
+        return InterviewTokenUsage(
+            new_token_usage=self.new_token_usage + other.new_token_usage,
+            cached_token_usage=self.cached_token_usage + other.cached_token_usage,
+        )
+
+    def __repr__(self):
+        return f"InterviewTokenUsage(new_token_usage={self.new_token_usage}, cached_token_usage={self.cached_token_usage})"
+
+    def cost(self, prices: TokenPricing):
+        return self.new_token_usage.cost(prices)
+
+    def saved(self, prices: TokenPricing):
+        return self.cached_token_usage.cost(prices)