edsl 0.1.29.dev3__py3-none-any.whl → 0.1.30__py3-none-any.whl

edsl/jobs/buckets/TokenBucket.py

@@ -1,8 +1,6 @@
 from typing import Union, List, Any
 import asyncio
 import time
-from collections import UserDict
-from matplotlib import pyplot as plt
 
 
 class TokenBucket:
@@ -19,11 +17,29 @@ class TokenBucket:
         self.bucket_name = bucket_name
         self.bucket_type = bucket_type
         self.capacity = capacity  # Maximum number of tokens
+        self._old_capacity = capacity
         self.tokens = capacity  # Current number of available tokens
         self.refill_rate = refill_rate  # Rate at which tokens are refilled
+        self._old_refill_rate = refill_rate
         self.last_refill = time.monotonic()  # Last refill time
-
         self.log: List[Any] = []
+        self.turbo_mode = False
+
+    def turbo_mode_on(self):
+        """Set the refill rate to infinity."""
+        if self.turbo_mode:
+            pass
+        else:
+            # pass
+            self.turbo_mode = True
+            self.capacity = float("inf")
+            self.refill_rate = float("inf")
+
+    def turbo_mode_off(self):
+        """Restore the refill rate to its original value."""
+        self.turbo_mode = False
+        self.capacity = self._old_capacity
+        self.refill_rate = self._old_refill_rate
 
     def __add__(self, other) -> "TokenBucket":
         """Combine two token buckets.
@@ -57,7 +73,17 @@ class TokenBucket:
         self.log.append((time.monotonic(), self.tokens))
 
     def refill(self) -> None:
-        """Refill the bucket with new tokens based on elapsed time."""
+        """Refill the bucket with new tokens based on elapsed time.
+
+
+
+        >>> bucket = TokenBucket(bucket_name="test", bucket_type="test", capacity=10, refill_rate=1)
+        >>> bucket.tokens = 0
+        >>> bucket.refill()
+        >>> bucket.tokens > 0
+        True
+
+        """
         now = time.monotonic()
         elapsed = now - self.last_refill
         refill_amount = elapsed * self.refill_rate
@@ -100,7 +126,7 @@ class TokenBucket:
             raise ValueError(msg)
         while self.tokens < amount:
             self.refill()
-            await asyncio.sleep(0.1)  # Sleep briefly to prevent busy waiting
+            await asyncio.sleep(0.01)  # Sleep briefly to prevent busy waiting
         self.tokens -= amount
 
         now = time.monotonic()
@@ -114,6 +140,7 @@ class TokenBucket:
         times, tokens = zip(*self.get_log())
         start_time = times[0]
         times = [t - start_time for t in times]  # Normalize time to start from 0
+        from matplotlib import pyplot as plt
 
         plt.figure(figsize=(10, 6))
         plt.plot(times, tokens, label="Tokens Available")

edsl/jobs/interviews/Interview.py

@@ -6,15 +6,9 @@ import asyncio
 import time
 from typing import Any, Type, List, Generator, Optional
 
-from edsl.agents import Agent
-from edsl.language_models import LanguageModel
-from edsl.scenarios import Scenario
-from edsl.surveys import Survey
-
 from edsl.jobs.Answers import Answers
 from edsl.surveys.base import EndOfSurvey
 from edsl.jobs.buckets.ModelBuckets import ModelBuckets
-
 from edsl.jobs.tasks.TaskCreators import TaskCreators
 
 from edsl.jobs.interviews.InterviewStatusLog import InterviewStatusLog
@@ -26,6 +20,12 @@ from edsl.jobs.interviews.retry_management import retry_strategy
 from edsl.jobs.interviews.InterviewTaskBuildingMixin import InterviewTaskBuildingMixin
 from edsl.jobs.interviews.InterviewStatusMixin import InterviewStatusMixin
 
+import asyncio
+
+
+def run_async(coro):
+    return asyncio.run(coro)
+
 
 class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
     """
@@ -36,14 +36,14 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
 
     def __init__(
         self,
-        agent: Agent,
-        survey: Survey,
-        scenario: Scenario,
-        model: Type[LanguageModel],
-        debug: bool = False,
+        agent: "Agent",
+        survey: "Survey",
+        scenario: "Scenario",
+        model: Type["LanguageModel"],
+        debug: Optional[bool] = False,
         iteration: int = 0,
-        cache: "Cache" = None,
-        sidecar_model: LanguageModel = None,
+        cache: Optional["Cache"] = None,
+        sidecar_model: Optional["LanguageModel"] = None,
     ):
         """Initialize the Interview instance.
 
@@ -51,6 +51,24 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         :param survey: the survey being administered to the agent.
         :param scenario: the scenario that populates the survey questions.
         :param model: the language model used to answer the questions.
+        :param debug: if True, run without calls to the language model.
+        :param iteration: the iteration number of the interview.
+        :param cache: the cache used to store the answers.
+        :param sidecar_model: a sidecar model used to answer questions.
+
+        >>> i = Interview.example()
+        >>> i.task_creators
+        {}
+
+        >>> i.exceptions
+        {}
+
+        >>> _ = asyncio.run(i.async_conduct_interview())
+        >>> i.task_status_logs['q0']
+        [{'log_time': ..., 'value': <TaskStatus.NOT_STARTED: 1>}, {'log_time': ..., 'value': <TaskStatus.WAITING_FOR_DEPENDENCIES: 2>}, {'log_time': ..., 'value': <TaskStatus.API_CALL_IN_PROGRESS: 7>}, {'log_time': ..., 'value': <TaskStatus.SUCCESS: 8>}]
+
+        >>> i.to_index
+        {'q0': 0, 'q1': 1, 'q2': 2}
 
         """
         self.agent = agent
@@ -70,7 +88,7 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         self.exceptions = InterviewExceptionCollection()
         self._task_status_log_dict = InterviewStatusLog()
 
-        # dictionary mapping question names to their index in the survey."""
+        # dictionary mapping question names to their index in the survey.
         self.to_index = {
             question_name: index
             for index, question_name in enumerate(self.survey.question_names)
@@ -82,14 +100,16 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         model_buckets: ModelBuckets = None,
         debug: bool = False,
         stop_on_exception: bool = False,
-        sidecar_model: Optional[LanguageModel] = None,
+        sidecar_model: Optional["LanguageModel"] = None,
     ) -> tuple["Answers", List[dict[str, Any]]]:
         """
         Conduct an Interview asynchronously.
+        It returns a tuple with the answers and a list of valid results.
 
         :param model_buckets: a dictionary of token buckets for the model.
         :param debug: run without calls to LLM.
         :param stop_on_exception: if True, stops the interview if an exception is raised.
+        :param sidecar_model: a sidecar model used to answer questions.
 
         Example usage:
 
@@ -97,16 +117,37 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         >>> result, _ = asyncio.run(i.async_conduct_interview())
         >>> result['q0']
         'yes'
+
+        >>> i = Interview.example(throw_exception = True)
+        >>> result, _ = asyncio.run(i.async_conduct_interview())
+        Attempt 1 failed with exception:This is a test error now waiting 1.00 seconds before retrying.Parameters: start=1.0, max=60.0, max_attempts=5.
+        <BLANKLINE>
+        <BLANKLINE>
+        Attempt 2 failed with exception:This is a test error now waiting 2.00 seconds before retrying.Parameters: start=1.0, max=60.0, max_attempts=5.
+        <BLANKLINE>
+        <BLANKLINE>
+        Attempt 3 failed with exception:This is a test error now waiting 4.00 seconds before retrying.Parameters: start=1.0, max=60.0, max_attempts=5.
+        <BLANKLINE>
+        <BLANKLINE>
+        Attempt 4 failed with exception:This is a test error now waiting 8.00 seconds before retrying.Parameters: start=1.0, max=60.0, max_attempts=5.
+        <BLANKLINE>
+        <BLANKLINE>
+
+        >>> i.exceptions
+        {'q0': [{'exception': "Exception('This is a test error')", 'time': ..., 'traceback': ...
+
+        >>> i = Interview.example()
+        >>> result, _ = asyncio.run(i.async_conduct_interview(stop_on_exception = True))
+        Traceback (most recent call last):
+        ...
+        asyncio.exceptions.CancelledError
         """
         self.sidecar_model = sidecar_model
 
         # if no model bucket is passed, create an 'infinity' bucket with no rate limits
-        # print("model_buckets", model_buckets)
         if model_buckets is None or hasattr(self.agent, "answer_question_directly"):
             model_buckets = ModelBuckets.infinity_bucket()
 
-        # model_buckets = ModelBuckets.infinity_bucket()
-
         ## build the tasks using the InterviewTaskBuildingMixin
         ## This is the key part---it creates a task for each question,
         ## with dependencies on the questions that must be answered before this one can be answered.
@@ -128,6 +169,14 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         It iterates through the tasks and invigilators, and yields the results of the tasks that are done.
         If a task is not done, it raises a ValueError.
         If an exception is raised in the task, it records the exception in the Interview instance except if the task was cancelled, which is expected behavior.
+
+        >>> i = Interview.example()
+        >>> result, _ = asyncio.run(i.async_conduct_interview())
+        >>> results = list(i._extract_valid_results())
+        >>> len(results) == len(i.survey)
+        True
+        >>> type(results[0])
+        <class 'edsl.data_transfer_models.AgentResponseDict'>
         """
         assert len(self.tasks) == len(self.invigilators)
 
@@ -145,7 +194,18 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
             yield result
 
     def _record_exception(self, task, exception: Exception) -> None:
-        """Record an exception in the Interview instance."""
+        """Record an exception in the Interview instance.
+
+        It records the exception in the Interview instance, with the task name and the exception entry.
+
+        >>> i = Interview.example()
+        >>> result, _ = asyncio.run(i.async_conduct_interview())
+        >>> i.exceptions
+        {}
+        >>> i._record_exception(i.tasks[0], Exception("An exception occurred."))
+        >>> i.exceptions
+        {'q0': [{'exception': "Exception('An exception occurred.')", 'time': ..., 'traceback': 'NoneType: None\\n'}]}
+        """
         exception_entry = InterviewExceptionEntry(
             exception=repr(exception),
             time=time.time(),
@@ -161,6 +221,10 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         It is used to determine the order in which questions should be answered.
         This reflects both agent 'memory' considerations and 'skip' logic.
         The 'textify' parameter is set to True, so that the question names are returned as strings rather than integer indices.
+
+        >>> i = Interview.example()
+        >>> i.dag == {'q2': {'q0'}, 'q1': {'q0'}}
+        True
         """
         return self.survey.dag(textify=True)
 
@@ -171,8 +235,15 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         """Return a string representation of the Interview instance."""
         return f"Interview(agent = {repr(self.agent)}, survey = {repr(self.survey)}, scenario = {repr(self.scenario)}, model = {repr(self.model)})"
 
-    def duplicate(self, iteration: int, cache: Cache) -> Interview:
-        """Duplicate the interview, but with a new iteration number and cache."""
+    def duplicate(self, iteration: int, cache: "Cache") -> Interview:
+        """Duplicate the interview, but with a new iteration number and cache.
+
+        >>> i = Interview.example()
+        >>> i2 = i.duplicate(1, None)
+        >>> i.iteration + 1 == i2.iteration
+        True
+
+        """
         return Interview(
             agent=self.agent,
             survey=self.survey,
@@ -183,7 +254,7 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         )
 
     @classmethod
-    def example(self):
+    def example(self, throw_exception: bool = False) -> Interview:
         """Return an example Interview instance."""
         from edsl.agents import Agent
         from edsl.surveys import Survey
@@ -198,66 +269,15 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         survey = Survey.example()
         scenario = Scenario.example()
         model = LanguageModel.example()
+        if throw_exception:
+            model = LanguageModel.example(test_model=True, throw_exception=True)
+            agent = Agent.example()
+            return Interview(agent=agent, survey=survey, scenario=scenario, model=model)
         return Interview(agent=agent, survey=survey, scenario=scenario, model=model)
 
 
 if __name__ == "__main__":
     import doctest
 
-    doctest.testmod()
-    # from edsl import Model
-    # from edsl.agents import Agent
-    # from edsl.surveys import Survey
-    # from edsl.scenarios import Scenario
-    # from edsl.questions import QuestionMultipleChoice
-
-    # # from edsl.jobs.Interview import Interview
-
-    # #  a survey with skip logic
-    # q0 = QuestionMultipleChoice(
-    #     question_text="Do you like school?",
-    #     question_options=["yes", "no"],
-    #     question_name="q0",
-    # )
-    # q1 = QuestionMultipleChoice(
-    #     question_text="Why not?",
-    #     question_options=["killer bees in cafeteria", "other"],
-    #     question_name="q1",
-    # )
-    # q2 = QuestionMultipleChoice(
-    #     question_text="Why?",
-    #     question_options=["**lack*** of killer bees in cafeteria", "other"],
-    #     question_name="q2",
-    # )
-    # s = Survey(questions=[q0, q1, q2])
-    # s = s.add_rule(q0, "q0 == 'yes'", q2)
-
-    # # create an interview
-    # a = Agent(traits=None)
-
-    # def direct_question_answering_method(self, question, scenario):
-    #     """Answer a question directly."""
-    #     raise Exception("Error!")
-    #     # return "yes"
-
-    # a.add_direct_question_answering_method(direct_question_answering_method)
-    # scenario = Scenario()
-    # m = Model()
-    # I = Interview(agent=a, survey=s, scenario=scenario, model=m)
-
-    # result = asyncio.run(I.async_conduct_interview())
-    # # # conduct five interviews
-    # # for _ in range(5):
-    # #     I.conduct_interview(debug=True)
-
-    # # # replace missing answers
-    # # I
-    # # repr(I)
-    # # eval(repr(I))
-    # # print(I.task_status_logs.status_matrix(20))
-    # status_matrix = I.task_status_logs.status_matrix(20)
-    # numerical_matrix = I.task_status_logs.numerical_matrix(20)
-    # I.task_status_logs.visualize()
-
-    # I.exceptions.print()
-    # I.exceptions.ascii_table()
+    # add ellipsis
+    doctest.testmod(optionflags=doctest.ELLIPSIS)

edsl/jobs/interviews/InterviewTaskBuildingMixin.py

@@ -5,17 +5,19 @@ import asyncio
 import time
 import traceback
 from typing import Generator, Union
+
 from edsl import CONFIG
 from edsl.exceptions import InterviewTimeoutError
-from edsl.data_transfer_models import AgentResponseDict
-from edsl.questions.QuestionBase import QuestionBase
+
+# from edsl.questions.QuestionBase import QuestionBase
 from edsl.surveys.base import EndOfSurvey
 from edsl.jobs.buckets.ModelBuckets import ModelBuckets
 from edsl.jobs.interviews.interview_exception_tracking import InterviewExceptionEntry
 from edsl.jobs.interviews.retry_management import retry_strategy
 from edsl.jobs.tasks.task_status_enum import TaskStatus
 from edsl.jobs.tasks.QuestionTaskCreator import QuestionTaskCreator
-from edsl.agents.InvigilatorBase import InvigilatorBase
+
+# from edsl.agents.InvigilatorBase import InvigilatorBase
 
 TIMEOUT = float(CONFIG.get("EDSL_API_TIMEOUT"))
 
@@ -23,7 +25,7 @@ TIMEOUT = float(CONFIG.get("EDSL_API_TIMEOUT"))
 class InterviewTaskBuildingMixin:
     def _build_invigilators(
         self, debug: bool
-    ) -> Generator[InvigilatorBase, None, None]:
+    ) -> Generator["InvigilatorBase", None, None]:
         """Create an invigilator for each question.
 
         :param debug: whether to use debug mode, in which case `InvigilatorDebug` is used.
@@ -33,7 +35,7 @@ class InterviewTaskBuildingMixin:
         for question in self.survey.questions:
             yield self._get_invigilator(question=question, debug=debug)
 
-    def _get_invigilator(self, question: QuestionBase, debug: bool) -> "Invigilator":
+    def _get_invigilator(self, question: "QuestionBase", debug: bool) -> "Invigilator":
         """Return an invigilator for the given question.
 
         :param question: the question to be answered
@@ -44,6 +46,7 @@ class InterviewTaskBuildingMixin:
             scenario=self.scenario,
             model=self.model,
             debug=debug,
+            survey=self.survey,
             memory_plan=self.survey.memory_plan,
             current_answers=self.answers,
             iteration=self.iteration,
@@ -81,7 +84,7 @@ class InterviewTaskBuildingMixin:
         return tuple(tasks)  # , invigilators
 
     def _get_tasks_that_must_be_completed_before(
-        self, *, tasks: list[asyncio.Task], question: QuestionBase
+        self, *, tasks: list[asyncio.Task], question: "QuestionBase"
     ) -> Generator[asyncio.Task, None, None]:
         """Return the tasks that must be completed before the given question can be answered.
 
@@ -97,7 +100,7 @@ class InterviewTaskBuildingMixin:
     def _create_question_task(
         self,
         *,
-        question: QuestionBase,
+        question: "QuestionBase",
         tasks_that_must_be_completed_before: list[asyncio.Task],
         model_buckets: ModelBuckets,
         debug: bool,
@@ -149,15 +152,17 @@ class InterviewTaskBuildingMixin:
     async def _answer_question_and_record_task(
         self,
         *,
-        question: QuestionBase,
+        question: "QuestionBase",
         debug: bool,
         task=None,
-    ) -> AgentResponseDict:
+    ) -> "AgentResponseDict":
         """Answer a question and records the task.
 
         This in turn calls the the passed-in agent's async_answer_question method, which returns a response dictionary.
         Note that is updates answers dictionary with the response.
         """
+        from edsl.data_transfer_models import AgentResponseDict
+
         try:
             invigilator = self._get_invigilator(question, debug=debug)
 
@@ -170,24 +175,14 @@ class InterviewTaskBuildingMixin:
 
             self._add_answer(response=response, question=question)
 
-            # With the answer to the question, we can now cancel any skipped questions
             self._cancel_skipped_questions(question)
             return AgentResponseDict(**response)
         except Exception as e:
             raise e
-            # import traceback
-            # print("Exception caught:")
-            # traceback.print_exc()
-
-            # # Extract and print the traceback info
-            # tb = e.__traceback__
-            # while tb is not None:
-            #     print(f"File {tb.tb_frame.f_code.co_filename}, line {tb.tb_lineno}, in {tb.tb_frame.f_code.co_name}")
-            #     tb = tb.tb_next
-            #     breakpoint()
-            #     raise e
-
-    def _add_answer(self, response: AgentResponseDict, question: QuestionBase) -> None:
+
+    def _add_answer(
+        self, response: "AgentResponseDict", question: "QuestionBase"
+    ) -> None:
         """Add the answer to the answers dictionary.
 
         :param response: the response to the question.
@@ -195,7 +190,7 @@ class InterviewTaskBuildingMixin:
         """
         self.answers.add_answer(response=response, question=question)
 
-    def _skip_this_question(self, current_question: QuestionBase) -> bool:
+    def _skip_this_question(self, current_question: "QuestionBase") -> bool:
         """Determine if the current question should be skipped.
 
         :param current_question: the question to be answered.

edsl/jobs/runners/JobsRunnerAsyncio.py

@@ -1,29 +1,17 @@
 from __future__ import annotations
 import time
 import asyncio
-import textwrap
+import time
 from contextlib import contextmanager
 
 from typing import Coroutine, List, AsyncGenerator, Optional, Union
 
-from rich.live import Live
-from rich.console import Console
-
 from edsl import shared_globals
-from edsl.results import Results, Result
-
 from edsl.jobs.interviews.Interview import Interview
-from edsl.utilities.decorators import jupyter_nb_handler
-
-# from edsl.jobs.Jobs import Jobs
 from edsl.jobs.runners.JobsRunnerStatusMixin import JobsRunnerStatusMixin
-from edsl.language_models import LanguageModel
-from edsl.data.Cache import Cache
-
 from edsl.jobs.tasks.TaskHistory import TaskHistory
 from edsl.jobs.buckets.BucketCollection import BucketCollection
-
-import time
+from edsl.utilities.decorators import jupyter_nb_handler
 
 
 class JobsRunnerAsyncio(JobsRunnerStatusMixin):
@@ -42,13 +30,13 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
 
     async def run_async_generator(
         self,
-        cache: Cache,
+        cache: "Cache",
         n: int = 1,
         debug: bool = False,
         stop_on_exception: bool = False,
         sidecar_model: "LanguageModel" = None,
         total_interviews: Optional[List["Interview"]] = None,
-    ) -> AsyncGenerator[Result, None]:
+    ) -> AsyncGenerator["Result", None]:
         """Creates the tasks, runs them asynchronously, and returns the results as a Results object.
 
         Completed tasks are yielded as they are completed.
@@ -100,6 +88,8 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
                     self.total_interviews.append(interview)
 
     async def run_async(self, cache=None) -> Results:
+        from edsl.results.Results import Results
+
         if cache is None:
             self.cache = Cache()
         else:
@@ -110,6 +100,8 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
         return Results(survey=self.jobs.survey, data=data)
 
     def simple_run(self):
+        from edsl.results.Results import Results
+
         data = asyncio.run(self.run_async())
         return Results(survey=self.jobs.survey, data=data)
 
@@ -169,6 +161,8 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
                 question_name + "_raw_model_response"
             ] = result["raw_model_response"]
 
+        from edsl.results.Result import Result
+
         result = Result(
             agent=interview.agent,
             scenario=interview.scenario,
@@ -197,6 +191,8 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
         print_exceptions: bool = True,
     ) -> "Coroutine":
         """Runs a collection of interviews, handling both async and sync contexts."""
+        from rich.console import Console
+
         console = Console()
         self.results = []
         self.start_time = time.monotonic()
@@ -204,6 +200,8 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
         self.cache = cache
         self.sidecar_model = sidecar_model
 
+        from edsl.results.Results import Results
+
         if not progress_bar:
             # print("Running without progress bar")
             with cache as c:
@@ -225,6 +223,8 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
             results = Results(survey=self.jobs.survey, data=self.results)
         else:
             # print("Running with progress bar")
+            from rich.live import Live
+            from rich.console import Console
 
             def generate_table():
                 return self.status_table(self.results, self.elapsed_time)

edsl/jobs/tasks/QuestionTaskCreator.py

@@ -144,12 +144,16 @@ class QuestionTaskCreator(UserList):
             self.task_status = TaskStatus.FAILED
             raise e
 
-        if "cached_response" in results:
-            if results["cached_response"]:
-                # Gives back the tokens b/c the API was not called.
-                self.tokens_bucket.add_tokens(requested_tokens)
-                self.requests_bucket.add_tokens(1)
-                self.from_cache = True
+        if results.get("cache_used", False):
+            self.tokens_bucket.add_tokens(requested_tokens)
+            self.requests_bucket.add_tokens(1)
+            self.from_cache = True
+            # Turbo mode means that we don't wait for tokens or requests.
+            self.tokens_bucket.turbo_mode_on()
+            self.requests_bucket.turbo_mode_on()
+        else:
+            self.tokens_bucket.turbo_mode_off()
+            self.requests_bucket.turbo_mode_off()
 
         _ = results.pop("cached_response", None)
 

edsl/jobs/tasks/TaskHistory.py

@@ -1,8 +1,5 @@
 from edsl.jobs.tasks.task_status_enum import TaskStatus
-from matplotlib import pyplot as plt
 from typing import List, Optional
-
-import matplotlib.pyplot as plt
 from io import BytesIO
 import base64
 
@@ -75,6 +72,8 @@ class TaskHistory:
 
     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]
@@ -126,6 +125,8 @@ class TaskHistory:
         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