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

edsl/jobs/JobsRunnerStatusMixin.py

@@ -1,115 +0,0 @@
-from typing import List
-import asyncio
-
-from rich.table import Table
-from rich.text import Text
-from rich.box import SIMPLE
-
-from edsl.jobs.token_tracking import TokenPricing
-
-pricing = {
-    "gpt-3.5-turbo": TokenPricing(
-        model_name="gpt-3.5-turbo",
-        prompt_token_price_per_k=0.0005,
-        completion_token_price_per_k=0.0015,
-    ),
-    "gpt-4-1106-preview": TokenPricing(
-        model_name="gpt-4",
-        prompt_token_price_per_k=0.01,
-        completion_token_price_per_k=0.03,
-    ),
-    "test": TokenPricing(
-        model_name="test",
-        prompt_token_price_per_k=0.0,
-        completion_token_price_per_k=0.0,
-    ),
-    "gemini_pro": TokenPricing(
-        model_name="gemini_pro",
-        prompt_token_price_per_k=0.0,
-        completion_token_price_per_k=0.0,
-    ),
-    "llama-2-13b-chat-hf": TokenPricing(
-        model_name="llama-2-13b-chat-hf",
-        prompt_token_price_per_k=0.0,
-        completion_token_price_per_k=0.0,
-    ),
-    "llama-2-70b-chat-hf": TokenPricing(
-        model_name="llama-2-70b-chat-hf",
-        prompt_token_price_per_k=0.0,
-        completion_token_price_per_k=0.0,
-    ),
-    "mixtral-8x7B-instruct-v0.1": TokenPricing(
-        model_name="mixtral-8x7B-instruct-v0.1",
-        prompt_token_price_per_k=0.0,
-        completion_token_price_per_k=0.0,
-    ),
-}
-
-
-class ModelStatus:
-    def __init__(self, model, TPM, RPM):
-        self.model = model
-        self.TPM = TPM
-        self.RPM = RPM
-
-
-from edsl.jobs.token_tracking import InterviewTokenUsage
-from edsl.jobs.task_management import InterviewStatusDictionary
-
-from collections import defaultdict
-
-
-class JobsRunnerStatusMixin:
-    def _generate_status_table(self, data: List[asyncio.Task], elapsed_time):
-        models_to_tokens = defaultdict(InterviewTokenUsage)
-        model_to_status = defaultdict(InterviewStatusDictionary)
-
-        for interview in self.interviews:
-            model = interview.model
-            models_to_tokens[model] += interview.token_usage
-            model_to_status[model] += interview.interview_status
-
-        pct_complete = len(data) / len(self.interviews) * 100
-        average_time = elapsed_time / len(data) if len(data) > 0 else 0
-
-        table = Table(
-            title="Job Status",
-            show_header=True,
-            header_style="bold magenta",
-            box=SIMPLE,
-        )
-        table.add_column("Key", style="dim", no_wrap=True)
-        table.add_column("Value")
-
-        # Add rows for each key-value pair
-        table.add_row(Text("Task status", style="bold red"), "")
-        table.add_row("Total interviews requested", str(len(self.interviews)))
-        table.add_row("Completed interviews", str(len(data)))
-        # table.add_row("Interviews from cache", str(num_from_cache))
-        table.add_row("Percent complete", f"{pct_complete:.2f}%")
-        table.add_row("", "")
-
-        # table.add_row(Text("Timing", style = "bold red"), "")
-        # table.add_row("Elapsed time (seconds)", f"{elapsed_time:.3f}")
-        # table.add_row("Average time/interview (seconds)", f"{average_time:.3f}")
-        # table.add_row("", "")
-
-        # table.add_row(Text("Model Queues", style = "bold red"), "")
-        # for model, num_waiting in waiting_dict.items():
-        #     if model.model not in pricing:
-        #         raise ValueError(f"Model {model.model} not found in pricing")
-        #     prices = pricing[model.model]
-        #     table.add_row(Text(f"{model.model}", style="blue"),"")
-        #     table.add_row(f"-TPM limit (k)", str(model.TPM/1000))
-        #     table.add_row(f"-RPM limit (k)", str(model.RPM/1000))
-        #     table.add_row(f"-Num tasks waiting", str(num_waiting))
-        #     token_usage = models_to_tokens[model]
-        #     for cache_status in ['new_token_usage', 'cached_token_usage']:
-        #         table.add_row(Text(f"{cache_status}", style="bold"), "")
-        #         token_usage = getattr(models_to_tokens[model], cache_status)
-        #         for token_type in ["prompt_tokens", "completion_tokens"]:
-        #             tokens = getattr(token_usage, token_type)
-        #             table.add_row(f"-{token_type}", str(tokens))
-        #             table.add_row("Cost", f"${token_usage.cost(prices):.5f}")
-
-        return table

edsl/jobs/base.py

@@ -1,47 +0,0 @@
-from collections import UserDict
-import importlib
-
-from edsl.jobs.runners.JobsRunnerAsyncio import JobsRunnerAsyncio
-from edsl.jobs.runners.JobsRunnerDryRun import JobsRunnerDryRun
-
-from edsl.exceptions import JobsRunError
-from edsl.jobs.JobsRunner import RegisterJobsRunnerMeta
-
-
-class JobsRunnersRegistryDict(UserDict):
-    def __getitem__(self, key):
-        try:
-            return super().__getitem__(key)
-        except KeyError:
-            raise JobsRunError(f"JobsRunner '{key}' not found in registry.")
-
-
-registry_data = RegisterJobsRunnerMeta.lookup()
-JobsRunnersRegistry = JobsRunnersRegistryDict(registry_data)
-
-
-class JobsRunnerDescriptor:
-    def validate(self, value: str) -> None:
-        """Validates the value. If it is invalid, raises an exception. If it is valid, does nothing."""
-        if value not in JobsRunnersRegistry:
-            raise ValueError(
-                f"JobsRunner must be one of {list(JobsRunnersRegistry.keys())}"
-            )
-
-    def __get__(self, instance, owner):
-        """"""
-        if self.name not in instance.__dict__:
-            return None
-        else:
-            return instance.__dict__[self.name]
-
-    def __set__(self, instance, value: str) -> None:
-        self.validate(value, instance)
-        instance.__dict__[self.name] = value
-
-    def __set_name__(self, owner, name: str) -> None:
-        self.name = "_" + name
-
-
-if __name__ == "__main__":
-    pass

edsl/jobs/buckets.py

@@ -1,178 +0,0 @@
-from typing import Union
-import asyncio
-import time
-from collections import UserDict
-from matplotlib import pyplot as plt
-
-
-class TokenBucket:
-    """This is a token bucket used to respect rate limits to services."""
-
-    def __init__(
-        self,
-        *,
-        bucket_name,
-        bucket_type: str,
-        capacity: Union[int, float],
-        refill_rate: Union[int, float],
-    ):
-        self.bucket_name = bucket_name
-        self.bucket_type = bucket_type
-        self.capacity = capacity  # Maximum number of tokens
-        self.tokens = capacity  # Current number of available tokens
-        self.refill_rate = refill_rate  # Rate at which tokens are refilled
-        self.last_refill = time.monotonic()  # Last refill time
-
-        self.log = []
-
-    def __add__(self, other) -> "TokenBucket":
-        """Combine two token buckets. The resulting bucket has the minimum capacity and refill rate of the two buckets.
-        This is useful, for example, if we have two calls to the same model on the same service but have different temperatures.
-        """
-        return TokenBucket(
-            bucket_name=self.bucket_name,
-            bucket_type=self.bucket_type,
-            capacity=min(self.capacity, other.capacity),
-            refil_rate=min(self.refill_rate, other.refill_rate),
-        )
-
-    def __repr__(self):
-        return f"TokenBucket(bucket_name={self.bucket_name}, bucket_type='{self.bucket_type}', capacity={self.capacity}, refill_rate={self.refill_rate})"
-
-    def add_tokens(self, tokens: Union[int, float]) -> None:
-        """Add tokens to the bucket, up to the maximum capacity."""
-        self.tokens = min(self.capacity, self.tokens + tokens)
-        self.log.append((time.monotonic(), self.tokens))
-
-    def refill(self) -> None:
-        """Refill the bucket with new tokens based on elapsed time."""
-        now = time.monotonic()
-        elapsed = now - self.last_refill
-        refill_amount = elapsed * self.refill_rate
-        self.tokens = min(self.capacity, self.tokens + refill_amount)
-        self.last_refill = now
-
-        self.log.append((now, self.tokens))
-
-    def wait_time(self, requested_tokens) -> float:
-        """Calculate the time to wait for the requested number of tokens."""
-        now = time.monotonic()
-        elapsed = now - self.last_refill
-        refill_amount = elapsed * self.refill_rate
-        available_tokens = min(self.capacity, self.tokens + refill_amount)
-        return max(0, requested_tokens - available_tokens) / self.refill_rate
-
-    async def get_tokens(self, amount=1) -> None:
-        """Wait for the specified number of tokens to become available.
-        Note that this method is a coroutine.
-        """
-        if amount > self.capacity:
-            raise ValueError(
-                f"Requested tokens exceed bucket capacity. Bucket capacity: {self.capacity}, requested amount: {amount}"
-            )
-        while self.tokens < amount:
-            self.refill()
-            await asyncio.sleep(0.1)  # Sleep briefly to prevent busy waiting
-        self.tokens -= amount
-
-        now = time.monotonic()
-        self.log.append((now, self.tokens))
-
-    def get_log(self) -> list[tuple]:
-        return self.log
-
-    def visualize(self):
-        """Visualize the token bucket over time."""
-        times, tokens = zip(*self.get_log())
-        start_time = times[0]
-        times = [t - start_time for t in times]  # Normalize time to start from 0
-
-        plt.figure(figsize=(10, 6))
-        plt.plot(times, tokens, label="Tokens Available")
-        plt.xlabel("Time (seconds)", fontsize=12)
-        plt.ylabel("Number of Tokens", fontsize=12)
-        details = f"{self.bucket_name} ({self.bucket_type}) Bucket Usage Over Time\nCapacity: {self.capacity:.1f}, Refill Rate: {self.refill_rate:.1f}/second"
-        plt.title(details, fontsize=14)
-
-        plt.legend()
-        plt.grid(True)
-        plt.tight_layout()
-        plt.show()
-
-
-class ModelBuckets:
-    """A class to represent the token and request buckets for a model.
-    Most LLM model services have limits both on requests-per-minute (RPM) and tokens-per-minute (TPM).
-    A request is one call to the service. The number of tokens required for a request depends on parameters.
-    """
-
-    def __init__(self, requests_bucket: TokenBucket, tokens_bucket: TokenBucket):
-        self.requests_bucket = requests_bucket
-        self.tokens_bucket = tokens_bucket
-
-    def __add__(self, other):
-        return ModelBuckets(
-            requests_bucket=self.requests_bucket + other.requests_bucket,
-            tokens_bucket=self.tokens_bucket + other.tokens_bucket,
-        )
-
-    @classmethod
-    def infinity_bucket(cls, model_name: str = "not_specified") -> "ModelBuckets":
-        """Create a bucket with infinite capacity and refill rate."""
-        return cls(
-            requests_bucket=TokenBucket(
-                bucket_name=model_name, bucket_type="requests", capacity=float("inf"), refill_rate=float("inf")
-            ),
-            tokens_bucket=TokenBucket(
-                bucket_name=model_name, bucket_type="tokens", capacity=float("inf"), refill_rate=float("inf")
-            ),
-        )
-
-    def visualize(self):
-        plot1 = self.requests_bucket.visualize()
-        plot2 = self.tokens_bucket.visualize()
-        return plot1, plot2
-
-    def __repr__(self):
-        return f"ModelBuckets(requests_bucket={self.requests_bucket}, tokens_bucket={self.tokens_bucket})"
-
-
-class BucketCollection(UserDict):
-    """A jobs object will have a whole collection of model buckets, as multiple models could be used.
-    The keys here are the models, and the values are the ModelBuckets objects.
-    Models themselves are hashable, so this works.
-    """
-
-    def __init__(self):
-        super().__init__()
-
-    def __repr__(self):
-        return f"BucketCollection({self.data})"
-
-    def add_model(self, model) -> None:
-        """Adds a model to the bucket collection. This will create the token and request buckets for the model."""
-        # compute the TPS and RPS from the model
-        TPS = model.TPM / 60.0
-        RPS = model.RPM / 60.0
-        # create the buckets
-        requests_bucket = TokenBucket(
-            bucket_name=model.model,
-            bucket_type="requests",
-            capacity=RPS,
-            refill_rate=RPS,
-        )
-        tokens_bucket = TokenBucket(
-            bucket_name=model.model, bucket_type="tokens", capacity=TPS, refill_rate=TPS
-        )
-        model_buckets = ModelBuckets(requests_bucket, tokens_bucket)
-        if model in self:
-            # it if already exists, combine the buckets
-            self[model] += model_buckets
-        else:
-            self[model] = model_buckets
-
-    def visualize(self) -> dict:
-        plots = {}
-        for model in self:
-            plots[model] = self[model].visualize()
-        return plots

edsl/jobs/runners/JobsRunnerDryRun.py

@@ -1,19 +0,0 @@
-import asyncio
-
-from edsl.jobs import Jobs
-from edsl.results import Results, Result
-from edsl.jobs.JobsRunner import JobsRunner
-
-
-class JobsRunnerDryRun(JobsRunner):
-    runner_name = "dryrun"
-
-    def __init__(self, jobs: Jobs):
-        super().__init__(jobs)
-
-    def run(
-        self, n=1, verbose=False, sleep=0, debug=False, progress_bar=False
-    ) -> Results:
-        """Runs a collection of interviews."""
-
-        print(f"This will run {len(self.interviews)} interviews.")

edsl/jobs/runners/JobsRunnerStreaming.py

@@ -1,54 +0,0 @@
-import json
-import threading
-import time
-import uuid
-from edsl.data import CRUD
-from edsl.jobs.JobsRunner import JobsRunner
-from edsl.results import Results
-
-
-class JobsRunnerStreaming(JobsRunner):
-    """This JobRunner conducts interviews serially."""
-
-    runner_name = "streaming"
-
-    def run(
-        self,
-        debug: bool = False,
-        sleep: int = 0,
-        n: int = 1,
-        verbose: bool = False,
-        progress_bar: bool = False,
-    ) -> Results:
-        """
-        Conducts Interviews **serially** and returns their results.
-        - `n`: how many times to run each interview
-        - `debug`: prints debug messages
-        - `verbose`: prints messages
-        - `progress_bar`: shows a progress bar
-        """
-        job_uuid = str(uuid.uuid4())
-        total_results = len(self.interviews)
-
-        def conduct_and_save_interviews():
-            for i, interview in enumerate(self.interviews):
-                answer = interview.conduct_interview(debug=debug)
-                CRUD.write_result(
-                    job_uuid=job_uuid,
-                    result_uuid=str(i),
-                    agent=json.dumps(interview.agent.to_dict()),
-                    scenario=json.dumps(interview.scenario.to_dict()),
-                    model=json.dumps(interview.model.to_dict()),
-                    answer=json.dumps(answer),
-                )
-                time.sleep(sleep)
-
-        interview_thread = threading.Thread(target=conduct_and_save_interviews)
-        interview_thread.start()
-
-        return Results(
-            survey=self.jobs.survey,
-            data=[],
-            job_uuid=job_uuid,
-            total_results=total_results,
-        )

edsl/jobs/task_management.py

@@ -1,215 +0,0 @@
-import asyncio
-import enum
-from typing import Callable
-from collections import UserDict, UserList
-
-from edsl.jobs.buckets import ModelBuckets
-from edsl.jobs.token_tracking import TokenUsage
-from edsl.questions import Question
-
-from edsl.exceptions import InterviewErrorPriorTaskCanceled
-
-
-class TaskStatus(enum.Enum):
-    "These are the possible statuses for a task."
-    NOT_STARTED = enum.auto()
-    WAITING_ON_DEPENDENCIES = enum.auto()
-    CANCELLED = enum.auto()
-    PARENT_FAILED = enum.auto()
-    DEPENDENCIES_COMPLETE = enum.auto()
-    WAITING_FOR_REQUEST_CAPCITY = enum.auto()
-    REQUEST_CAPACITY_ACQUIRED = enum.auto()
-    WAITING_FOR_TOKEN_CAPCITY = enum.auto()
-    TOKEN_CAPACITY_ACQUIRED = enum.auto()
-    API_CALL_IN_PROGRESS = enum.auto()
-    API_CALL_COMPLETE = enum.auto()
-
-
-class InterviewStatusDictionary(UserDict):
-    def __init__(self, data=None):
-        if data:
-            assert all([task_status in data for task_status in TaskStatus])
-            super().__init__(data)
-        else:
-            d = {}
-            for task_status in TaskStatus:
-                d[task_status] = 0
-            d["number_from_cache"] = 0
-            super().__init__(d)
-
-    def __add__(
-        self, other: "InterviewStatusDictionary"
-    ) -> "InterviewStatusDictionary":
-        if not isinstance(other, InterviewStatusDictionary):
-            raise ValueError(f"Can't add {type(other)} to InterviewStatusDictionary")
-        new_dict = {}
-        for key in self.keys():
-            new_dict[key] = self[key] + other[key]
-        return InterviewStatusDictionary(new_dict)
-
-    def __repr__(self):
-        return f"InterviewStatusDictionary({self.data})"
-
-
-
-class TaskStatusDescriptor:
-    def __init__(self):
-        self._task_status = None
-
-    def __get__(self, instance, owner):
-        return self._task_status
-
-    def __set__(self, instance, value):
-        if not isinstance(value, TaskStatus):
-            raise ValueError("Value must be an instance of TaskStatus enum")
-        # logging.info(f"TaskStatus changed for {instance} from {self._task_status} to {value}")
-        self._task_status = value
-
-    def __delete__(self, instance):
-        self._task_status = None
-
-
-class QuestionTaskCreator(UserList):
-    """Class to create and manage question tasks with dependencies.
-    It is a UserList with all the tasks that must be completed before the focal task can be run.
-    When called, it returns an asyncio.Task that depends on the tasks that must be completed before it can be run.
-    """
-
-    task_status = TaskStatusDescriptor()
-
-    def __init__(
-        self,
-        *,
-        question: Question,
-        answer_question_func: Callable,
-        model_buckets: ModelBuckets,
-        token_estimator: Callable = None,
-    ):
-        super().__init__([])
-        self.answer_question_func = answer_question_func
-        self.question = question
-
-        self.model_buckets = model_buckets
-        self.requests_bucket = self.model_buckets.requests_bucket
-        self.tokens_bucket = self.model_buckets.tokens_bucket
-        self.token_estimator = token_estimator
-
-        self.from_cache = False
-
-        self.cached_token_usage = TokenUsage(from_cache=True)
-        self.new_token_usage = TokenUsage(from_cache=False)
-
-        self.task_status = TaskStatus.NOT_STARTED
-
-    def add_dependency(self, task) -> None:
-        """Adds a task dependency to the list of dependencies."""
-        self.append(task)
-
-    def __repr__(self):
-        return f"QuestionTaskCreator for {self.question.question_name}"
-
-    def generate_task(self, debug) -> asyncio.Task:
-        """Creates a task that depends on the passed-in dependencies."""
-        task = asyncio.create_task(self._run_task_async(debug))
-        task.edsl_name = self.question.question_name
-        task.depends_on = [x.edsl_name for x in self]
-        return task
-
-    def estimated_tokens(self) -> int:
-        """Estimates the number of tokens that will be required to run the focal task."""
-        token_estimate = self.token_estimator(self.question)
-        return token_estimate
-
-    def token_usage(self) -> dict:
-        """Returns the token usage for the task."""
-        return {
-            "cached_tokens": self.cached_token_usage,
-            "new_tokens": self.new_token_usage,
-        }
-
-    async def _run_focal_task(self, debug) -> "Answers":
-        """Runs the focal task i.e., the question that we are interested in answering.
-        It is only called after all the dependency tasks are completed.
-        """
-
-        requested_tokens = self.estimated_tokens()
-        if (estimated_wait_time := self.tokens_bucket.wait_time(requested_tokens)) > 0:
-            self.task_status = TaskStatus.WAITING_FOR_TOKEN_CAPCITY
-
-        await self.tokens_bucket.get_tokens(requested_tokens)
-        self.task_status = TaskStatus.TOKEN_CAPACITY_ACQUIRED
-
-        if (estimated_wait_time := self.requests_bucket.wait_time(1)) > 0:
-            self.waiting = True
-            self.task_status = TaskStatus.WAITING_FOR_REQUEST_CAPCITY
-
-        await self.requests_bucket.get_tokens(1)
-        self.task_status = TaskStatus.REQUEST_CAPACITY_ACQUIRED
-
-        self.task_status = TaskStatus.API_CALL_IN_PROGRESS
-        results = await self.answer_question_func(self.question, debug)
-        self.task_status = TaskStatus.API_CALL_COMPLETE
-
-        if "cached_response" in results:
-            if results["cached_response"]:
-                self.tokens_bucket.add_tokens(requested_tokens)
-                self.requests_bucket.add_tokens(1)
-                self.from_cache = True
-
-        tracker = self.cached_token_usage if self.from_cache else self.new_token_usage
-
-        # TODO: This is hacky. The 'func' call should return an object that definitely has a 'usage' key.
-
-        usage = results.get("usage", {"prompt_tokens": 0, "completion_tokens": 0})
-        prompt_tokens = usage.get("prompt_tokens", 0)
-        completion_tokens = usage.get("completion_tokens", 0)
-        tracker.add_tokens(
-            prompt_tokens=prompt_tokens, completion_tokens=completion_tokens
-        )
-
-        return results
-
-    async def _run_task_async(self, debug) -> None:
-        """Runs the task asynchronously, awaiting the tasks that must be completed before this one can be run."""
-        # logger.info(f"Running task for {self.question.question_name}")
-        try:
-            # This is waiting for the tasks that must be completed before this one can be run.
-            # This does *not* use the return_exceptions = True flag, so if any of the tasks fail,
-            # it throws the exception immediately, which is what we want.
-            self.task_status = TaskStatus.WAITING_ON_DEPENDENCIES
-            await asyncio.gather(*self)
-        except asyncio.CancelledError:
-            self.status = TaskStatus.CANCELLED
-            # logger.info(f"Task for {self.question.question_name} was cancelled, most likely because it was skipped.")
-            raise
-        except Exception as e:
-            self.task_status = TaskStatus.PARENT_FAILED
-            # logger.error(f"Required tasks for {self.question.question_name} failed: {e}")
-            # turns the parent exception into a custom exception
-            # So the task gets canceled but this InterviewErrorPriorTaskCanceled exception
-            # So we never get the question details we need.
-            raise InterviewErrorPriorTaskCanceled(
-                f"Required tasks failed for {self.question.question_name}"
-            ) from e
-        else:
-            # logger.info(f"Tasks for {self.question.question_name} completed")
-            # This is the actual task that we want to run.
-            self.task_status = TaskStatus.DEPENDENCIES_COMPLETE
-            return await self._run_focal_task(debug)
-
-
-class TasksList(UserList):
-    def status(self, debug=False):
-        if debug:
-            for task in self:
-                print(f"Task {task.edsl_name}")
-                print(f"\t DEPENDS ON: {task.depends_on}")
-                print(f"\t DONE: {task.done()}")
-                print(f"\t CANCELLED: {task.cancelled()}")
-                if not task.cancelled():
-                    if task.done():
-                        print(f"\t RESULT: {task.result()}")
-                    else:
-                        print(f"\t RESULT: None - Not done yet")
-
-            print("---------------------")