edsl 0.1.36.dev2__py3-none-any.whl → 0.1.36.dev6__py3-none-any.whl

edsl/jobs/runners/JobsRunnerAsyncio.py

@@ -1,343 +1,337 @@
-from __future__ import annotations
-import time
-import math
-import asyncio
-import functools
-import threading
-from typing import Coroutine, List, AsyncGenerator, Optional, Union, Generator
-from contextlib import contextmanager
-from collections import UserList
-
-from rich.live import Live
-from rich.console import Console
-
-from edsl.results.Results import Results
-from edsl import shared_globals
-from edsl.jobs.interviews.Interview import Interview
-from edsl.jobs.runners.JobsRunnerStatus import JobsRunnerStatus
-
-from edsl.jobs.tasks.TaskHistory import TaskHistory
-from edsl.jobs.buckets.BucketCollection import BucketCollection
-from edsl.utilities.decorators import jupyter_nb_handler
-from edsl.data.Cache import Cache
-from edsl.results.Result import Result
-from edsl.results.Results import Results
-from edsl.language_models.LanguageModel import LanguageModel
-from edsl.data.Cache import Cache
-
-
-class StatusTracker(UserList):
-    def __init__(self, total_tasks: int):
-        self.total_tasks = total_tasks
-        super().__init__()
-
-    def current_status(self):
-        return print(f"Completed: {len(self.data)} of {self.total_tasks}", end="\r")
-
-
-class JobsRunnerAsyncio:
-    """A class for running a collection of interviews asynchronously.
-
-    It gets instaniated from a Jobs object.
-    The Jobs object is a collection of interviews that are to be run.
-    """
-
-    def __init__(self, jobs: "Jobs"):
-        self.jobs = jobs
-        self.interviews: List["Interview"] = jobs.interviews()
-        self.bucket_collection: "BucketCollection" = jobs.bucket_collection
-        self.total_interviews: List["Interview"] = []
-
-        # self.jobs_runner_status = JobsRunnerStatus(self, n=1)
-
-    async def run_async_generator(
-        self,
-        cache: Cache,
-        n: int = 1,
-        stop_on_exception: bool = False,
-        sidecar_model: Optional[LanguageModel] = None,
-        total_interviews: Optional[List["Interview"]] = None,
-        raise_validation_errors: bool = False,
-    ) -> AsyncGenerator["Result", None]:
-        """Creates the tasks, runs them asynchronously, and returns the results as a Results object.
-
-        Completed tasks are yielded as they are completed.
-
-        :param n: how many times to run each interview
-        :param stop_on_exception: Whether to stop the interview if an exception is raised
-        :param sidecar_model: a language model to use in addition to the interview's model
-        :param total_interviews: A list of interviews to run can be provided instead.
-        :param raise_validation_errors: Whether to raise validation errors
-        """
-        tasks = []
-        if total_interviews:  # was already passed in total interviews
-            self.total_interviews = total_interviews
-        else:
-            self.total_interviews = list(
-                self._populate_total_interviews(n=n)
-            )  # Populate self.total_interviews before creating tasks
-
-        for interview in self.total_interviews:
-            interviewing_task = self._build_interview_task(
-                interview=interview,
-                stop_on_exception=stop_on_exception,
-                sidecar_model=sidecar_model,
-                raise_validation_errors=raise_validation_errors,
-            )
-            tasks.append(asyncio.create_task(interviewing_task))
-
-        for task in asyncio.as_completed(tasks):
-            result = await task
-            self.jobs_runner_status.add_completed_interview(result)
-            yield result
-
-    def _populate_total_interviews(
-        self, n: int = 1
-    ) -> Generator["Interview", None, None]:
-        """Populates self.total_interviews with n copies of each interview.
-
-        :param n: how many times to run each interview.
-        """
-        for interview in self.interviews:
-            for iteration in range(n):
-                if iteration > 0:
-                    yield interview.duplicate(iteration=iteration, cache=self.cache)
-                else:
-                    interview.cache = self.cache
-                    yield interview
-
-    async def run_async(self, cache: Optional[Cache] = None, n: int = 1) -> Results:
-        """Used for some other modules that have a non-standard way of running interviews."""
-        self.jobs_runner_status = JobsRunnerStatus(self, n=n)
-        self.cache = Cache() if cache is None else cache
-        data = []
-        async for result in self.run_async_generator(cache=self.cache, n=n):
-            data.append(result)
-        return Results(survey=self.jobs.survey, data=data)
-
-    def simple_run(self):
-        data = asyncio.run(self.run_async())
-        return Results(survey=self.jobs.survey, data=data)
-
-    async def _build_interview_task(
-        self,
-        *,
-        interview: Interview,
-        stop_on_exception: bool = False,
-        sidecar_model: Optional["LanguageModel"] = None,
-        raise_validation_errors: bool = False,
-    ) -> "Result":
-        """Conducts an interview and returns the result.
-
-        :param interview: the interview to conduct
-        :param stop_on_exception: stops the interview if an exception is raised
-        :param sidecar_model: a language model to use in addition to the interview's model
-        """
-        # the model buckets are used to track usage rates
-        model_buckets = self.bucket_collection[interview.model]
-
-        # get the results of the interview
-        answer, valid_results = await interview.async_conduct_interview(
-            model_buckets=model_buckets,
-            stop_on_exception=stop_on_exception,
-            sidecar_model=sidecar_model,
-            raise_validation_errors=raise_validation_errors,
-        )
-
-        question_results = {}
-        for result in valid_results:
-            question_results[result.question_name] = result
-
-        answer_key_names = list(question_results.keys())
-
-        generated_tokens_dict = {
-            k + "_generated_tokens": question_results[k].generated_tokens
-            for k in answer_key_names
-        }
-        comments_dict = {
-            k + "_comment": question_results[k].comment for k in answer_key_names
-        }
-
-        # we should have a valid result for each question
-        answer_dict = {k: answer[k] for k in answer_key_names}
-        assert len(valid_results) == len(answer_key_names)
-
-        # TODO: move this down into Interview
-        question_name_to_prompts = dict({})
-        for result in valid_results:
-            question_name = result.question_name
-            question_name_to_prompts[question_name] = {
-                "user_prompt": result.prompts["user_prompt"],
-                "system_prompt": result.prompts["system_prompt"],
-            }
-
-        prompt_dictionary = {}
-        for answer_key_name in answer_key_names:
-            prompt_dictionary[
-                answer_key_name + "_user_prompt"
-            ] = question_name_to_prompts[answer_key_name]["user_prompt"]
-            prompt_dictionary[
-                answer_key_name + "_system_prompt"
-            ] = question_name_to_prompts[answer_key_name]["system_prompt"]
-
-        raw_model_results_dictionary = {}
-        for result in valid_results:
-            question_name = result.question_name
-            raw_model_results_dictionary[
-                question_name + "_raw_model_response"
-            ] = result.raw_model_response
-            raw_model_results_dictionary[question_name + "_cost"] = result.cost
-            one_use_buys = (
-                "NA"
-                if isinstance(result.cost, str)
-                or result.cost == 0
-                or result.cost is None
-                else 1.0 / result.cost
-            )
-            raw_model_results_dictionary[question_name + "_one_usd_buys"] = one_use_buys
-
-        result = Result(
-            agent=interview.agent,
-            scenario=interview.scenario,
-            model=interview.model,
-            iteration=interview.iteration,
-            answer=answer_dict,
-            prompt=prompt_dictionary,
-            raw_model_response=raw_model_results_dictionary,
-            survey=interview.survey,
-            generated_tokens=generated_tokens_dict,
-            comments_dict=comments_dict,
-        )
-        result.interview_hash = hash(interview)
-
-        return result
-
-    @property
-    def elapsed_time(self):
-        return time.monotonic() - self.start_time
-
-    def process_results(
-        self, raw_results: Results, cache: Cache, print_exceptions: bool
-    ):
-        interview_lookup = {
-            hash(interview): index
-            for index, interview in enumerate(self.total_interviews)
-        }
-        interview_hashes = list(interview_lookup.keys())
-
-        results = Results(
-            survey=self.jobs.survey,
-            data=sorted(
-                raw_results, key=lambda x: interview_hashes.index(x.interview_hash)
-            ),
-        )
-        results.cache = cache
-        results.task_history = TaskHistory(
-            self.total_interviews, include_traceback=False
-        )
-        results.has_unfixed_exceptions = results.task_history.has_unfixed_exceptions
-        results.bucket_collection = self.bucket_collection
-
-        if results.has_unfixed_exceptions and print_exceptions:
-            from edsl.scenarios.FileStore import HTMLFileStore
-            from edsl.config import CONFIG
-            from edsl.coop.coop import Coop
-
-            msg = f"Exceptions were raised in {len(results.task_history.indices)} out of {len(self.total_interviews)} interviews.\n"
-
-            if len(results.task_history.indices) > 5:
-                msg += f"Exceptions were raised in the following interviews: {results.task_history.indices}.\n"
-
-            print(msg)
-            # this is where exceptions are opening up
-            filepath = results.task_history.html(
-                cta="Open report to see details.",
-                open_in_browser=True,
-                return_link=True,
-            )
-
-            try:
-                coop = Coop()
-                user_edsl_settings = coop.edsl_settings
-                remote_logging = user_edsl_settings["remote_logging"]
-            except Exception as e:
-                print(e)
-                remote_logging = False
-            if remote_logging:
-                filestore = HTMLFileStore(filepath)
-                coop_details = filestore.push(description="Error report")
-                print(coop_details)
-
-            print("Also see: https://docs.expectedparrot.com/en/latest/exceptions.html")
-
-        return results
-
-    @jupyter_nb_handler
-    async def run(
-        self,
-        cache: Union[Cache, False, None],
-        n: int = 1,
-        stop_on_exception: bool = False,
-        progress_bar: bool = False,
-        sidecar_model: Optional[LanguageModel] = None,
-        print_exceptions: bool = True,
-        raise_validation_errors: bool = False,
-    ) -> "Coroutine":
-        """Runs a collection of interviews, handling both async and sync contexts."""
-
-        self.results = []
-        self.start_time = time.monotonic()
-        self.completed = False
-        self.cache = cache
-        self.sidecar_model = sidecar_model
-
-        self.jobs_runner_status = JobsRunnerStatus(self, n=n)
-
-        stop_event = threading.Event()
-
-        async def process_results(cache):
-            """Processes results from interviews."""
-            async for result in self.run_async_generator(
-                n=n,
-                stop_on_exception=stop_on_exception,
-                cache=cache,
-                sidecar_model=sidecar_model,
-                raise_validation_errors=raise_validation_errors,
-            ):
-                self.results.append(result)
-            self.completed = True
-
-        def run_progress_bar(stop_event):
-            """Runs the progress bar in a separate thread."""
-            self.jobs_runner_status.update_progress(stop_event)
-
-        if progress_bar:
-            progress_thread = threading.Thread(
-                target=run_progress_bar, args=(stop_event,)
-            )
-            progress_thread.start()
-
-        exception_to_raise = None
-        try:
-            with cache as c:
-                await process_results(cache=c)
-        except KeyboardInterrupt:
-            print("Keyboard interrupt received. Stopping gracefully...")
-            stop_event.set()
-        except Exception as e:
-            if stop_on_exception:
-                exception_to_raise = e
-            stop_event.set()
-        finally:
-            stop_event.set()
-            if progress_bar:
-                # self.jobs_runner_status.stop_event.set()
-                if progress_thread:
-                    progress_thread.join()
-
-            if exception_to_raise:
-                raise exception_to_raise
-
-            return self.process_results(
-                raw_results=self.results, cache=cache, print_exceptions=print_exceptions
-            )
+from __future__ import annotations
+import time
+import asyncio
+import threading
+from typing import Coroutine, List, AsyncGenerator, Optional, Union, Generator
+from contextlib import contextmanager
+from collections import UserList
+
+from edsl.results.Results import Results
+from edsl.jobs.interviews.Interview import Interview
+from edsl.jobs.runners.JobsRunnerStatus import JobsRunnerStatus
+
+from edsl.jobs.tasks.TaskHistory import TaskHistory
+from edsl.jobs.buckets.BucketCollection import BucketCollection
+from edsl.utilities.decorators import jupyter_nb_handler
+from edsl.data.Cache import Cache
+from edsl.results.Result import Result
+from edsl.results.Results import Results
+from edsl.language_models.LanguageModel import LanguageModel
+from edsl.data.Cache import Cache
+
+class StatusTracker(UserList):
+    def __init__(self, total_tasks: int):
+        self.total_tasks = total_tasks
+        super().__init__()
+
+    def current_status(self):
+        return print(f"Completed: {len(self.data)} of {self.total_tasks}", end="\r")
+
+
+class JobsRunnerAsyncio:
+    """A class for running a collection of interviews asynchronously.
+
+    It gets instaniated from a Jobs object.
+    The Jobs object is a collection of interviews that are to be run.
+    """
+
+    def __init__(self, jobs: "Jobs"):
+        self.jobs = jobs
+        self.interviews: List["Interview"] = jobs.interviews()
+        self.bucket_collection: "BucketCollection" = jobs.bucket_collection
+        self.total_interviews: List["Interview"] = []
+
+    async def run_async_generator(
+        self,
+        cache: Cache,
+        n: int = 1,
+        stop_on_exception: bool = False,
+        sidecar_model: Optional[LanguageModel] = None,
+        total_interviews: Optional[List["Interview"]] = None,
+        raise_validation_errors: bool = False,
+    ) -> AsyncGenerator["Result", None]:
+        """Creates the tasks, runs them asynchronously, and returns the results as a Results object.
+
+        Completed tasks are yielded as they are completed.
+
+        :param n: how many times to run each interview
+        :param stop_on_exception: Whether to stop the interview if an exception is raised
+        :param sidecar_model: a language model to use in addition to the interview's model
+        :param total_interviews: A list of interviews to run can be provided instead.
+        :param raise_validation_errors: Whether to raise validation errors
+        """
+        tasks = []
+        if total_interviews:  # was already passed in total interviews
+            self.total_interviews = total_interviews
+        else:
+            self.total_interviews = list(
+                self._populate_total_interviews(n=n)
+            )  # Populate self.total_interviews before creating tasks
+
+        for interview in self.total_interviews:
+            interviewing_task = self._build_interview_task(
+                interview=interview,
+                stop_on_exception=stop_on_exception,
+                sidecar_model=sidecar_model,
+                raise_validation_errors=raise_validation_errors,
+            )
+            tasks.append(asyncio.create_task(interviewing_task))
+
+        for task in asyncio.as_completed(tasks):
+            result = await task
+            self.jobs_runner_status.add_completed_interview(result)
+            yield result
+
+    def _populate_total_interviews(
+        self, n: int = 1
+    ) -> Generator["Interview", None, None]:
+        """Populates self.total_interviews with n copies of each interview.
+
+        :param n: how many times to run each interview.
+        """
+        for interview in self.interviews:
+            for iteration in range(n):
+                if iteration > 0:
+                    yield interview.duplicate(iteration=iteration, cache=self.cache)
+                else:
+                    interview.cache = self.cache
+                    yield interview
+
+    async def run_async(self, cache: Optional[Cache] = None, n: int = 1) -> Results:
+        """Used for some other modules that have a non-standard way of running interviews."""
+        self.jobs_runner_status = JobsRunnerStatus(self, n=n)
+        self.cache = Cache() if cache is None else cache
+        data = []
+        async for result in self.run_async_generator(cache=self.cache, n=n):
+            data.append(result)
+        return Results(survey=self.jobs.survey, data=data)
+
+    def simple_run(self):
+        data = asyncio.run(self.run_async())
+        return Results(survey=self.jobs.survey, data=data)
+
+    async def _build_interview_task(
+        self,
+        *,
+        interview: Interview,
+        stop_on_exception: bool = False,
+        sidecar_model: Optional["LanguageModel"] = None,
+        raise_validation_errors: bool = False,
+    ) -> "Result":
+        """Conducts an interview and returns the result.
+
+        :param interview: the interview to conduct
+        :param stop_on_exception: stops the interview if an exception is raised
+        :param sidecar_model: a language model to use in addition to the interview's model
+        """
+        # the model buckets are used to track usage rates
+        model_buckets = self.bucket_collection[interview.model]
+
+        # get the results of the interview
+        answer, valid_results = await interview.async_conduct_interview(
+            model_buckets=model_buckets,
+            stop_on_exception=stop_on_exception,
+            sidecar_model=sidecar_model,
+            raise_validation_errors=raise_validation_errors,
+        )
+
+        question_results = {}
+        for result in valid_results:
+            question_results[result.question_name] = result
+
+        answer_key_names = list(question_results.keys())
+
+        generated_tokens_dict = {
+            k + "_generated_tokens": question_results[k].generated_tokens
+            for k in answer_key_names
+        }
+        comments_dict = {
+            k + "_comment": question_results[k].comment for k in answer_key_names
+        }
+
+        # we should have a valid result for each question
+        answer_dict = {k: answer[k] for k in answer_key_names}
+        assert len(valid_results) == len(answer_key_names)
+
+        # TODO: move this down into Interview
+        question_name_to_prompts = dict({})
+        for result in valid_results:
+            question_name = result.question_name
+            question_name_to_prompts[question_name] = {
+                "user_prompt": result.prompts["user_prompt"],
+                "system_prompt": result.prompts["system_prompt"],
+            }
+
+        prompt_dictionary = {}
+        for answer_key_name in answer_key_names:
+            prompt_dictionary[answer_key_name + "_user_prompt"] = (
+                question_name_to_prompts[answer_key_name]["user_prompt"]
+            )
+            prompt_dictionary[answer_key_name + "_system_prompt"] = (
+                question_name_to_prompts[answer_key_name]["system_prompt"]
+            )
+
+        raw_model_results_dictionary = {}
+        cache_used_dictionary = {}
+        for result in valid_results:
+            question_name = result.question_name
+            raw_model_results_dictionary[question_name + "_raw_model_response"] = (
+                result.raw_model_response
+            )
+            raw_model_results_dictionary[question_name + "_cost"] = result.cost
+            one_use_buys = (
+                "NA"
+                if isinstance(result.cost, str)
+                or result.cost == 0
+                or result.cost is None
+                else 1.0 / result.cost
+            )
+            raw_model_results_dictionary[question_name + "_one_usd_buys"] = one_use_buys
+            cache_used_dictionary[question_name] = result.cache_used
+
+        result = Result(
+            agent=interview.agent,
+            scenario=interview.scenario,
+            model=interview.model,
+            iteration=interview.iteration,
+            answer=answer_dict,
+            prompt=prompt_dictionary,
+            raw_model_response=raw_model_results_dictionary,
+            survey=interview.survey,
+            generated_tokens=generated_tokens_dict,
+            comments_dict=comments_dict,
+            cache_used_dict=cache_used_dictionary,
+        )
+        result.interview_hash = hash(interview)
+
+        return result
+
+    @property
+    def elapsed_time(self):
+        return time.monotonic() - self.start_time
+
+    def process_results(
+        self, raw_results: Results, cache: Cache, print_exceptions: bool
+    ):
+        interview_lookup = {
+            hash(interview): index
+            for index, interview in enumerate(self.total_interviews)
+        }
+        interview_hashes = list(interview_lookup.keys())
+
+        task_history = TaskHistory(self.total_interviews, include_traceback=False)
+
+        results = Results(
+            survey=self.jobs.survey,
+            data=sorted(
+                raw_results, key=lambda x: interview_hashes.index(x.interview_hash)
+            ),
+            task_history=task_history,
+            cache=cache,
+        )
+        results.bucket_collection = self.bucket_collection
+
+        if results.has_unfixed_exceptions and print_exceptions:
+            from edsl.scenarios.FileStore import HTMLFileStore
+            from edsl.config import CONFIG
+            from edsl.coop.coop import Coop
+
+            msg = f"Exceptions were raised in {len(results.task_history.indices)} out of {len(self.total_interviews)} interviews.\n"
+
+            if len(results.task_history.indices) > 5:
+                msg += f"Exceptions were raised in the following interviews: {results.task_history.indices}.\n"
+
+            print(msg)
+            # this is where exceptions are opening up
+            filepath = results.task_history.html(
+                cta="Open report to see details.",
+                open_in_browser=True,
+                return_link=True,
+            )
+
+            try:
+                coop = Coop()
+                user_edsl_settings = coop.edsl_settings
+                remote_logging = user_edsl_settings["remote_logging"]
+            except Exception as e:
+                print(e)
+                remote_logging = False
+
+            if remote_logging:
+                filestore = HTMLFileStore(filepath)
+                coop_details = filestore.push(description="Error report")
+                print(coop_details)
+
+            print("Also see: https://docs.expectedparrot.com/en/latest/exceptions.html")
+
+        return results
+
+    @jupyter_nb_handler
+    async def run(
+        self,
+        cache: Union[Cache, False, None],
+        n: int = 1,
+        stop_on_exception: bool = False,
+        progress_bar: bool = False,
+        sidecar_model: Optional[LanguageModel] = None,
+        print_exceptions: bool = True,
+        raise_validation_errors: bool = False,
+    ) -> "Coroutine":
+        """Runs a collection of interviews, handling both async and sync contexts."""
+
+        self.results = []
+        self.start_time = time.monotonic()
+        self.completed = False
+        self.cache = cache
+        self.sidecar_model = sidecar_model
+
+        self.jobs_runner_status = JobsRunnerStatus(self, n=n)
+
+        stop_event = threading.Event()
+
+        async def process_results(cache):
+            """Processes results from interviews."""
+            async for result in self.run_async_generator(
+                n=n,
+                stop_on_exception=stop_on_exception,
+                cache=cache,
+                sidecar_model=sidecar_model,
+                raise_validation_errors=raise_validation_errors,
+            ):
+                self.results.append(result)
+            self.completed = True
+
+        def run_progress_bar(stop_event):
+            """Runs the progress bar in a separate thread."""
+            self.jobs_runner_status.update_progress(stop_event)
+
+        if progress_bar:
+            progress_thread = threading.Thread(
+                target=run_progress_bar, args=(stop_event,)
+            )
+            progress_thread.start()
+
+        exception_to_raise = None
+        try:
+            with cache as c:
+                await process_results(cache=c)
+        except KeyboardInterrupt:
+            print("Keyboard interrupt received. Stopping gracefully...")
+            stop_event.set()
+        except Exception as e:
+            if stop_on_exception:
+                exception_to_raise = e
+            stop_event.set()
+        finally:
+            stop_event.set()
+            if progress_bar:
+                # self.jobs_runner_status.stop_event.set()
+                if progress_thread:
+                    progress_thread.join()
+
+            if exception_to_raise:
+                raise exception_to_raise
+
+            return self.process_results(
+                raw_results=self.results, cache=cache, print_exceptions=print_exceptions
+            )