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

edsl/jobs/runners/JobsRunnerAsyncio.py

@@ -4,6 +4,7 @@ import asyncio
 import threading
 import warnings
 from typing import Coroutine, List, AsyncGenerator, Optional, Union, Generator, Type
+from uuid import UUID
 from collections import UserList
 
 from edsl.results.Results import Results
@@ -36,6 +37,8 @@ class JobsRunnerAsyncio:
     The Jobs object is a collection of interviews that are to be run.
     """
 
+    MAX_CONCURRENT_DEFAULT = 500
+
     def __init__(self, jobs: "Jobs"):
         self.jobs = jobs
         self.interviews: List["Interview"] = jobs.interviews()
@@ -43,6 +46,53 @@ class JobsRunnerAsyncio:
         self.total_interviews: List["Interview"] = []
         self._initialized = threading.Event()
 
+        from edsl.config import CONFIG
+
+        self.MAX_CONCURRENT = int(CONFIG.get("EDSL_MAX_CONCURRENT_TASKS"))
+        # print(f"MAX_CONCURRENT: {self.MAX_CONCURRENT}")
+
+    # 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
+    #     self._initialized.set()  # Signal that we're ready
+
+    #     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
+
     async def run_async_generator(
         self,
         cache: Cache,
@@ -52,9 +102,10 @@ class JobsRunnerAsyncio:
         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.
+        """Creates and processes tasks asynchronously, yielding results as they complete.
 
-        Completed tasks are yielded as they are completed.
+        Tasks are created and processed in a streaming fashion rather than building the full list upfront.
+        Results are yielded as soon as they are available.
 
         :param n: how many times to run each interview
         :param stop_on_exception: Whether to stop the interview if an exception is raised
@@ -62,29 +113,70 @@ class JobsRunnerAsyncio:
         :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
+        # Initialize interviews iterator
+        if total_interviews:
+            interviews_iter = iter(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
+            interviews_iter = self._populate_total_interviews(n=n)
+            self.total_interviews = list(interviews_iter)
+            interviews_iter = iter(self.total_interviews)  # Create fresh iterator
 
         self._initialized.set()  # Signal that we're ready
 
-        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))
+        # Keep track of active tasks
+        active_tasks = set()
 
-        for task in asyncio.as_completed(tasks):
-            result = await task
-            self.jobs_runner_status.add_completed_interview(result)
-            yield result
+        try:
+            while True:
+                # Add new tasks if we're below max_concurrent and there are more interviews
+                while len(active_tasks) < self.MAX_CONCURRENT:
+                    try:
+                        interview = next(interviews_iter)
+                        task = asyncio.create_task(
+                            self._build_interview_task(
+                                interview=interview,
+                                stop_on_exception=stop_on_exception,
+                                sidecar_model=sidecar_model,
+                                raise_validation_errors=raise_validation_errors,
+                            )
+                        )
+                        active_tasks.add(task)
+                        # Add callback to remove task from set when done
+                        task.add_done_callback(active_tasks.discard)
+                    except StopIteration:
+                        break
+
+                if not active_tasks:
+                    break
+
+                # Wait for next completed task
+                done, _ = await asyncio.wait(
+                    active_tasks, return_when=asyncio.FIRST_COMPLETED
+                )
+
+                # Process completed tasks
+                for task in done:
+                    try:
+                        result = await task
+                        self.jobs_runner_status.add_completed_interview(result)
+                        yield result
+                    except Exception as e:
+                        if stop_on_exception:
+                            # Cancel remaining tasks
+                            for t in active_tasks:
+                                if not t.done():
+                                    t.cancel()
+                            raise
+                        else:
+                            # Log error and continue
+                            # logger.error(f"Task failed with error: {e}")
+                            continue
+        finally:
+            # Ensure we cancel any remaining tasks if we exit early
+            for task in active_tasks:
+                if not task.done():
+                    task.cancel()
 
     def _populate_total_interviews(
         self, n: int = 1
@@ -168,20 +260,20 @@ class JobsRunnerAsyncio:
 
         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 = {}
         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 + "_raw_model_response"
+            ] = result.raw_model_response
             raw_model_results_dictionary[question_name + "_cost"] = result.cost
             one_use_buys = (
                 "NA"
@@ -245,11 +337,25 @@ class JobsRunnerAsyncio:
             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
+            import sys
+
+            print(msg, file=sys.stderr)
+            from edsl.config import CONFIG
+
+            if CONFIG.get("EDSL_OPEN_EXCEPTION_REPORT_URL") == "True":
+                open_in_browser = True
+            elif CONFIG.get("EDSL_OPEN_EXCEPTION_REPORT_URL") == "False":
+                open_in_browser = False
+            else:
+                raise Exception(
+                    "EDSL_OPEN_EXCEPTION_REPORT_URL", "must be either True or False"
+                )
+
+            # print("open_in_browser", open_in_browser)
+
             filepath = results.task_history.html(
                 cta="Open report to see details.",
-                open_in_browser=True,
+                open_in_browser=open_in_browser,
                 return_link=True,
             )
 
@@ -279,6 +385,7 @@ class JobsRunnerAsyncio:
         progress_bar: bool = False,
         sidecar_model: Optional[LanguageModel] = None,
         jobs_runner_status: Optional[Type[JobsRunnerStatusBase]] = None,
+        job_uuid: Optional[UUID] = None,
         print_exceptions: bool = True,
         raise_validation_errors: bool = False,
     ) -> "Coroutine":
@@ -297,13 +404,11 @@ class JobsRunnerAsyncio:
 
         if jobs_runner_status is not None:
             self.jobs_runner_status = jobs_runner_status(
-                self, n=n, endpoint_url=endpoint_url
+                self, n=n, endpoint_url=endpoint_url, job_uuid=job_uuid
             )
         else:
             self.jobs_runner_status = JobsRunnerStatus(
-                self,
-                n=n,
-                endpoint_url=endpoint_url,
+                self, n=n, endpoint_url=endpoint_url, job_uuid=job_uuid
             )
 
         stop_event = threading.Event()

edsl/jobs/runners/JobsRunnerStatus.py

@@ -239,7 +239,6 @@ class JobsRunnerStatusBase(ABC):
         return stat_definitions[stat_name]()
 
     def update_progress(self, stop_event):
-
         while not stop_event.is_set():
             self.send_status_update()
             time.sleep(self.refresh_rate)
@@ -248,7 +247,6 @@ class JobsRunnerStatusBase(ABC):
 
 
 class JobsRunnerStatus(JobsRunnerStatusBase):
-
     @property
     def create_url(self) -> str:
         return f"{self.base_url}/api/v0/local-job"

edsl/jobs/tasks/TaskHistory.py

@@ -8,7 +8,12 @@ from edsl.jobs.tasks.task_status_enum import TaskStatus
 
 
 class TaskHistory:
-    def __init__(self, interviews: List["Interview"], include_traceback: bool = False):
+    def __init__(
+        self,
+        interviews: List["Interview"],
+        include_traceback: bool = False,
+        max_interviews: int = 10,
+    ):
         """
         The structure of a TaskHistory exception
 
@@ -22,6 +27,7 @@ class TaskHistory:
         self.include_traceback = include_traceback
 
         self._interviews = {index: i for index, i in enumerate(self.total_interviews)}
+        self.max_interviews = max_interviews
 
     @classmethod
     def example(cls):
@@ -75,13 +81,6 @@ class TaskHistory:
 
     def to_dict(self, add_edsl_version=True):
         """Return the TaskHistory as a dictionary."""
-        # return {
-        #     "exceptions": [
-        #         e.to_dict(include_traceback=self.include_traceback)
-        #         for e in self.exceptions
-        #     ],
-        #     "indices": self.indices,
-        # }
         d = {
             "interviews": [
                 i.to_dict(add_edsl_version=add_edsl_version)
@@ -124,10 +123,11 @@ class TaskHistory:
 
     def _repr_html_(self):
         """Return an HTML representation of the TaskHistory."""
-        from edsl.utilities.utilities import data_to_html
+        d = self.to_dict(add_edsl_version=False)
+        data = [[k, v] for k, v in d.items()]
+        from tabulate import tabulate
 
-        newdata = self.to_dict()["exceptions"]
-        return data_to_html(newdata, replace_new_lines=True)
+        return tabulate(data, headers=["keys", "values"], tablefmt="html")
 
     def show_exceptions(self, tracebacks=False):
         """Print the exceptions."""
@@ -257,8 +257,6 @@ class TaskHistory:
             for question_name, exceptions in interview.exceptions.items():
                 for exception in exceptions:
                     exception_type = exception.exception.__class__.__name__
-                    # exception_type = exception["exception"]
-                    # breakpoint()
                     if exception_type in exceptions_by_type:
                         exceptions_by_type[exception_type] += 1
                     else:
@@ -345,9 +343,9 @@ class TaskHistory:
 
         env = Environment(loader=TemplateLoader("edsl", "templates/error_reporting"))
 
-        # Load and render a template
+        # Get current memory usage at this point
+
         template = env.get_template("base.html")
-        # rendered_template = template.render(your_data=your_data)
 
         # Render the template with data
         output = template.render(
@@ -361,6 +359,7 @@ class TaskHistory:
             exceptions_by_model=self.exceptions_by_model,
             exceptions_by_service=self.exceptions_by_service,
             models_used=models_used,
+            max_interviews=self.max_interviews,
         )
         return output
 
@@ -370,7 +369,7 @@ class TaskHistory:
         return_link=False,
         css=None,
         cta="Open Report in New Tab",
-        open_in_browser=True,
+        open_in_browser=False,
     ):
         """Return an HTML report."""
 

edsl/language_models/LanguageModel.py

@@ -41,17 +41,19 @@ from edsl.data_transfer_models import (
     AgentResponseDict,
 )
 
+if TYPE_CHECKING:
+    from edsl.data.Cache import Cache
+    from edsl.scenarios.FileStore import FileStore
+    from edsl.questions.QuestionBase import QuestionBase
 
 from edsl.config import CONFIG
 from edsl.utilities.decorators import sync_wrapper, jupyter_nb_handler
-from edsl.utilities.decorators import add_edsl_version, remove_edsl_version
-from edsl.language_models.repair import repair
-from edsl.enums import InferenceServiceType
-from edsl.Base import RichPrintingMixin, PersistenceMixin
-from edsl.language_models.RegisterLanguageModelsMeta import RegisterLanguageModelsMeta
-from edsl.exceptions.language_models import LanguageModelBadResponseError
+from edsl.utilities.decorators import remove_edsl_version
 
+from edsl.Base import PersistenceMixin
+from edsl.language_models.RegisterLanguageModelsMeta import RegisterLanguageModelsMeta
 from edsl.language_models.KeyLookup import KeyLookup
+from edsl.exceptions.language_models import LanguageModelBadResponseError
 
 TIMEOUT = float(CONFIG.get("EDSL_API_TIMEOUT"))
 
@@ -116,29 +118,11 @@ def handle_key_error(func):
 
 
 class LanguageModel(
-    RichPrintingMixin, PersistenceMixin, ABC, metaclass=RegisterLanguageModelsMeta
+    PersistenceMixin,
+    ABC,
+    metaclass=RegisterLanguageModelsMeta,
 ):
-    """ABC for LLM subclasses.
-
-    TODO:
-
-    1) Need better, more descriptive names for functions
-
-    get_model_response_no_cache  (currently called async_execute_model_call)
-
-    get_model_response (currently called async_get_raw_response; uses cache & adds tracking info)
-      Calls:
-        - async_execute_model_call
-        - _updated_model_response_with_tracking
-
-    get_answer (currently called async_get_response)
-        This parses out the answer block and does some error-handling.
-        Calls:
-            - async_get_raw_response
-            - parse_response
-
-
-    """
+    """ABC for Language Models."""
 
     _model_ = None
     key_sequence = (
@@ -196,7 +180,7 @@ class LanguageModel(
         system_prompt = "You are a helpful agent pretending to be a human."
         return self.execute_model_call(user_prompt, system_prompt)
 
-    def set_key_lookup(self, key_lookup: KeyLookup):
+    def set_key_lookup(self, key_lookup: KeyLookup) -> None:
         del self._api_token
         self.key_lookup = key_lookup
 
@@ -211,10 +195,14 @@ class LanguageModel(
     def __getitem__(self, key):
         return getattr(self, key)
 
-    def _repr_html_(self):
-        from edsl.utilities.utilities import data_to_html
+    def _repr_html_(self) -> str:
+        d = {"model": self.model}
+        d.update(self.parameters)
+        data = [[k, v] for k, v in d.items()]
+        from tabulate import tabulate
 
-        return data_to_html(self.to_dict())
+        table = str(tabulate(data, headers=["keys", "values"], tablefmt="html"))
+        return f"<pre>{table}</pre>"
 
     def hello(self, verbose=False):
         """Runs a simple test to check if the model is working."""
@@ -235,7 +223,6 @@ class LanguageModel(
         This method is used to check if the model has a valid API key.
         """
         from edsl.enums import service_to_api_keyname
-        import os
 
         if self._model_ == "test":
             return True
@@ -248,9 +235,9 @@ class LanguageModel(
         """Allow the model to be used as a key in a dictionary."""
         from edsl.utilities.utilities import dict_hash
 
-        return dict_hash(self.to_dict())
+        return dict_hash(self.to_dict(add_edsl_version=False))
 
-    def __eq__(self, other):
+    def __eq__(self, other) -> bool:
         """Check is two models are the same.
 
         >>> m1 = LanguageModel.example()
@@ -278,15 +265,11 @@ class LanguageModel(
     @property
     def RPM(self):
         """Model's requests-per-minute limit."""
-        # self._set_rate_limits()
-        # return self._safety_factor * self.__rate_limits["rpm"]
         return self._rpm
 
     @property
     def TPM(self):
         """Model's tokens-per-minute limit."""
-        # self._set_rate_limits()
-        # return self._safety_factor * self.__rate_limits["tpm"]
         return self._tpm
 
     @property
@@ -314,8 +297,6 @@ class LanguageModel(
         >>> LanguageModel._overide_default_parameters(passed_parameter_dict={"temperature": 0.5}, default_parameter_dict={"temperature":0.9, "max_tokens": 1000})
         {'temperature': 0.5, 'max_tokens': 1000}
         """
-        # parameters = dict({})
-
         # this is the case when data is loaded from a dict after serialization
         if "parameters" in passed_parameter_dict:
             passed_parameter_dict = passed_parameter_dict["parameters"]
@@ -429,9 +410,10 @@ class LanguageModel(
         self,
         user_prompt: str,
         system_prompt: str,
-        cache: "Cache",
+        cache: Cache,
         iteration: int = 0,
-        files_list=None,
+        files_list: Optional[List[FileStore]] = None,
+        invigilator=None,
     ) -> ModelResponse:
         """Handle caching of responses.
 
@@ -455,7 +437,6 @@ class LanguageModel(
 
         if files_list:
             files_hash = "+".join([str(hash(file)) for file in files_list])
-            # print(f"Files hash: {files_hash}")
             user_prompt_with_hashes = user_prompt + f" {files_hash}"
         else:
             user_prompt_with_hashes = user_prompt
@@ -481,9 +462,7 @@ class LanguageModel(
                 "user_prompt": user_prompt,
                 "system_prompt": system_prompt,
                 "files_list": files_list,
-                # **({"encoded_image": encoded_image} if encoded_image else {}),
             }
-            # response = await f(**params)
             response = await asyncio.wait_for(f(**params), timeout=TIMEOUT)
             new_cache_key = cache.store(
                 **cache_call_params, response=response
@@ -504,11 +483,9 @@ class LanguageModel(
         _async_get_intended_model_call_outcome
     )
 
-    # get_raw_response = sync_wrapper(async_get_raw_response)
-
     def simple_ask(
         self,
-        question: "QuestionBase",
+        question: QuestionBase,
         system_prompt="You are a helpful agent pretending to be a human.",
         top_logprobs=2,
     ):
@@ -523,9 +500,10 @@ class LanguageModel(
         self,
         user_prompt: str,
         system_prompt: str,
-        cache: "Cache",
+        cache: Cache,
         iteration: int = 1,
-        files_list: Optional[List["File"]] = None,
+        files_list: Optional[List[FileStore]] = None,
+        **kwargs,
     ) -> dict:
         """Get response, parse, and return as string.
 
@@ -543,6 +521,9 @@ class LanguageModel(
             "cache": cache,
             "files_list": files_list,
         }
+        if "invigilator" in kwargs:
+            params.update({"invigilator": kwargs["invigilator"]})
+
         model_inputs = ModelInputs(user_prompt=user_prompt, system_prompt=system_prompt)
         model_outputs = await self._async_get_intended_model_call_outcome(**params)
         edsl_dict = self.parse_response(model_outputs.response)
@@ -553,8 +534,6 @@ class LanguageModel(
         )
         return agent_response_dict
 
-        # return await self._async_prepare_response(model_call_outcome, cache=cache)
-
     get_response = sync_wrapper(async_get_response)
 
     def cost(self, raw_response: dict[str, Any]) -> Union[float, str]:
@@ -604,10 +583,7 @@ class LanguageModel(
 
         return input_cost + output_cost
 
-    #######################
-    # SERIALIZATION METHODS
-    #######################
-    def to_dict(self, add_edsl_version=True) -> dict[str, Any]:
+    def to_dict(self, add_edsl_version: bool = True) -> dict[str, Any]:
         """Convert instance to a dictionary
 
         >>> m = LanguageModel.example()
@@ -629,18 +605,8 @@ class LanguageModel(
         from edsl.language_models.registry import get_model_class
 
         model_class = get_model_class(data["model"])
-        # data["use_cache"] = True
         return model_class(**data)
 
-    #######################
-    # DUNDER METHODS
-    #######################
-    def print(self):
-        from rich import print_json
-        import json
-
-        print_json(json.dumps(self.to_dict()))
-
     def __repr__(self) -> str:
         """Return a string representation of the object."""
         param_string = ", ".join(
@@ -654,33 +620,21 @@ class LanguageModel(
 
     def __add__(self, other_model: Type[LanguageModel]) -> Type[LanguageModel]:
         """Combine two models into a single model (other_model takes precedence over self)."""
-        print(
+        import warnings
+
+        warnings.warn(
             f"""Warning: one model is replacing another. If you want to run both models, use a single `by` e.g., 
               by(m1, m2, m3) not by(m1).by(m2).by(m3)."""
         )
         return other_model or self
 
-    def rich_print(self):
-        """Display an object as a table."""
-        from rich.table import Table
-
-        table = Table(title="Language Model")
-        table.add_column("Attribute", style="bold")
-        table.add_column("Value")
-
-        to_display = self.__dict__.copy()
-        for attr_name, attr_value in to_display.items():
-            table.add_row(attr_name, repr(attr_value))
-
-        return table
-
     @classmethod
     def example(
         cls,
         test_model: bool = False,
         canned_response: str = "Hello world",
         throw_exception: bool = False,
-    ):
+    ) -> LanguageModel:
         """Return a default instance of the class.
 
         >>> from edsl.language_models import LanguageModel
@@ -691,11 +645,17 @@ class LanguageModel(
         >>> q = QuestionFreeText(question_text = "What is your name?", question_name = 'example')
         >>> q.by(m).run(cache = False, disable_remote_cache = True, disable_remote_inference = True).select('example').first()
         'WOWZA!'
+        >>> m = LanguageModel.example(test_model = True, canned_response = "WOWZA!", throw_exception = True)
+        >>> r = q.by(m).run(cache = False, disable_remote_cache = True, disable_remote_inference = True, print_exceptions = True)
+        Exception report saved to ...
+        Also see: ...
         """
         from edsl import Model
 
         if test_model:
-            m = Model("test", canned_response=canned_response)
+            m = Model(
+                "test", canned_response=canned_response, throw_exception=throw_exception
+            )
             return m
         else:
             return Model(skip_api_key_check=True)