edsl 0.1.39__py3-none-any.whl → 0.1.39.dev1__py3-none-any.whl

edsl/jobs/runners/JobsRunnerStatus.py

@@ -3,12 +3,21 @@ from __future__ import annotations
 import os
 import time
 import requests
+import warnings
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
+
+from typing import Any, List, DefaultDict, Optional, Dict
 from collections import defaultdict
-from typing import Any, Dict, Optional
 from uuid import UUID
 
+from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
+
+InterviewTokenUsageMapping = DefaultDict[str, InterviewTokenUsage]
+
+from edsl.jobs.interviews.InterviewStatistic import InterviewStatistic
+from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
+
 
 @dataclass
 class ModelInfo:
@@ -19,44 +28,11 @@ class ModelInfo:
     token_usage_info: dict
 
 
-class StatisticsTracker:
-    def __init__(self, total_interviews: int, distinct_models: list[str]):
-        self.start_time = time.time()
-        self.total_interviews = total_interviews
-        self.completed_count = 0
-        self.completed_by_model = defaultdict(int)
-        self.distinct_models = distinct_models
-        self.total_exceptions = 0
-        self.unfixed_exceptions = 0
-
-    def add_completed_interview(
-        self, model: str, num_exceptions: int = 0, num_unfixed: int = 0
-    ):
-        self.completed_count += 1
-        self.completed_by_model[model] += 1
-        self.total_exceptions += num_exceptions
-        self.unfixed_exceptions += num_unfixed
-
-    def get_elapsed_time(self) -> float:
-        return time.time() - self.start_time
-
-    def get_average_time_per_interview(self) -> float:
-        return (
-            self.get_elapsed_time() / self.completed_count
-            if self.completed_count > 0
-            else 0
-        )
-
-    def get_throughput(self) -> float:
-        elapsed = self.get_elapsed_time()
-        return self.completed_count / elapsed if elapsed > 0 else 0
-
-    def get_estimated_time_remaining(self) -> float:
-        if self.completed_count == 0:
-            return 0
-        avg_time = self.get_average_time_per_interview()
-        remaining = self.total_interviews - self.completed_count
-        return avg_time * remaining
+@dataclass
+class ModelTokenUsageStats:
+    token_usage_type: str
+    details: List[dict]
+    cost: str
 
 
 class JobsRunnerStatusBase(ABC):
@@ -70,39 +46,48 @@ class JobsRunnerStatusBase(ABC):
         api_key: str = None,
     ):
         self.jobs_runner = jobs_runner
+
+        # The uuid of the job on Coop
         self.job_uuid = job_uuid
+
         self.base_url = f"{endpoint_url}"
+
+        self.start_time = time.time()
+        self.completed_interviews = []
         self.refresh_rate = refresh_rate
         self.statistics = [
             "elapsed_time",
             "total_interviews_requested",
             "completed_interviews",
+            #            "percent_complete",
             "average_time_per_interview",
+            #            "task_remaining",
             "estimated_time_remaining",
             "exceptions",
             "unfixed_exceptions",
             "throughput",
         ]
-        self.num_total_interviews = n * len(self.jobs_runner)
+        self.num_total_interviews = n * len(self.jobs_runner.interviews)
 
         self.distinct_models = list(
-            set(model.model for model in self.jobs_runner.jobs.models)
+            set(i.model.model for i in self.jobs_runner.interviews)
         )
 
-        self.stats_tracker = StatisticsTracker(
-            total_interviews=self.num_total_interviews,
-            distinct_models=self.distinct_models,
-        )
+        self.completed_interview_by_model = defaultdict(list)
 
         self.api_key = api_key or os.getenv("EXPECTED_PARROT_API_KEY")
 
     @abstractmethod
     def has_ep_api_key(self):
-        """Checks if the user has an Expected Parrot API key."""
+        """
+        Checks if the user has an Expected Parrot API key.
+        """
         pass
 
     def get_status_dict(self) -> Dict[str, Any]:
-        """Converts current status into a JSON-serializable dictionary."""
+        """
+        Converts current status into a JSON-serializable dictionary.
+        """
         # Get all statistics
         stats = {}
         for stat_name in self.statistics:
@@ -110,46 +95,42 @@ class JobsRunnerStatusBase(ABC):
             name, value = list(stat.items())[0]
             stats[name] = value
 
+        # Calculate overall progress
+        total_interviews = len(self.jobs_runner.total_interviews)
+        completed = len(self.completed_interviews)
+
         # Get model-specific progress
         model_progress = {}
-        target_per_model = int(self.num_total_interviews / len(self.distinct_models))
-
         for model in self.distinct_models:
-            completed = self.stats_tracker.completed_by_model[model]
+            completed_for_model = len(self.completed_interview_by_model[model])
+            target_for_model = int(
+                self.num_total_interviews / len(self.distinct_models)
+            )
             model_progress[model] = {
-                "completed": completed,
-                "total": target_per_model,
+                "completed": completed_for_model,
+                "total": target_for_model,
                 "percent": (
-                    (completed / target_per_model * 100) if target_per_model > 0 else 0
+                    (completed_for_model / target_for_model * 100)
+                    if target_for_model > 0
+                    else 0
                 ),
             }
 
         status_dict = {
             "overall_progress": {
-                "completed": self.stats_tracker.completed_count,
-                "total": self.num_total_interviews,
+                "completed": completed,
+                "total": total_interviews,
                 "percent": (
-                    (
-                        self.stats_tracker.completed_count
-                        / self.num_total_interviews
-                        * 100
-                    )
-                    if self.num_total_interviews > 0
-                    else 0
+                    (completed / total_interviews * 100) if total_interviews > 0 else 0
                 ),
             },
             "language_model_progress": model_progress,
             "statistics": stats,
-            "status": (
-                "completed"
-                if self.stats_tracker.completed_count >= self.num_total_interviews
-                else "running"
-            ),
+            "status": "completed" if completed >= total_interviews else "running",
         }
 
         model_queues = {}
-        # for model, bucket in self.jobs_runner.bucket_collection.items():
-        for model, bucket in self.jobs_runner.environment.bucket_collection.items():
+        for model, bucket in self.jobs_runner.bucket_collection.items():
             model_name = model.model
             model_queues[model_name] = {
                 "language_model_name": model_name,
@@ -171,67 +152,98 @@ class JobsRunnerStatusBase(ABC):
         status_dict["language_model_queues"] = model_queues
         return status_dict
 
-    def add_completed_interview(self, result):
-        """Records a completed interview without storing the full interview data."""
-        self.stats_tracker.add_completed_interview(
-            model=result.model.model,
-            num_exceptions=(
-                len(result.exceptions) if hasattr(result, "exceptions") else 0
-            ),
-            num_unfixed=(
-                result.exceptions.num_unfixed() if hasattr(result, "exceptions") else 0
-            ),
-        )
-
-    def _compute_statistic(self, stat_name: str):
-        """Computes individual statistics based on the stats tracker."""
-        if stat_name == "elapsed_time":
-            value = self.stats_tracker.get_elapsed_time()
-            return {"elapsed_time": (value, 1, "sec.")}
-
-        elif stat_name == "total_interviews_requested":
-            return {"total_interviews_requested": (self.num_total_interviews, None, "")}
+    @abstractmethod
+    def setup(self):
+        """
+        Conducts any setup that needs to happen prior to sending status updates.
 
-        elif stat_name == "completed_interviews":
-            return {
-                "completed_interviews": (self.stats_tracker.completed_count, None, "")
-            }
+        Ex. For a local job, creates a job in the Coop database.
+        """
+        pass
 
-        elif stat_name == "average_time_per_interview":
-            value = self.stats_tracker.get_average_time_per_interview()
-            return {"average_time_per_interview": (value, 2, "sec.")}
+    @abstractmethod
+    def send_status_update(self):
+        """
+        Updates the current status of the job.
+        """
+        pass
 
-        elif stat_name == "estimated_time_remaining":
-            value = self.stats_tracker.get_estimated_time_remaining()
-            return {"estimated_time_remaining": (value, 1, "sec.")}
+    def add_completed_interview(self, result):
+        self.completed_interviews.append(result.interview_hash)
 
-        elif stat_name == "exceptions":
-            return {"exceptions": (self.stats_tracker.total_exceptions, None, "")}
+        relevant_model = result.model.model
+        self.completed_interview_by_model[relevant_model].append(result.interview_hash)
 
-        elif stat_name == "unfixed_exceptions":
-            return {
-                "unfixed_exceptions": (self.stats_tracker.unfixed_exceptions, None, "")
-            }
+    def _compute_statistic(self, stat_name: str):
+        completed_tasks = self.completed_interviews
+        elapsed_time = time.time() - self.start_time
+        interviews = self.jobs_runner.total_interviews
 
-        elif stat_name == "throughput":
-            value = self.stats_tracker.get_throughput()
-            return {"throughput": (value, 2, "interviews/sec.")}
+        stat_definitions = {
+            "elapsed_time": lambda: InterviewStatistic(
+                "elapsed_time", value=elapsed_time, digits=1, units="sec."
+            ),
+            "total_interviews_requested": lambda: InterviewStatistic(
+                "total_interviews_requested", value=len(interviews), units=""
+            ),
+            "completed_interviews": lambda: InterviewStatistic(
+                "completed_interviews", value=len(completed_tasks), units=""
+            ),
+            "percent_complete": lambda: InterviewStatistic(
+                "percent_complete",
+                value=(
+                    len(completed_tasks) / len(interviews) * 100
+                    if len(interviews) > 0
+                    else 0
+                ),
+                digits=1,
+                units="%",
+            ),
+            "average_time_per_interview": lambda: InterviewStatistic(
+                "average_time_per_interview",
+                value=elapsed_time / len(completed_tasks) if completed_tasks else 0,
+                digits=2,
+                units="sec.",
+            ),
+            "task_remaining": lambda: InterviewStatistic(
+                "task_remaining", value=len(interviews) - len(completed_tasks), units=""
+            ),
+            "estimated_time_remaining": lambda: InterviewStatistic(
+                "estimated_time_remaining",
+                value=(
+                    (len(interviews) - len(completed_tasks))
+                    * (elapsed_time / len(completed_tasks))
+                    if len(completed_tasks) > 0
+                    else 0
+                ),
+                digits=1,
+                units="sec.",
+            ),
+            "exceptions": lambda: InterviewStatistic(
+                "exceptions",
+                value=sum(len(i.exceptions) for i in interviews),
+                units="",
+            ),
+            "unfixed_exceptions": lambda: InterviewStatistic(
+                "unfixed_exceptions",
+                value=sum(i.exceptions.num_unfixed() for i in interviews),
+                units="",
+            ),
+            "throughput": lambda: InterviewStatistic(
+                "throughput",
+                value=len(completed_tasks) / elapsed_time if elapsed_time > 0 else 0,
+                digits=2,
+                units="interviews/sec.",
+            ),
+        }
+        return stat_definitions[stat_name]()
 
     def update_progress(self, stop_event):
         while not stop_event.is_set():
             self.send_status_update()
             time.sleep(self.refresh_rate)
-        self.send_status_update()
 
-    @abstractmethod
-    def setup(self):
-        """Conducts any setup needed prior to sending status updates."""
-        pass
-
-    @abstractmethod
-    def send_status_update(self):
-        """Updates the current status of the job."""
-        pass
+        self.send_status_update()
 
 
 class JobsRunnerStatus(JobsRunnerStatusBase):
@@ -248,35 +260,49 @@ class JobsRunnerStatus(JobsRunnerStatusBase):
         return f"{self.base_url}/api/v0/local-job/{str(self.job_uuid)}"
 
     def setup(self) -> None:
-        """Creates a local job on Coop if one does not already exist."""
-        headers = {
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {self.api_key or 'None'}",
-        }
+        """
+        Creates a local job on Coop if one does not already exist.
+        """
+
+        headers = {"Content-Type": "application/json"}
+
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+        else:
+            headers["Authorization"] = f"Bearer None"
 
         if self.job_uuid is None:
+            # Create a new local job
             response = requests.post(
                 self.create_url,
                 headers=headers,
                 timeout=1,
             )
-            response.raise_for_status()
-            data = response.json()
-            self.job_uuid = data.get("job_uuid")
+        response.raise_for_status()
+        data = response.json()
+        self.job_uuid = data.get("job_uuid")
 
         print(f"Running with progress bar. View progress at {self.viewing_url}")
 
     def send_status_update(self) -> None:
-        """Sends current status to the web endpoint using the instance's job_uuid."""
+        """
+        Sends current status to the web endpoint using the instance's job_uuid.
+        """
         try:
+            # Get the status dictionary and add the job_id
             status_dict = self.get_status_dict()
+
+            # Make the UUID JSON serializable
             status_dict["job_id"] = str(self.job_uuid)
 
-            headers = {
-                "Content-Type": "application/json",
-                "Authorization": f"Bearer {self.api_key or 'None'}",
-            }
+            headers = {"Content-Type": "application/json"}
+
+            if self.api_key:
+                headers["Authorization"] = f"Bearer {self.api_key}"
+            else:
+                headers["Authorization"] = f"Bearer None"
 
+            # Send the update
             response = requests.patch(
                 self.update_url,
                 json=status_dict,
@@ -288,8 +314,14 @@ class JobsRunnerStatus(JobsRunnerStatusBase):
             print(f"Failed to send status update for job {self.job_uuid}: {e}")
 
     def has_ep_api_key(self) -> bool:
-        """Returns True if the user has an Expected Parrot API key."""
-        return self.api_key is not None
+        """
+        Returns True if the user has an Expected Parrot API key. Otherwise, returns False.
+        """
+
+        if self.api_key is not None:
+            return True
+        else:
+            return False
 
 
 if __name__ == "__main__":

edsl/jobs/tasks/QuestionTaskCreator.py

@@ -1,17 +1,17 @@
 import asyncio
-from typing import Callable, Union, List, TYPE_CHECKING
+from typing import Callable, Union, List
 from collections import UserList, UserDict
 
-from edsl.exceptions.jobs import InterviewErrorPriorTaskCanceled
+from edsl.jobs.buckets import ModelBuckets
+from edsl.exceptions import InterviewErrorPriorTaskCanceled
 
+from edsl.jobs.interviews.InterviewStatusDictionary import InterviewStatusDictionary
 from edsl.jobs.tasks.task_status_enum import TaskStatus, TaskStatusDescriptor
 from edsl.jobs.tasks.TaskStatusLog import TaskStatusLog
+from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
 from edsl.jobs.tokens.TokenUsage import TokenUsage
 from edsl.jobs.Answers import Answers
-
-if TYPE_CHECKING:
-    from edsl.questions.QuestionBase import QuestionBase
-    from edsl.jobs.buckets import ModelBuckets
+from edsl.questions.QuestionBase import QuestionBase
 
 
 class TokensUsed(UserDict):
@@ -24,6 +24,7 @@ class TokensUsed(UserDict):
 
 class QuestionTaskCreator(UserList):
     """Class to create and manage a single question and its dependencies.
+    The class is an instance of a UserList of tasks that must be completed before the focal task can be run.
 
     It is a UserList with all the tasks that must be completed before the focal task can be run.
     The focal task is the question that we are interested in answering.
@@ -34,9 +35,9 @@ class QuestionTaskCreator(UserList):
     def __init__(
         self,
         *,
-        question: "QuestionBase",
+        question: QuestionBase,
         answer_question_func: Callable,
-        model_buckets: "ModelBuckets",
+        model_buckets: ModelBuckets,
         token_estimator: Union[Callable, None] = None,
         iteration: int = 0,
     ):
@@ -50,15 +51,14 @@ class QuestionTaskCreator(UserList):
 
         """
         super().__init__([])
+        # answer_question_func is the 'interview.answer_question_and_record_task" method
         self.answer_question_func = answer_question_func
         self.question = question
         self.iteration = iteration
 
         self.model_buckets = model_buckets
-
         self.requests_bucket = self.model_buckets.requests_bucket
         self.tokens_bucket = self.model_buckets.tokens_bucket
-
         self.status_log = TaskStatusLog()
 
         def fake_token_estimator(question):
@@ -125,13 +125,11 @@ class QuestionTaskCreator(UserList):
 
         await self.tokens_bucket.get_tokens(requested_tokens)
 
-        if (estimated_wait_time := self.model_buckets.requests_bucket.wait_time(1)) > 0:
+        if (estimated_wait_time := self.requests_bucket.wait_time(1)) > 0:
             self.waiting = True  #  do we need this?
             self.task_status = TaskStatus.WAITING_FOR_REQUEST_CAPACITY
 
-        await self.model_buckets.requests_bucket.get_tokens(
-            1, cheat_bucket_capacity=True
-        )
+        await self.requests_bucket.get_tokens(1, cheat_bucket_capacity=True)
 
         self.task_status = TaskStatus.API_CALL_IN_PROGRESS
         try:
@@ -144,22 +142,22 @@ class QuestionTaskCreator(UserList):
             raise e
 
         if results.cache_used:
-            self.model_buckets.tokens_bucket.add_tokens(requested_tokens)
-            self.model_buckets.requests_bucket.add_tokens(1)
+            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.model_buckets.tokens_bucket.turbo_mode_on()
-            self.model_buckets.requests_bucket.turbo_mode_on()
+            self.tokens_bucket.turbo_mode_on()
+            self.requests_bucket.turbo_mode_on()
         else:
-            self.model_buckets.tokens_bucket.turbo_mode_off()
-            self.model_buckets.requests_bucket.turbo_mode_off()
+            self.tokens_bucket.turbo_mode_off()
+            self.requests_bucket.turbo_mode_off()
 
         return results
 
     @classmethod
     def example(cls):
         """Return an example instance of the class."""
-        from edsl.questions.QuestionFreeText import QuestionFreeText
+        from edsl import QuestionFreeText
         from edsl.jobs.buckets.ModelBuckets import ModelBuckets
 
         m = ModelBuckets.infinity_bucket()

edsl/jobs/tasks/TaskHistory.py

@@ -1,17 +1,18 @@
 from typing import List, Optional
 from io import BytesIO
+import webbrowser
+import os
 import base64
+from importlib import resources
 from edsl.jobs.tasks.task_status_enum import TaskStatus
-from edsl.Base import RepresentationMixin
 
 
-class TaskHistory(RepresentationMixin):
+class TaskHistory:
     def __init__(
         self,
-        interviews: List["Interview"] = None,
+        interviews: List["Interview"],
         include_traceback: bool = False,
         max_interviews: int = 10,
-        interviews_with_exceptions_only: bool = False,
     ):
         """
         The structure of a TaskHistory exception
@@ -21,33 +22,13 @@ class TaskHistory(RepresentationMixin):
         >>> _ = 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.total_interviews = interviews
         self.include_traceback = include_traceback
 
-        # self._interviews = {index: i for index, i in enumerate(self.total_interviews)}
+        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):
         """ """
@@ -140,6 +121,14 @@ class TaskHistory(RepresentationMixin):
         """Return True if there are any exceptions."""
         return len(self.unfixed_exceptions) > 0
 
+    def _repr_html_(self):
+        """Return an HTML representation of the TaskHistory."""
+        d = self.to_dict(add_edsl_version=False)
+        data = [[k, v] for k, v in d.items()]
+        from tabulate import tabulate
+
+        return tabulate(data, headers=["keys", "values"], tablefmt="html")
+
     def show_exceptions(self, tracebacks=False):
         """Print the exceptions."""
         for index in self.indices:
@@ -251,15 +240,11 @@ class TaskHistory(RepresentationMixin):
             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
@@ -296,7 +281,7 @@ class TaskHistory(RepresentationMixin):
         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_type = interview.survey.get_question(
                     question_name
                 ).question_type
                 if (question_name, question_type) not in exceptions_by_question_name:
@@ -345,11 +330,8 @@ class TaskHistory(RepresentationMixin):
         }
         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 = ""
+    def generate_html_report(self, css: Optional[str]):
+        performance_plot_html = self.plot(num_periods=100, get_embedded_html=True)
 
         if css is None:
             css = self.css()
@@ -427,8 +409,6 @@ class TaskHistory(RepresentationMixin):
             print(f"Exception report saved to {filename}")
 
         if open_in_browser:
-            import webbrowser
-
             webbrowser.open(f"file://{os.path.abspath(filename)}")
 
         if return_link:

edsl/jobs/tasks/task_status_enum.py

@@ -3,6 +3,8 @@ from collections import UserDict
 import enum
 import time
 
+# from edsl.jobs.tasks.TaskStatusLogEntry import TaskStatusLogEntry
+
 
 class TaskStatus(enum.Enum):
     "These are the possible states a task can be in."