edsl 0.1.38.dev1__py3-none-any.whl → 0.1.38.dev2__py3-none-any.whl

edsl/jobs/runners/JobsRunnerStatus.py

@@ -1,33 +1,21 @@
 from __future__ import annotations
 
+import os
 import time
-from dataclasses import dataclass, asdict
-
-from typing import List, DefaultDict, Optional, Type, Literal
-from collections import UserDict, defaultdict
-
-from rich.text import Text
-from rich.box import SIMPLE
-from rich.table import Table
-from rich.live import Live
-from rich.panel import Panel
-from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn
-from rich.layout import Layout
-from rich.console import Group
-from rich import box
-
-from edsl.jobs.interviews.InterviewStatusDictionary import InterviewStatusDictionary
+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 uuid import UUID
+
 from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
-from edsl.jobs.tokens.TokenUsage import TokenUsage
-from edsl.enums import get_token_pricing
-from edsl.jobs.tasks.task_status_enum import TaskStatus
 
 InterviewTokenUsageMapping = DefaultDict[str, InterviewTokenUsage]
 
 from edsl.jobs.interviews.InterviewStatistic import InterviewStatistic
-from edsl.jobs.interviews.InterviewStatisticsCollection import (
-    InterviewStatisticsCollection,
-)
 from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
 
 
@@ -47,16 +35,23 @@ class ModelTokenUsageStats:
     cost: str
 
 
-class Stats:
-    def elapsed_time(self):
-        InterviewStatistic("elapsed_time", value=elapsed_time, digits=1, units="sec.")
-
-
-class JobsRunnerStatus:
+class JobsRunnerStatusBase(ABC):
     def __init__(
-        self, jobs_runner: "JobsRunnerAsyncio", n: int, refresh_rate: float = 0.25
+        self,
+        jobs_runner: "JobsRunnerAsyncio",
+        n: int,
+        refresh_rate: float = 1,
+        endpoint_url: Optional[str] = "http://localhost:8000",
+        job_uuid: Optional[UUID] = None,
+        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
@@ -80,6 +75,99 @@ class JobsRunnerStatus:
 
         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.
+        """
+        pass
+
+    def get_status_dict(self) -> Dict[str, Any]:
+        """
+        Converts current status into a JSON-serializable dictionary.
+        """
+        # Get all statistics
+        stats = {}
+        for stat_name in self.statistics:
+            stat = self._compute_statistic(stat_name)
+            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 = {}
+        for model in self.distinct_models:
+            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_for_model,
+                "total": target_for_model,
+                "percent": (
+                    (completed_for_model / target_for_model * 100)
+                    if target_for_model > 0
+                    else 0
+                ),
+            }
+
+        status_dict = {
+            "overall_progress": {
+                "completed": completed,
+                "total": total_interviews,
+                "percent": (
+                    (completed / total_interviews * 100) if total_interviews > 0 else 0
+                ),
+            },
+            "language_model_progress": model_progress,
+            "statistics": stats,
+            "status": "completed" if completed >= total_interviews else "running",
+        }
+
+        model_queues = {}
+        for model, bucket in self.jobs_runner.bucket_collection.items():
+            model_name = model.model
+            model_queues[model_name] = {
+                "language_model_name": model_name,
+                "requests_bucket": {
+                    "completed": bucket.requests_bucket.num_released,
+                    "requested": bucket.requests_bucket.num_requests,
+                    "tokens_returned": bucket.requests_bucket.tokens_returned,
+                    "target_rate": round(bucket.requests_bucket.target_rate, 1),
+                    "current_rate": round(bucket.requests_bucket.get_throughput(), 1),
+                },
+                "tokens_bucket": {
+                    "completed": bucket.tokens_bucket.num_released,
+                    "requested": bucket.tokens_bucket.num_requests,
+                    "tokens_returned": bucket.tokens_bucket.tokens_returned,
+                    "target_rate": round(bucket.tokens_bucket.target_rate, 1),
+                    "current_rate": round(bucket.tokens_bucket.get_throughput(), 1),
+                },
+            }
+        status_dict["language_model_queues"] = model_queues
+        return status_dict
+
+    @abstractmethod
+    def setup(self):
+        """
+        Conducts any setup that needs to happen prior to sending status updates.
+
+        Ex. For a local job, creates a job in the Coop database.
+        """
+        pass
+
+    @abstractmethod
+    def send_status_update(self):
+        """
+        Updates the current status of the job.
+        """
+        pass
+
     def add_completed_interview(self, result):
         self.completed_interviews.append(result.interview_hash)
 
@@ -150,180 +238,92 @@ class JobsRunnerStatus:
         }
         return stat_definitions[stat_name]()
 
-    def create_progress_bar(self):
-        return Progress(
-            TextColumn("[progress.description]{task.description}"),
-            BarColumn(),
-            TaskProgressColumn(),
-            TextColumn("{task.completed}/{task.total}"),
-        )
+    def update_progress(self, stop_event):
 
-    def generate_model_queues_table(self):
-        table = Table(show_header=False, box=box.SIMPLE)
-        table.add_column("Info", style="cyan")
-        table.add_column("Value", style="magenta")
-        # table.add_row("Bucket collection", str(self.jobs_runner.bucket_collection))
-        for model, bucket in self.jobs_runner.bucket_collection.items():
-            table.add_row(Text(model.model, style="bold blue"), "")
-            bucket_types = ["requests_bucket", "tokens_bucket"]
-            for bucket_type in bucket_types:
-                table.add_row(Text(" " + bucket_type, style="green"), "")
-                # table.add_row(
-                #     f"  Current level (capacity = {round(getattr(bucket, bucket_type).capacity, 3)})",
-                #     str(round(getattr(bucket, bucket_type).tokens, 3)),
-                # )
-                num_requests = getattr(bucket, bucket_type).num_requests
-                num_released = getattr(bucket, bucket_type).num_released
-                tokens_returned = getattr(bucket, bucket_type).tokens_returned
-                # table.add_row(
-                #     f"  Requested",
-                #     str(num_requests),
-                # )
-                # table.add_row(
-                #     f"  Completed",
-                #     str(num_released),
-                # )
-                table.add_row(
-                    "  Completed vs. Requested", f"{num_released} vs. {num_requests}"
-                )
-                table.add_row(
-                    "  Added tokens (from cache)",
-                    str(tokens_returned),
-                )
-                if bucket_type == "tokens_bucket":
-                    rate_name = "TPM"
-                else:
-                    rate_name = "RPM"
-                target_rate = round(getattr(bucket, bucket_type).target_rate, 1)
-                table.add_row(
-                    f"  Empirical {rate_name} (target = {target_rate})",
-                    str(round(getattr(bucket, bucket_type).get_throughput(), 0)),
-                )
-
-        return table
-
-    def generate_layout(self):
-        progress = self.create_progress_bar()
-        task_ids = []
-        for model in self.distinct_models:
-            task_id = progress.add_task(
-                f"[cyan]{model}...",
-                total=int(self.num_total_interviews / len(self.distinct_models)),
-            )
-            task_ids.append((model, task_id))
-
-        progress_height = min(5, 2 + len(self.distinct_models))
-        layout = Layout()
-
-        # Create the top row with only the progress panel
-        layout.split_column(
-            Layout(
-                Panel(
-                    progress,
-                    title="Interview Progress",
-                    border_style="cyan",
-                    box=box.ROUNDED,
-                ),
-                name="progress",
-                size=progress_height,  # Adjusted size
-            ),
-            Layout(name="bottom_row"),  # Adjusted size
-        )
+        while not stop_event.is_set():
+            self.send_status_update()
+            time.sleep(self.refresh_rate)
 
-        # Split the bottom row into two columns for metrics and model queues
-        layout["bottom_row"].split_row(
-            Layout(
-                Panel(
-                    self.generate_metrics_table(),
-                    title="Metrics",
-                    border_style="magenta",
-                    box=box.ROUNDED,
-                ),
-                name="metrics",
-            ),
-            Layout(
-                Panel(
-                    self.generate_model_queues_table(),
-                    title="Model Queues",
-                    border_style="yellow",
-                    box=box.ROUNDED,
-                ),
-                name="model_queues",
-            ),
-        )
+        self.send_status_update()
 
-        return layout, progress, task_ids
 
-    def generate_metrics_table(self):
-        table = Table(show_header=True, header_style="bold magenta", box=box.SIMPLE)
-        table.add_column("Metric", style="cyan", no_wrap=True)
-        table.add_column("Value", justify="right")
+class JobsRunnerStatus(JobsRunnerStatusBase):
 
-        for stat_name in self.statistics:
-            pretty_name, value = list(self._compute_statistic(stat_name).items())[0]
-            # breakpoint()
-            table.add_row(pretty_name, value)
-        return table
+    @property
+    def create_url(self) -> str:
+        return f"{self.base_url}/api/v0/local-job"
 
-    def update_progress(self, stop_event):
-        layout, progress, task_ids = self.generate_layout()
-
-        with Live(
-            layout, refresh_per_second=int(1 / self.refresh_rate), transient=True
-        ) as live:
-            while (
-                len(self.completed_interviews) < len(self.jobs_runner.total_interviews)
-                and not stop_event.is_set()
-            ):
-                completed_tasks = len(self.completed_interviews)
-                total_tasks = len(self.jobs_runner.total_interviews)
-
-                for model, task_id in task_ids:
-                    completed_tasks = len(self.completed_interview_by_model[model])
-                    progress.update(
-                        task_id,
-                        completed=completed_tasks,
-                        description=f"[cyan]Conducting interviews for {model}...",
-                    )
-
-                layout["metrics"].update(
-                    Panel(
-                        self.generate_metrics_table(),
-                        title="Metrics",
-                        border_style="magenta",
-                        box=box.ROUNDED,
-                    )
-                )
-                layout["model_queues"].update(
-                    Panel(
-                        self.generate_model_queues_table(),
-                        title="Final Model Queues",
-                        border_style="yellow",
-                        box=box.ROUNDED,
-                    )
-                )
-
-                time.sleep(self.refresh_rate)
-
-            # Final update
-            for model, task_id in task_ids:
-                completed_tasks = len(self.completed_interview_by_model[model])
-                progress.update(
-                    task_id,
-                    completed=completed_tasks,
-                    description=f"[cyan]Conducting interviews for {model}...",
-                )
-
-            layout["metrics"].update(
-                Panel(
-                    self.generate_metrics_table(),
-                    title="Final Metrics",
-                    border_style="magenta",
-                    box=box.ROUNDED,
-                )
+    @property
+    def viewing_url(self) -> str:
+        return f"{self.base_url}/home/local-job-progress/{str(self.job_uuid)}"
+
+    @property
+    def update_url(self) -> str:
+        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"}
+
+        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")
+
+        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.
+        """
+        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"}
+
+            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,
+                headers=headers,
+                timeout=1,
             )
-            live.update(layout)
-            time.sleep(1)  # Show final state for 1 second
+            response.raise_for_status()
+        except requests.exceptions.RequestException as e:
+            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. Otherwise, returns False.
+        """
+
+        if self.api_key is not None:
+            return True
+        else:
+            return False
 
 
 if __name__ == "__main__":

edsl/jobs/tasks/TaskHistory.py

@@ -73,7 +73,7 @@ class TaskHistory:
         """Return a string representation of the TaskHistory."""
         return f"TaskHistory(interviews={self.total_interviews})."
 
-    def to_dict(self):
+    def to_dict(self, add_edsl_version=True):
         """Return the TaskHistory as a dictionary."""
         # return {
         #     "exceptions": [
@@ -82,10 +82,19 @@ class TaskHistory:
         #     ],
         #     "indices": self.indices,
         # }
-        return {
-            "interviews": [i._to_dict() for i in self.total_interviews],
+        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):

edsl/language_models/LanguageModel.py

@@ -607,18 +607,20 @@ class LanguageModel(
     #######################
     # SERIALIZATION METHODS
     #######################
-    def _to_dict(self) -> dict[str, Any]:
-        return {"model": self.model, "parameters": self.parameters}
-
-    @add_edsl_version
-    def to_dict(self) -> dict[str, Any]:
-        """Convert instance to a dictionary.
+    def to_dict(self, add_edsl_version=True) -> dict[str, Any]:
+        """Convert instance to a dictionary
 
         >>> m = LanguageModel.example()
         >>> m.to_dict()
         {'model': '...', 'parameters': {'temperature': ..., 'max_tokens': ..., 'top_p': ..., 'frequency_penalty': ..., 'presence_penalty': ..., 'logprobs': False, 'top_logprobs': ...}, 'edsl_version': '...', 'edsl_class_name': 'LanguageModel'}
         """
-        return self._to_dict()
+        d = {"model": self.model, "parameters": self.parameters}
+        if add_edsl_version:
+            from edsl import __version__
+
+            d["edsl_version"] = __version__
+            d["edsl_class_name"] = self.__class__.__name__
+        return d
 
     @classmethod
     @remove_edsl_version

edsl/language_models/ModelList.py

@@ -46,14 +46,30 @@ class ModelList(Base, UserList):
         """
         from edsl.utilities.utilities import dict_hash
 
-        return dict_hash(self._to_dict(sort=True))
+        return dict_hash(self.to_dict(sort=True, add_edsl_version=False))
 
-    def _to_dict(self, sort=False):
+    def to_dict(self, sort=False, add_edsl_version=True):
         if sort:
             model_list = sorted([model for model in self], key=lambda x: hash(x))
-            return {"models": [model._to_dict() for model in model_list]}
+            d = {
+                "models": [
+                    model.to_dict(add_edsl_version=add_edsl_version)
+                    for model in model_list
+                ]
+            }
         else:
-            return {"models": [model._to_dict() for model in self]}
+            d = {
+                "models": [
+                    model.to_dict(add_edsl_version=add_edsl_version) for model in self
+                ]
+            }
+        if add_edsl_version:
+            from edsl import __version__
+
+            d["edsl_version"] = __version__
+            d["edsl_class_name"] = "ModelList"
+
+        return d
 
     @classmethod
     def from_names(self, *args, **kwargs):
@@ -62,15 +78,6 @@ class ModelList(Base, UserList):
             args = args[0]
         return ModelList([Model(model_name, **kwargs) for model_name in args])
 
-    @add_edsl_version
-    def to_dict(self):
-        """
-        Convert the ModelList to a dictionary.
-        >>> ModelList.example().to_dict()
-        {'models': [...], 'edsl_version': '...', 'edsl_class_name': 'ModelList'}
-        """
-        return self._to_dict()
-
     @classmethod
     @remove_edsl_version
     def from_dict(cls, data):

edsl/notebooks/Notebook.py

@@ -102,18 +102,17 @@ class Notebook(Base):
 
         return dict_hash(self.data["cells"])
 
-    def _to_dict(self) -> dict:
+    def to_dict(self, add_edsl_version=False) -> dict:
         """
         Serialize to a dictionary.
         """
-        return {"name": self.name, "data": self.data}
+        d = {"name": self.name, "data": self.data}
+        if add_edsl_version:
+            from edsl import __version__
 
-    @add_edsl_version
-    def to_dict(self) -> dict:
-        """
-        Convert a Notebook to a dictionary.
-        """
-        return self._to_dict()
+            d["edsl_version"] = __version__
+            d["edsl_class_name"] = self.__class__.__name__
+        return d
 
     @classmethod
     @remove_edsl_version

edsl/questions/QuestionBase.py

@@ -82,8 +82,7 @@ class QuestionBase(
         if not hasattr(self, "_fake_data_factory"):
             from polyfactory.factories.pydantic_factory import ModelFactory
 
-            class FakeData(ModelFactory[self.response_model]):
-                ...
+            class FakeData(ModelFactory[self.response_model]): ...
 
             self._fake_data_factory = FakeData
         return self._fake_data_factory
@@ -135,7 +134,7 @@ class QuestionBase(
         """
         from edsl.utilities.utilities import dict_hash
 
-        return dict_hash(self._to_dict())
+        return dict_hash(self.to_dict(add_edsl_version=False))
 
     @property
     def data(self) -> dict:
@@ -147,13 +146,15 @@ class QuestionBase(
         """
         exclude_list = [
             "question_type",
-            "_include_comment",
+            # "_include_comment",
             "_fake_data_factory",
-            "_use_code",
+            # "_use_code",
             "_model_instructions",
         ]
         only_if_not_na_list = ["_answering_instructions", "_question_presentation"]
 
+        only_if_not_default_list = {"_include_comment": True, "_use_code": False}
+
         def ok(key, value):
             if not key.startswith("_"):
                 return False
@@ -161,6 +162,12 @@ class QuestionBase(
                 return False
             if key in only_if_not_na_list and value is None:
                 return False
+            if (
+                key in only_if_not_default_list
+                and value == only_if_not_default_list[key]
+            ):
+                return False
+
             return True
 
         candidate_data = {
@@ -175,25 +182,22 @@ class QuestionBase(
 
         return candidate_data
 
-    def _to_dict(self):
+    def to_dict(self, add_edsl_version=True):
         """Convert the question to a dictionary that includes the question type (used in deserialization).
 
-        >>> from edsl import QuestionFreeText as Q; Q.example()._to_dict()
+        >>> from edsl import QuestionFreeText as Q; Q.example().to_dict(add_edsl_version = False)
         {'question_name': 'how_are_you', 'question_text': 'How are you?', 'question_type': 'free_text'}
         """
         candidate_data = self.data.copy()
         candidate_data["question_type"] = self.question_type
-        return {
-            key: value for key, value in candidate_data.items() if value is not None
-        }
+        d = {key: value for key, value in candidate_data.items() if value is not None}
+        if add_edsl_version:
+            from edsl import __version__
 
-    @add_edsl_version
-    def to_dict(self) -> dict[str, Any]:
-        """Convert the question to a dictionary that includes the question type (used in deserialization).
-        >>> from edsl import QuestionFreeText as Q; Q.example().to_dict()
-        {'question_name': 'how_are_you', 'question_text': 'How are you?', 'question_type': 'free_text', 'edsl_version': '...'}
-        """
-        return self._to_dict()
+            d["edsl_version"] = __version__
+            d["edsl_class_name"] = "QuestionBase"
+
+        return d
 
     @classmethod
     @remove_edsl_version

edsl/questions/QuestionBaseGenMixin.py

@@ -138,7 +138,7 @@ class QuestionBaseGenMixin:
         if exclude_components is None:
             exclude_components = ["question_name", "question_type"]
 
-        d = copy.deepcopy(self._to_dict())
+        d = copy.deepcopy(self.to_dict(add_edsl_version=False))
         for key, value in d.items():
             if key in exclude_components:
                 continue