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

edsl/jobs/runners/JobsRunnerAsyncio.py

@@ -1,27 +1,42 @@
 from __future__ import annotations
 import time
-import math
 import asyncio
-import functools
-import threading
-from typing import Coroutine, List, AsyncGenerator, Optional, Union, Generator
+import time
 from contextlib import contextmanager
-from collections import UserList
 
-from edsl.results.Results import Results
-from rich.live import Live
-from rich.console import Console
+from typing import Coroutine, List, AsyncGenerator, Optional, Union
 
 from edsl import shared_globals
 from edsl.jobs.interviews.Interview import Interview
-from edsl.jobs.runners.JobsRunnerStatus import JobsRunnerStatus
-
+from edsl.jobs.runners.JobsRunnerStatusMixin import JobsRunnerStatusMixin
 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
+
+import time
+import functools
+
+
+def cache_with_timeout(timeout):
+    def decorator(func):
+        cached_result = {}
+        last_computation_time = [0]  # Using list to store mutable value
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            current_time = time.time()
+            if (current_time - last_computation_time[0]) >= timeout:
+                cached_result["value"] = func(*args, **kwargs)
+                last_computation_time[0] = current_time
+            return cached_result["value"]
+
+        return wrapper
+
+    return decorator
+
+
+# from queue import Queue
+from collections import UserList
 
 
 class StatusTracker(UserList):
@@ -33,87 +48,99 @@ class StatusTracker(UserList):
         return print(f"Completed: {len(self.data)} of {self.total_tasks}", end="\r")
 
 
-class JobsRunnerAsyncio:
+class JobsRunnerAsyncio(JobsRunnerStatusMixin):
     """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"):
+    def __init__(self, jobs: Jobs):
         self.jobs = jobs
+        # this creates the interviews, which can take a while
         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,
+        debug: bool = False,
         stop_on_exception: bool = False,
-        sidecar_model: Optional["LanguageModel"] = None,
+        sidecar_model: "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 debug:
         :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
+        if total_interviews:
             self.total_interviews = total_interviews
         else:
-            self.total_interviews = list(
-                self._populate_total_interviews(n=n)
+            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,
+                debug=debug,
                 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]:
+    def _populate_total_interviews(self, n: int = 1) -> None:
         """Populates self.total_interviews with n copies of each interview.
 
         :param n: how many times to run each interview.
         """
+        # TODO: Why not return a list of interviews instead of modifying the object?
+
+        self.total_interviews = []
         for interview in self.interviews:
             for iteration in range(n):
                 if iteration > 0:
-                    yield interview.duplicate(iteration=iteration, cache=self.cache)
+                    new_interview = interview.duplicate(
+                        iteration=iteration, cache=self.cache
+                    )
+                    self.total_interviews.append(new_interview)
                 else:
-                    interview.cache = self.cache
-                    yield interview
+                    interview.cache = (
+                        self.cache
+                    )  # set the cache for the first interview
+                    self.total_interviews.append(interview)
+
+    async def run_async(self, cache=None, n=1) -> Results:
+        from edsl.results.Results import Results
 
-    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
+        # breakpoint()
+        # tracker = StatusTracker(total_tasks=len(self.interviews))
+
+        if cache is None:
+            self.cache = Cache()
+        else:
+            self.cache = 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):
+        from edsl.results.Results import Results
+
         data = asyncio.run(self.run_async())
         return Results(survey=self.jobs.survey, data=data)
 
@@ -121,13 +148,14 @@ class JobsRunnerAsyncio:
         self,
         *,
         interview: Interview,
+        debug: bool,
         stop_on_exception: bool = False,
-        sidecar_model: Optional["LanguageModel"] = None,
-        raise_validation_errors: bool = False,
-    ) -> "Result":
+        sidecar_model: Optional[LanguageModel] = None,
+    ) -> Result:
         """Conducts an interview and returns the result.
 
         :param interview: the interview to conduct
+        :param debug: prints debug messages
         :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
         """
@@ -136,37 +164,24 @@ class JobsRunnerAsyncio:
 
         # get the results of the interview
         answer, valid_results = await interview.async_conduct_interview(
+            debug=debug,
             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}
+        answer_key_names = {k for k in set(answer.keys()) if not k.endswith("_comment")}
+
         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 = result["question_name"]
             question_name_to_prompts[question_name] = {
-                "user_prompt": result.prompts["user_prompt"],
-                "system_prompt": result.prompts["system_prompt"],
+                "user_prompt": result["prompts"]["user_prompt"],
+                "system_prompt": result["prompts"]["system_prompt"],
             }
 
         prompt_dictionary = {}
@@ -180,31 +195,22 @@ class JobsRunnerAsyncio:
 
         raw_model_results_dictionary = {}
         for result in valid_results:
-            question_name = result.question_name
+            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["raw_model_response"]
+
+        from edsl.results.Result import Result
 
         result = Result(
             agent=interview.agent,
             scenario=interview.scenario,
             model=interview.model,
             iteration=interview.iteration,
-            answer=answer_dict,
+            answer=answer,
             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)
 
@@ -214,109 +220,132 @@ class JobsRunnerAsyncio:
     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,
+        debug: bool = False,
         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."""
+        from rich.console import Console
 
+        console = Console()
         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)
+        from edsl.results.Results import Results
+        from rich.live import Live
+        from rich.console import Console
+
+        @cache_with_timeout(1)
+        def generate_table():
+            return self.status_table(self.results, self.elapsed_time)
 
-        async def process_results(cache):
+        async def process_results(cache, progress_bar_context=None):
             """Processes results from interviews."""
             async for result in self.run_async_generator(
                 n=n,
+                debug=debug,
                 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():
-            """Runs the progress bar in a separate thread."""
-            self.jobs_runner_status.update_progress()
+                if progress_bar_context:
+                    progress_bar_context.update(generate_table())
+                self.completed = True
+
+        async def update_progress_bar(progress_bar_context):
+            """Updates the progress bar at fixed intervals."""
+            if progress_bar_context is None:
+                return
+
+            while True:
+                progress_bar_context.update(generate_table())
+                await asyncio.sleep(0.1)  # Update interval
+                if self.completed:
+                    break
+
+        @contextmanager
+        def conditional_context(condition, context_manager):
+            if condition:
+                with context_manager as cm:
+                    yield cm
+            else:
+                yield
+
+        with conditional_context(
+            progress_bar, Live(generate_table(), console=console, refresh_per_second=1)
+        ) as progress_bar_context:
+            with cache as c:
+                progress_task = asyncio.create_task(
+                    update_progress_bar(progress_bar_context)
+                )
+
+                try:
+                    await asyncio.gather(
+                        progress_task,
+                        process_results(
+                            cache=c, progress_bar_context=progress_bar_context
+                        ),
+                    )
+                except asyncio.CancelledError:
+                    pass
+                finally:
+                    progress_task.cancel()  # Cancel the progress_task when process_results is done
+                    await progress_task
+
+                    await asyncio.sleep(1)  # short delay to show the final status
+
+                    if progress_bar_context:
+                        progress_bar_context.update(generate_table())
+
+        # puts results in the same order as the total interviews
+        interview_hashes = [hash(interview) for interview in self.total_interviews]
+        self.results = sorted(
+            self.results, key=lambda x: interview_hashes.index(x.interview_hash)
+        )
 
-        if progress_bar:
-            progress_thread = threading.Thread(target=run_progress_bar)
-            progress_thread.start()
+        results = Results(survey=self.jobs.survey, data=self.results)
+        task_history = TaskHistory(self.total_interviews, include_traceback=False)
+        results.task_history = task_history
+
+        results.has_exceptions = task_history.has_exceptions
+
+        if results.has_exceptions:
+            # put the failed interviews in the results object as a list
+            failed_interviews = [
+                interview.duplicate(
+                    iteration=interview.iteration, cache=interview.cache
+                )
+                for interview in self.total_interviews
+                if interview.has_exceptions
+            ]
+            from edsl.jobs.Jobs import Jobs
+
+            results.failed_jobs = Jobs.from_interviews(
+                [interview for interview in failed_interviews]
+            )
+            if print_exceptions:
+                msg = f"Exceptions were raised in {len(results.task_history.indices)} out of {len(self.total_interviews)} interviews.\n"
 
-        with cache as c:
-            await process_results(cache=c)
+                if len(results.task_history.indices) > 5:
+                    msg += f"Exceptions were raised in the following interviews: {results.task_history.indices}.\n"
 
-        if progress_bar:
-            progress_thread.join()
+                shared_globals["edsl_runner_exceptions"] = task_history
+                print(msg)
+                # this is where exceptions are opening up
+                task_history.html(cta="Open report to see details.")
+                print(
+                    "Also see: https://docs.expectedparrot.com/en/latest/exceptions.html"
+                )
 
-        return self.process_results(
-            raw_results=self.results, cache=cache, print_exceptions=print_exceptions
-        )
+        return results