edsl 0.1.39__py3-none-any.whl → 0.1.39.dev2__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/TaskHistory.py

@@ -8,10 +8,9 @@ from edsl.Base import RepresentationMixin
 class TaskHistory(RepresentationMixin):
     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 +20,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):
         """ """

edsl/language_models/LanguageModel.py

@@ -44,8 +44,6 @@ if TYPE_CHECKING:
     from edsl.questions.QuestionBase import QuestionBase
     from edsl.language_models.key_management.KeyLookup import KeyLookup
 
-from edsl.enums import InferenceServiceType
-
 from edsl.utilities.decorators import (
     sync_wrapper,
     jupyter_nb_handler,
@@ -157,9 +155,7 @@ class LanguageModel(
             return klc.get(("config", "env"))
 
     def set_key_lookup(self, key_lookup: "KeyLookup") -> None:
-        """Set the key lookup, later"""
-        if hasattr(self, "_api_token"):
-            del self._api_token
+        del self._api_token
         self.key_lookup = key_lookup
 
     def ask_question(self, question: "QuestionBase") -> str:
@@ -497,10 +493,7 @@ class LanguageModel(
         >>> m.to_dict()
         {'model': '...', 'parameters': {'temperature': ..., 'max_tokens': ..., 'top_p': ..., 'frequency_penalty': ..., 'presence_penalty': ..., 'logprobs': False, 'top_logprobs': ...}, 'edsl_version': '...', 'edsl_class_name': 'LanguageModel'}
         """
-        d = {
-            "model": self.model,
-            "parameters": self.parameters,
-        }
+        d = {"model": self.model, "parameters": self.parameters}
         if add_edsl_version:
             from edsl import __version__
 
@@ -512,7 +505,7 @@ class LanguageModel(
     @remove_edsl_version
     def from_dict(cls, data: dict) -> Type[LanguageModel]:
         """Convert dictionary to a LanguageModel child instance."""
-        from edsl.language_models.model import get_model_class
+        from edsl.language_models.registry import get_model_class
 
         model_class = get_model_class(data["model"])
         return model_class(**data)
@@ -560,7 +553,7 @@ class LanguageModel(
         Exception report saved to ...
         Also see: ...
         """
-        from edsl.language_models.model import Model
+        from edsl.language_models.registry import Model
 
         if test_model:
             m = Model(
@@ -570,54 +563,6 @@ class LanguageModel(
         else:
             return Model(skip_api_key_check=True)
 
-    def from_cache(self, cache: "Cache") -> LanguageModel:
-
-        from copy import deepcopy
-        from types import MethodType
-        from edsl import Cache
-
-        new_instance = deepcopy(self)
-        print("Cache entries", len(cache))
-        new_instance.cache = Cache(
-            data={k: v for k, v in cache.items() if v.model == self.model}
-        )
-        print("Cache entries with same model", len(new_instance.cache))
-
-        new_instance.user_prompts = [
-            ce.user_prompt for ce in new_instance.cache.values()
-        ]
-        new_instance.system_prompts = [
-            ce.system_prompt for ce in new_instance.cache.values()
-        ]
-
-        async def async_execute_model_call(self, user_prompt: str, system_prompt: str):
-            cache_call_params = {
-                "model": str(self.model),
-                "parameters": self.parameters,
-                "system_prompt": system_prompt,
-                "user_prompt": user_prompt,
-                "iteration": 1,
-            }
-            cached_response, cache_key = cache.fetch(**cache_call_params)
-            response = json.loads(cached_response)
-            cost = 0
-            return ModelResponse(
-                response=response,
-                cache_used=True,
-                cache_key=cache_key,
-                cached_response=cached_response,
-                cost=cost,
-            )
-
-        # Bind the new method to the copied instance
-        setattr(
-            new_instance,
-            "async_execute_model_call",
-            MethodType(async_execute_model_call, new_instance),
-        )
-
-        return new_instance
-
 
 if __name__ == "__main__":
     """Run the module's test suite."""

edsl/language_models/ModelList.py

@@ -1,22 +1,18 @@
-from typing import Optional, List, TYPE_CHECKING
+from typing import Optional, List
 from collections import UserList
 
 from edsl.Base import Base
-from edsl.language_models.model import Model
+from edsl.language_models.registry import Model
 
-#
+# from edsl.language_models import LanguageModel
 from edsl.utilities.remove_edsl_version import remove_edsl_version
 from edsl.utilities.is_valid_variable_name import is_valid_variable_name
 
-if TYPE_CHECKING:
-    from edsl.inference_services.data_structures import AvailableModels
-    from edsl.language_models import LanguageModel
-
 
 class ModelList(Base, UserList):
     __documentation__ = """https://docs.expectedparrot.com/en/latest/language_models.html#module-edsl.language_models.ModelList"""
 
-    def __init__(self, data: Optional["LanguageModel"] = None):
+    def __init__(self, data: Optional[list] = None):
         """Initialize the ScenarioList class.
 
         >>> from edsl import Model
@@ -37,6 +33,9 @@ class ModelList(Base, UserList):
         """
         return set([model.model for model in self])
 
+    def rich_print(self):
+        pass
+
     def __repr__(self):
         return f"ModelList({super().__repr__()})"
 
@@ -88,7 +87,7 @@ class ModelList(Base, UserList):
             .table(*fields, tablefmt=tablefmt, pretty_labels=pretty_labels)
         )
 
-    def to_list(self) -> list:
+    def to_list(self):
         return self.to_scenario_list().to_list()
 
     def to_dict(self, sort=False, add_edsl_version=True):
@@ -121,16 +120,6 @@ class ModelList(Base, UserList):
             args = args[0]
         return ModelList([Model(model_name, **kwargs) for model_name in args])
 
-    @classmethod
-    def from_available_models(self, available_models_list: "AvailableModels"):
-        """Create a ModelList from an AvailableModels object"""
-        return ModelList(
-            [
-                Model(model.model_name, service_name=model.service_name)
-                for model in available_models_list
-            ]
-        )
-
     @classmethod
     @remove_edsl_version
     def from_dict(cls, data):

edsl/language_models/__init__.py

@@ -1,2 +1,2 @@
 from edsl.language_models.LanguageModel import LanguageModel
-from edsl.language_models.model import Model
+from edsl.language_models.registry import Model