edsl 0.1.27.dev2__py3-none-any.whl → 0.1.29__py3-none-any.whl

edsl/jobs/interviews/Interview.py

@@ -6,15 +6,9 @@ import asyncio
 import time
 from typing import Any, Type, List, Generator, Optional
 
-from edsl.agents import Agent
-from edsl.language_models import LanguageModel
-from edsl.scenarios import Scenario
-from edsl.surveys import Survey
-
 from edsl.jobs.Answers import Answers
 from edsl.surveys.base import EndOfSurvey
 from edsl.jobs.buckets.ModelBuckets import ModelBuckets
-
 from edsl.jobs.tasks.TaskCreators import TaskCreators
 
 from edsl.jobs.interviews.InterviewStatusLog import InterviewStatusLog
@@ -60,9 +54,9 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         self.debug = debug
         self.iteration = iteration
         self.cache = cache
-        self.answers: dict[
-            str, str
-        ] = Answers()  # will get filled in as interview progresses
+        self.answers: dict[str, str] = (
+            Answers()
+        )  # will get filled in as interview progresses
         self.sidecar_model = sidecar_model
 
         # Trackers
@@ -101,9 +95,12 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         self.sidecar_model = sidecar_model
 
         # if no model bucket is passed, create an 'infinity' bucket with no rate limits
-        if model_buckets is None:
+        # print("model_buckets", model_buckets)
+        if model_buckets is None or hasattr(self.agent, "answer_question_directly"):
             model_buckets = ModelBuckets.infinity_bucket()
 
+        # model_buckets = ModelBuckets.infinity_bucket()
+
         ## build the tasks using the InterviewTaskBuildingMixin
         ## This is the key part---it creates a task for each question,
         ## with dependencies on the questions that must be answered before this one can be answered.

edsl/jobs/interviews/InterviewStatusMixin.py

@@ -17,9 +17,9 @@ class InterviewStatusMixin:
         The keys are the question names; the values are the lists of status log changes for each task.
         """
         for task_creator in self.task_creators.values():
-            self._task_status_log_dict[
-                task_creator.question.question_name
-            ] = task_creator.status_log
+            self._task_status_log_dict[task_creator.question.question_name] = (
+                task_creator.status_log
+            )
         return self._task_status_log_dict
 
     @property

edsl/jobs/interviews/InterviewTaskBuildingMixin.py

@@ -5,17 +5,19 @@ import asyncio
 import time
 import traceback
 from typing import Generator, Union
+
 from edsl import CONFIG
 from edsl.exceptions import InterviewTimeoutError
-from edsl.data_transfer_models import AgentResponseDict
-from edsl.questions.QuestionBase import QuestionBase
+
+# from edsl.questions.QuestionBase import QuestionBase
 from edsl.surveys.base import EndOfSurvey
 from edsl.jobs.buckets.ModelBuckets import ModelBuckets
 from edsl.jobs.interviews.interview_exception_tracking import InterviewExceptionEntry
 from edsl.jobs.interviews.retry_management import retry_strategy
 from edsl.jobs.tasks.task_status_enum import TaskStatus
 from edsl.jobs.tasks.QuestionTaskCreator import QuestionTaskCreator
-from edsl.agents.InvigilatorBase import InvigilatorBase
+
+# from edsl.agents.InvigilatorBase import InvigilatorBase
 
 TIMEOUT = float(CONFIG.get("EDSL_API_TIMEOUT"))
 
@@ -44,6 +46,7 @@ class InterviewTaskBuildingMixin:
             scenario=self.scenario,
             model=self.model,
             debug=debug,
+            survey=self.survey,
             memory_plan=self.survey.memory_plan,
             current_answers=self.answers,
             iteration=self.iteration,
@@ -149,29 +152,45 @@ class InterviewTaskBuildingMixin:
     async def _answer_question_and_record_task(
         self,
         *,
-        question: QuestionBase,
+        question: "QuestionBase",
         debug: bool,
         task=None,
-    ) -> AgentResponseDict:
+    ) -> "AgentResponseDict":
         """Answer a question and records the task.
 
         This in turn calls the the passed-in agent's async_answer_question method, which returns a response dictionary.
         Note that is updates answers dictionary with the response.
         """
-        invigilator = self._get_invigilator(question, debug=debug)
+        from edsl.data_transfer_models import AgentResponseDict
+
+        try:
+            invigilator = self._get_invigilator(question, debug=debug)
 
-        if self._skip_this_question(question):
-            return invigilator.get_failed_task_result()
+            if self._skip_this_question(question):
+                return invigilator.get_failed_task_result()
 
-        response: AgentResponseDict = await self._attempt_to_answer_question(
-            invigilator, task
-        )
+            response: AgentResponseDict = await self._attempt_to_answer_question(
+                invigilator, task
+            )
 
-        self._add_answer(response=response, question=question)
+            self._add_answer(response=response, question=question)
 
-        # With the answer to the question, we can now cancel any skipped questions
-        self._cancel_skipped_questions(question)
-        return AgentResponseDict(**response)
+            # With the answer to the question, we can now cancel any skipped questions
+            self._cancel_skipped_questions(question)
+            return AgentResponseDict(**response)
+        except Exception as e:
+            raise e
+            # import traceback
+            # print("Exception caught:")
+            # traceback.print_exc()
+
+            # # Extract and print the traceback info
+            # tb = e.__traceback__
+            # while tb is not None:
+            #     print(f"File {tb.tb_frame.f_code.co_filename}, line {tb.tb_lineno}, in {tb.tb_frame.f_code.co_name}")
+            #     tb = tb.tb_next
+            #     breakpoint()
+            #     raise e
 
     def _add_answer(self, response: AgentResponseDict, question: QuestionBase) -> None:
         """Add the answer to the answers dictionary.
@@ -239,11 +258,11 @@ class InterviewTaskBuildingMixin:
         """
         current_question_index: int = self.to_index[current_question.question_name]
 
-        next_question: Union[
-            int, EndOfSurvey
-        ] = self.survey.rule_collection.next_question(
-            q_now=current_question_index,
-            answers=self.answers | self.scenario | self.agent["traits"],
+        next_question: Union[int, EndOfSurvey] = (
+            self.survey.rule_collection.next_question(
+                q_now=current_question_index,
+                answers=self.answers | self.scenario | self.agent["traits"],
+            )
         )
 
         next_question_index = next_question.next_q

edsl/jobs/interviews/retry_management.py

@@ -13,17 +13,17 @@ EDSL_MAX_BACKOFF_SEC = float(CONFIG.get("EDSL_MAX_BACKOFF_SEC"))
 EDSL_MAX_ATTEMPTS = int(CONFIG.get("EDSL_MAX_ATTEMPTS"))
 
 
-def print_retry(retry_state, print_to_terminal=False):
+def print_retry(retry_state, print_to_terminal=True):
     "Prints details on tenacity retries."
     attempt_number = retry_state.attempt_number
     exception = retry_state.outcome.exception()
     wait_time = retry_state.next_action.sleep
-    # breakpoint()
     if print_to_terminal:
         print(
-            f"Attempt {attempt_number} failed with exception: {repr(exception)}; "
+            f"Attempt {attempt_number} failed with exception:" f"{exception}",
             f"now waiting {wait_time:.2f} seconds before retrying."
-            f" Parameters: start={EDSL_BACKOFF_START_SEC}, max={EDSL_MAX_BACKOFF_SEC}, max_attempts={EDSL_MAX_ATTEMPTS}."
+            f"Parameters: start={EDSL_BACKOFF_START_SEC}, max={EDSL_MAX_BACKOFF_SEC}, max_attempts={EDSL_MAX_ATTEMPTS}."
+            "\n\n",
         )
 
 

edsl/jobs/runners/JobsRunnerAsyncio.py

@@ -1,26 +1,17 @@
 from __future__ import annotations
 import time
 import asyncio
-import textwrap
+import time
 from contextlib import contextmanager
 
 from typing import Coroutine, List, AsyncGenerator, Optional, Union
 
-from rich.live import Live
-from rich.console import Console
-
 from edsl import shared_globals
-from edsl.results import Results, Result
-
 from edsl.jobs.interviews.Interview import Interview
-from edsl.utilities.decorators import jupyter_nb_handler
-from edsl.jobs.Jobs import Jobs
 from edsl.jobs.runners.JobsRunnerStatusMixin import JobsRunnerStatusMixin
-from edsl.language_models import LanguageModel
-from edsl.data.Cache import Cache
-
 from edsl.jobs.tasks.TaskHistory import TaskHistory
 from edsl.jobs.buckets.BucketCollection import BucketCollection
+from edsl.utilities.decorators import jupyter_nb_handler
 
 
 class JobsRunnerAsyncio(JobsRunnerStatusMixin):
@@ -32,20 +23,20 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
 
     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"] = []
 
-    async def run_async(
+    async def run_async_generator(
         self,
-        cache: Cache,
+        cache: "Cache",
         n: int = 1,
         debug: bool = False,
         stop_on_exception: bool = False,
         sidecar_model: "LanguageModel" = None,
         total_interviews: Optional[List["Interview"]] = None,
-    ) -> AsyncGenerator[Result, None]:
+    ) -> 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.
@@ -96,6 +87,20 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
                     )  # set the cache for the first interview
                     self.total_interviews.append(interview)
 
+    async def run_async(self, cache=None) -> Results:
+        if cache is None:
+            self.cache = Cache()
+        else:
+            self.cache = cache
+        data = []
+        async for result in self.run_async_generator(cache=self.cache):
+            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,
         *,
@@ -138,19 +143,21 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
 
         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"]
+            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 + "_raw_model_response"] = (
+                result["raw_model_response"]
+            )
+
+        from edsl.results.Result import Result
 
         result = Result(
             agent=interview.agent,
@@ -180,6 +187,8 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
         print_exceptions: bool = True,
     ) -> "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()
@@ -187,39 +196,15 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
         self.cache = cache
         self.sidecar_model = sidecar_model
 
-        def generate_table():
-            return self.status_table(self.results, self.elapsed_time)
+        from edsl.results.Results import Results
 
-        @contextmanager
-        def no_op_cm():
-            """A no-op context manager with a dummy update method."""
-            yield DummyLive()
-
-        class DummyLive:
-            def update(self, *args, **kwargs):
-                """A dummy update method that does nothing."""
-                pass
-
-        progress_bar_context = (
-            Live(generate_table(), console=console, refresh_per_second=5)
-            if progress_bar
-            else no_op_cm()
-        )
-
-        with cache as c:
-            with progress_bar_context as live:
-
-                async def update_progress_bar():
-                    """Updates the progress bar at fixed intervals."""
-                    while True:
-                        live.update(generate_table())
-                        await asyncio.sleep(0.1)  # Update interval
-                        if self.completed:
-                            break
+        if not progress_bar:
+            # print("Running without progress bar")
+            with cache as c:
 
                 async def process_results():
                     """Processes results from interviews."""
-                    async for result in self.run_async(
+                    async for result in self.run_async_generator(
                         n=n,
                         debug=debug,
                         stop_on_exception=stop_on_exception,
@@ -227,25 +212,76 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
                         sidecar_model=sidecar_model,
                     ):
                         self.results.append(result)
-                        live.update(generate_table())
                     self.completed = True
 
-                progress_task = asyncio.create_task(update_progress_bar())
+                await asyncio.gather(process_results())
 
-                try:
-                    await asyncio.gather(process_results(), progress_task)
-                except asyncio.CancelledError:
+            results = Results(survey=self.jobs.survey, data=self.results)
+        else:
+            # print("Running with progress bar")
+            from rich.live import Live
+            from rich.console import Console
+
+            def generate_table():
+                return self.status_table(self.results, self.elapsed_time)
+
+            @contextmanager
+            def no_op_cm():
+                """A no-op context manager with a dummy update method."""
+                yield DummyLive()
+
+            class DummyLive:
+                def update(self, *args, **kwargs):
+                    """A dummy update method that does nothing."""
                     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
+            progress_bar_context = (
+                Live(generate_table(), console=console, refresh_per_second=5)
+                if progress_bar
+                else no_op_cm()
+            )
 
-                    # one more update
-                    live.update(generate_table())
+            with cache as c:
+                with progress_bar_context as live:
+
+                    async def update_progress_bar():
+                        """Updates the progress bar at fixed intervals."""
+                        while True:
+                            live.update(generate_table())
+                            await asyncio.sleep(0.00001)  # Update interval
+                            if self.completed:
+                                break
+
+                    async def process_results():
+                        """Processes results from interviews."""
+                        async for result in self.run_async_generator(
+                            n=n,
+                            debug=debug,
+                            stop_on_exception=stop_on_exception,
+                            cache=c,
+                            sidecar_model=sidecar_model,
+                        ):
+                            self.results.append(result)
+                            live.update(generate_table())
+                        self.completed = True
+
+                    progress_task = asyncio.create_task(update_progress_bar())
+
+                    try:
+                        await asyncio.gather(process_results(), progress_task)
+                    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
+
+                        # one more update
+                        live.update(generate_table())
+
+            results = Results(survey=self.jobs.survey, data=self.results)
 
-        results = Results(survey=self.jobs.survey, data=self.results)
         task_history = TaskHistory(self.total_interviews, include_traceback=False)
         results.task_history = task_history
 
@@ -259,6 +295,8 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
                 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]
             )

edsl/jobs/runners/JobsRunnerStatusData.py

@@ -81,7 +81,7 @@ class JobsRunnerStatusData:
         >>> completed_tasks = []
         >>> elapsed_time = 0
         >>> JobsRunnerStatusData().generate_status_summary(completed_tasks, elapsed_time, interviews)
-        {'Elapsed time': '0.0 sec.', 'Total interviews requested': '1 ', 'Completed interviews': '0 ', 'Percent complete': '0 %', 'Average time per interview': 'NA', 'Task remaining': '1 ', 'Estimated time remaining': 'NA', 'model_queues': [{'model_name': 'gpt-4-1106-preview', 'TPM_limit_k': 1600.0, 'RPM_limit_k': 8.0, 'num_tasks_waiting': 0, 'token_usage_info': [{'cache_status': 'new_token_usage', 'details': [{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], 'cost': '$0.00000'}, {'cache_status': 'cached_token_usage', 'details': [{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], 'cost': '$0.00000'}]}]}
+        {'Elapsed time': '0.0 sec.', 'Total interviews requested': '1 ', 'Completed interviews': '0 ', 'Percent complete': '0 %', 'Average time per interview': 'NA', 'Task remaining': '1 ', 'Estimated time remaining': 'NA', 'model_queues': [{'model_name': '...', 'TPM_limit_k': ..., 'RPM_limit_k': ..., 'num_tasks_waiting': 0, 'token_usage_info': [{'cache_status': 'new_token_usage', 'details': [{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], 'cost': '$0.00000'}, {'cache_status': 'cached_token_usage', 'details': [{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], 'cost': '$0.00000'}]}]}
         """
 
         models_to_tokens = defaultdict(InterviewTokenUsage)
@@ -176,7 +176,7 @@ class JobsRunnerStatusData:
         >>> model = interviews[0].model
         >>> num_waiting = 0
         >>> JobsRunnerStatusData()._get_model_info(model, num_waiting, models_to_tokens)
-        {'model_name': 'gpt-4-1106-preview', 'TPM_limit_k': 1600.0, 'RPM_limit_k': 8.0, 'num_tasks_waiting': 0, 'token_usage_info': [{'cache_status': 'new_token_usage', 'details': [{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], 'cost': '$0.00000'}, {'cache_status': 'cached_token_usage', 'details': [{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], 'cost': '$0.00000'}]}
+        {'model_name': 'gpt-4-1106-preview', 'TPM_limit_k': ..., 'RPM_limit_k': ..., 'num_tasks_waiting': 0, 'token_usage_info': [{'cache_status': 'new_token_usage', 'details': [{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], 'cost': '$0.00000'}, {'cache_status': 'cached_token_usage', 'details': [{'type': 'prompt_tokens', 'tokens': 0}, {'type': 'completion_tokens', 'tokens': 0}], 'cost': '$0.00000'}]}
         """
 
         prices = get_token_pricing(model.model)
@@ -234,4 +234,4 @@ class JobsRunnerStatusData:
 if __name__ == "__main__":
     import doctest
 
-    doctest.testmod()
+    doctest.testmod(optionflags=doctest.ELLIPSIS)

edsl/jobs/tasks/QuestionTaskCreator.py

@@ -1,7 +1,7 @@
 import asyncio
 from typing import Callable, Union, List
 from collections import UserList, UserDict
-
+import time
 
 from edsl.jobs.buckets import ModelBuckets
 from edsl.exceptions import InterviewErrorPriorTaskCanceled
@@ -132,7 +132,7 @@ class QuestionTaskCreator(UserList):
             self.waiting = True
             self.task_status = TaskStatus.WAITING_FOR_REQUEST_CAPACITY
 
-        await self.requests_bucket.get_tokens(1)
+        await self.tokens_bucket.get_tokens(1)
 
         self.task_status = TaskStatus.API_CALL_IN_PROGRESS
         try:
@@ -151,6 +151,8 @@ class QuestionTaskCreator(UserList):
                 self.requests_bucket.add_tokens(1)
                 self.from_cache = True
 
+        _ = results.pop("cached_response", None)
+
         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.

edsl/jobs/tasks/TaskHistory.py

@@ -1,8 +1,5 @@
 from edsl.jobs.tasks.task_status_enum import TaskStatus
-from matplotlib import pyplot as plt
 from typing import List, Optional
-
-import matplotlib.pyplot as plt
 from io import BytesIO
 import base64
 
@@ -75,6 +72,8 @@ class TaskHistory:
 
     def plot_completion_times(self):
         """Plot the completion times for each task."""
+        import matplotlib.pyplot as plt
+
         updates = self.get_updates()
 
         elapsed = [update.max_time - update.min_time for update in updates]
@@ -126,6 +125,8 @@ class TaskHistory:
         rows = int(len(TaskStatus) ** 0.5) + 1
         cols = (len(TaskStatus) + rows - 1) // rows  # Ensure all plots fit
 
+        import matplotlib.pyplot as plt
+
         fig, axes = plt.subplots(rows, cols, figsize=(15, 10))
         axes = axes.flatten()  # Flatten in case of a single row/column