edsl 0.1.28__py3-none-any.whl → 0.1.29__py3-none-any.whl

edsl/jobs/buckets/TokenBucket.py

@@ -1,8 +1,6 @@
 from typing import Union, List, Any
 import asyncio
 import time
-from collections import UserDict
-from matplotlib import pyplot as plt
 
 
 class TokenBucket:
@@ -114,6 +112,7 @@ class TokenBucket:
         times, tokens = zip(*self.get_log())
         start_time = times[0]
         times = [t - start_time for t in times]  # Normalize time to start from 0
+        from matplotlib import pyplot as plt
 
         plt.figure(figsize=(10, 6))
         plt.plot(times, tokens, label="Tokens Available")

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

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,15 +152,17 @@ 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.
         """
+        from edsl.data_transfer_models import AgentResponseDict
+
         try:
             invigilator = self._get_invigilator(question, debug=debug)
 
@@ -253,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/runners/JobsRunnerAsyncio.py

@@ -1,29 +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
-
-import time
+from edsl.utilities.decorators import jupyter_nb_handler
 
 
 class JobsRunnerAsyncio(JobsRunnerStatusMixin):
@@ -42,13 +30,13 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
 
     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.
@@ -155,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,
@@ -197,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()
@@ -204,6 +196,8 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
         self.cache = cache
         self.sidecar_model = sidecar_model
 
+        from edsl.results.Results import Results
+
         if not progress_bar:
             # print("Running without progress bar")
             with cache as c:
@@ -225,6 +219,8 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
             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)

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
 

edsl/language_models/LanguageModel.py

@@ -7,26 +7,18 @@ import asyncio
 import json
 import time
 import os
-
 from typing import Coroutine, Any, Callable, Type, List, get_type_hints
-
-from abc import ABC, abstractmethod, ABCMeta
-
-from rich.table import Table
+from abc import ABC, abstractmethod
 
 from edsl.config import CONFIG
 
-from edsl.utilities.utilities import clean_json
 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.exceptions.language_models import LanguageModelAttributeTypeError
 from edsl.enums import InferenceServiceType
 from edsl.Base import RichPrintingMixin, PersistenceMixin
-from edsl.data.Cache import Cache
 from edsl.enums import service_to_api_keyname
-
-
 from edsl.exceptions import MissingAPIKeyError
 from edsl.language_models.RegisterLanguageModelsMeta import RegisterLanguageModelsMeta
 
@@ -291,7 +283,7 @@ class LanguageModel(
         self,
         user_prompt: str,
         system_prompt: str,
-        cache,
+        cache: "Cache",
         iteration: int = 0,
         encoded_image=None,
     ) -> tuple[dict, bool, str]:
@@ -490,6 +482,8 @@ class LanguageModel(
 
     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")

edsl/language_models/ModelList.py

@@ -5,7 +5,7 @@ from edsl import Model
 from edsl.language_models import LanguageModel
 from edsl.Base import Base
 from edsl.utilities.decorators import add_edsl_version, remove_edsl_version
-from edsl.utilities import is_valid_variable_name
+from edsl.utilities.utilities import is_valid_variable_name
 from edsl.utilities.utilities import dict_hash
 
 
@@ -56,11 +56,11 @@ class ModelList(Base, UserList):
             return {"models": [model._to_dict() for model in self]}
 
     @classmethod
-    def from_names(self, *args):
+    def from_names(self, *args, **kwargs):
         """A a model list from a list of names"""
         if len(args) == 1 and isinstance(args[0], list):
             args = args[0]
-        return ModelList([Model(model_name) for model_name in args])
+        return ModelList([Model(model_name, **kwargs) for model_name in args])
 
     @add_edsl_version
     def to_dict(self):

edsl/language_models/repair.py

@@ -1,18 +1,13 @@
 import json
 import asyncio
 import warnings
-from rich import print
-from rich.console import Console
-from rich.syntax import Syntax
-
-from edsl.utilities.utilities import clean_json
-
-from edsl.utilities.repair_functions import extract_json_from_string
 
 
 async def async_repair(
     bad_json, error_message="", user_prompt=None, system_prompt=None, cache=None
 ):
+    from edsl.utilities.utilities import clean_json
+
     s = clean_json(bad_json)
 
     try:
@@ -27,6 +22,8 @@ async def async_repair(
         return valid_dict, success
 
     try:
+        from edsl.utilities.repair_functions import extract_json_from_string
+
         valid_dict = extract_json_from_string(s)
         success = True
     except ValueError:
@@ -98,6 +95,10 @@ async def async_repair(
     except json.JSONDecodeError:
         valid_dict = {}
         success = False
+        from rich import print
+        from rich.console import Console
+        from rich.syntax import Syntax
+
         console = Console()
         error_message = (
             f"All repairs. failed. LLM Model given [red]{str(bad_json)}[/red]"

edsl/notebooks/Notebook.py

@@ -1,10 +1,9 @@
 """A Notebook is a utility class that allows you to easily share/pull ipynbs from Coop."""
 
 import json
-import nbformat
-from nbconvert import HTMLExporter
 from typing import Dict, List, Optional
-from rich.table import Table
+
+
 from edsl.Base import Base
 from edsl.utilities.decorators import (
     add_edsl_version,
@@ -34,6 +33,8 @@ class Notebook(Base):
         If no path is provided, assume this code is run in a notebook and try to load the current notebook from file.
         :param name: A name for the Notebook.
         """
+        import nbformat
+
         # Load current notebook path as fallback (VS Code only)
         path = path or globals().get("__vsc_ipynb_file__")
         if data is not None:
@@ -56,6 +57,37 @@ class Notebook(Base):
 
         self.name = name or self.default_name
 
+    @classmethod
+    def from_script(cls, path: str, name: Optional[str] = None) -> "Notebook":
+        # Read the script file
+        with open(path, "r") as script_file:
+            script_content = script_file.read()
+
+        # Create a new Jupyter notebook
+        nb = nbformat.v4.new_notebook()
+
+        # Add the script content to the first cell
+        first_cell = nbformat.v4.new_code_cell(script_content)
+        nb.cells.append(first_cell)
+
+        # Create a Notebook instance with the notebook data
+        notebook_instance = cls(nb)
+
+        return notebook_instance
+
+    @classmethod
+    def from_current_script(cls) -> "Notebook":
+        import inspect
+        import os
+
+        # Get the path to the current file
+        current_frame = inspect.currentframe()
+        caller_frame = inspect.getouterframes(current_frame, 2)
+        current_file_path = os.path.abspath(caller_frame[1].filename)
+
+        # Use from_script to create the notebook
+        return cls.from_script(current_file_path)
+
     def __eq__(self, other):
         """
         Check if two Notebooks are equal.
@@ -103,6 +135,9 @@ class Notebook(Base):
         """
         Return HTML representation of Notebook.
         """
+        from nbconvert import HTMLExporter
+        import nbformat
+
         notebook = nbformat.from_dict(self.data)
         html_exporter = HTMLExporter(template_name="basic")
         (body, _) = html_exporter.from_notebook_node(notebook)
@@ -143,6 +178,8 @@ class Notebook(Base):
         """
         Display a Notebook as a rich table.
         """
+        from rich.table import Table
+
         table_data, column_names = self._table()
         table = Table(title=f"{self.__class__.__name__} Attributes")
         for column in column_names:

edsl/prompts/Prompt.py

@@ -1,12 +1,16 @@
-"""Class for creating prompts to be used in a survey."""
-
 from __future__ import annotations
 from typing import Optional
 from abc import ABC
 from typing import Any, List
 
 from rich.table import Table
-from jinja2 import Template, Environment, meta, TemplateSyntaxError
+from jinja2 import Template, Environment, meta, TemplateSyntaxError, Undefined
+
+
+class PreserveUndefined(Undefined):
+    def __str__(self):
+        return "{{ " + self._undefined_name + " }}"
+
 
 from edsl.exceptions.prompts import TemplateRenderError
 from edsl.prompts.prompt_config import (
@@ -35,6 +39,10 @@ class PromptBase(
 
         return data_to_html(self.to_dict())
 
+    def __len__(self):
+        """Return the length of the prompt text."""
+        return len(self.text)
+
     @classmethod
     def prompt_attributes(cls) -> List[str]:
         """Return the prompt class attributes."""
@@ -75,10 +83,10 @@ class PromptBase(
         >>> p = Prompt("Hello, {{person}}")
         >>> p2 = Prompt("How are you?")
         >>> p + p2
-        Prompt(text='Hello, {{person}}How are you?')
+        Prompt(text=\"""Hello, {{person}}How are you?\""")
 
         >>> p + "How are you?"
-        Prompt(text='Hello, {{person}}How are you?')
+        Prompt(text=\"""Hello, {{person}}How are you?\""")
         """
         if isinstance(other_prompt, str):
             return self.__class__(self.text + other_prompt)
@@ -114,7 +122,7 @@ class PromptBase(
         Example:
         >>> p = Prompt("Hello, {{person}}")
         >>> p
-        Prompt(text='Hello, {{person}}')
+        Prompt(text=\"""Hello, {{person}}\""")
         """
         return f'Prompt(text="""{self.text}""")'
 
@@ -137,7 +145,7 @@ class PromptBase(
         :param template: The template to find the variables in.
 
         """
-        env = Environment()
+        env = Environment(undefined=PreserveUndefined)
         ast = env.parse(template)
         return list(meta.find_undeclared_variables(ast))
 
@@ -186,13 +194,16 @@ class PromptBase(
 
         >>> p = Prompt("Hello, {{person}}")
         >>> p.render({"person": "John"})
-        'Hello, John'
+        Prompt(text=\"""Hello, John\""")
 
         >>> p.render({"person": "Mr. {{last_name}}", "last_name": "Horton"})
-        'Hello, Mr. Horton'
+        Prompt(text=\"""Hello, Mr. Horton\""")
 
         >>> p.render({"person": "Mr. {{last_name}}", "last_name": "Ho{{letter}}ton"}, max_nesting = 1)
-        'Hello, Mr. Horton'
+        Prompt(text=\"""Hello, Mr. Ho{{ letter }}ton\""")
+
+        >>> p.render({"person": "Mr. {{last_name}}"})
+        Prompt(text=\"""Hello, Mr. {{ last_name }}\""")
         """
         new_text = self._render(
             self.text, primary_replacement, **additional_replacements
@@ -216,12 +227,13 @@ class PromptBase(
         >>> codebook = {"age": "Age"}
         >>> p = Prompt("You are an agent named {{ name }}. {{ codebook['age']}}: {{ age }}")
         >>> p.render({"name": "John", "age": 44}, codebook=codebook)
-        'You are an agent named John. Age: 44'
+        Prompt(text=\"""You are an agent named John. Age: 44\""")
         """
+        env = Environment(undefined=PreserveUndefined)
         try:
             previous_text = None
             for _ in range(MAX_NESTING):
-                rendered_text = Template(text).render(
+                rendered_text = env.from_string(text).render(
                     primary_replacement, **additional_replacements
                 )
                 if rendered_text == previous_text:
@@ -258,7 +270,7 @@ class PromptBase(
         >>> p = Prompt("Hello, {{person}}")
         >>> p2 = Prompt.from_dict(p.to_dict())
         >>> p2
-        Prompt(text='Hello, {{person}}')
+        Prompt(text=\"""Hello, {{person}}\""")
 
         """
         class_name = data["class_name"]
@@ -290,6 +302,12 @@ class Prompt(PromptBase):
     component_type = ComponentTypes.GENERIC
 
 
+if __name__ == "__main__":
+    print("Running doctests...")
+    import doctest
+
+    doctest.testmod()
+
 from edsl.prompts.library.question_multiple_choice import *
 from edsl.prompts.library.agent_instructions import *
 from edsl.prompts.library.agent_persona import *
@@ -302,9 +320,3 @@ from edsl.prompts.library.question_numerical import *
 from edsl.prompts.library.question_rank import *
 from edsl.prompts.library.question_extract import *
 from edsl.prompts.library.question_list import *
-
-
-if __name__ == "__main__":
-    import doctest
-
-    doctest.testmod()