edsl 0.1.29.dev6__py3-none-any.whl → 0.1.30.dev1__py3-none-any.whl

edsl/config.py

@@ -1,11 +1,11 @@
 """This module provides a Config class that loads environment variables from a .env file and sets them as class attributes."""
 
 import os
-from dotenv import load_dotenv, find_dotenv
 from edsl.exceptions import (
     InvalidEnvironmentVariableError,
     MissingEnvironmentVariableError,
 )
+from dotenv import load_dotenv, find_dotenv
 
 # valid values for EDSL_RUN_MODE
 EDSL_RUN_MODES = ["development", "development-testrun", "production"]
@@ -96,6 +96,7 @@ class Config:
         Loads the .env
         - Overrides existing env vars unless EDSL_RUN_MODE=="development-testrun"
         """
+
         override = True
         if self.EDSL_RUN_MODE == "development-testrun":
             override = False

edsl/coop/utils.py

@@ -10,7 +10,7 @@ from edsl import (
     Study,
 )
 from edsl.questions import QuestionBase
-from typing import Literal, Type, Union
+from typing import Literal, Optional, Type, Union
 
 EDSLObject = Union[
     Agent,
@@ -94,3 +94,30 @@ class ObjectRegistry:
         if EDSL_object is None:
             raise ValueError(f"EDSL class not found for {object_type=}")
         return EDSL_object
+
+    @classmethod
+    def get_registry(
+        cls,
+        subclass_registry: Optional[dict] = None,
+        exclude_classes: Optional[list] = None,
+    ) -> dict:
+        """
+        Return the registry of objects.
+
+        Exclude objects that are already registered in subclass_registry.
+        This allows the user to isolate Coop-only objects.
+
+        Also exclude objects if their class name is in the exclude_classes list.
+        """
+
+        if subclass_registry is None:
+            subclass_registry = {}
+        if exclude_classes is None:
+            exclude_classes = []
+
+        return {
+            class_name: o["edsl_class"]
+            for o in cls.objects
+            if (class_name := o["edsl_class"].__name__) not in subclass_registry
+            and class_name not in exclude_classes
+        }

edsl/data/Cache.py

@@ -7,13 +7,13 @@ import json
 import os
 import warnings
 from typing import Optional, Union
-
+import time
 from edsl.config import CONFIG
 from edsl.data.CacheEntry import CacheEntry
-from edsl.data.SQLiteDict import SQLiteDict
+
+# from edsl.data.SQLiteDict import SQLiteDict
 from edsl.Base import Base
 from edsl.utilities.utilities import dict_hash
-
 from edsl.utilities.decorators import (
     add_edsl_version,
     remove_edsl_version,
@@ -38,7 +38,7 @@ class Cache(Base):
         self,
         *,
         filename: Optional[str] = None,
-        data: Optional[Union[SQLiteDict, dict]] = None,
+        data: Optional[Union["SQLiteDict", dict]] = None,
         immediate_write: bool = True,
         method=None,
     ):
@@ -104,6 +104,8 @@ class Cache(Base):
 
     def _perform_checks(self):
         """Perform checks on the cache."""
+        from edsl.data.CacheEntry import CacheEntry
+
         if any(not isinstance(value, CacheEntry) for value in self.data.values()):
             raise Exception("Not all values are CacheEntry instances")
         if self.method is not None:
@@ -138,6 +140,8 @@ class Cache(Base):
 
 
         """
+        from edsl.data.CacheEntry import CacheEntry
+
         key = CacheEntry.gen_key(
             model=model,
             parameters=parameters,
@@ -171,6 +175,7 @@ class Cache(Base):
         * If `immediate_write` is True , the key-value pair is added to `self.data`
         * If `immediate_write` is False, the key-value pair is added to `self.new_entries_to_write_later`
         """
+
         entry = CacheEntry(
             model=model,
             parameters=parameters,
@@ -188,13 +193,14 @@ class Cache(Base):
         return key
 
     def add_from_dict(
-        self, new_data: dict[str, CacheEntry], write_now: Optional[bool] = True
+        self, new_data: dict[str, "CacheEntry"], write_now: Optional[bool] = True
     ) -> None:
         """
         Add entries to the cache from a dictionary.
 
         :param write_now: Whether to write to the cache immediately (similar to `immediate_write`).
         """
+
         for key, value in new_data.items():
             if key in self.data:
                 if value != self.data[key]:
@@ -231,6 +237,8 @@ class Cache(Base):
 
         :param write_now: Whether to write to the cache immediately (similar to `immediate_write`).
         """
+        from edsl.data.SQLiteDict import SQLiteDict
+
         db = SQLiteDict(db_path)
         new_data = {}
         for key, value in db.items():
@@ -242,6 +250,8 @@ class Cache(Base):
         """
         Construct a Cache from a SQLite database.
         """
+        from edsl.data.SQLiteDict import SQLiteDict
+
         return cls(data=SQLiteDict(db_path))
 
     @classmethod
@@ -268,6 +278,8 @@ class Cache(Base):
         * If `db_path` is provided, the cache will be stored in an SQLite database.
         """
         # if a file doesn't exist at jsonfile, throw an error
+        from edsl.data.SQLiteDict import SQLiteDict
+
         if not os.path.exists(jsonlfile):
             raise FileNotFoundError(f"File {jsonlfile} not found")
 
@@ -286,6 +298,8 @@ class Cache(Base):
         """
         ## TODO: Check to make sure not over-writing (?)
         ## Should be added to SQLiteDict constructor (?)
+        from edsl.data.SQLiteDict import SQLiteDict
+
         new_data = SQLiteDict(db_path)
         for key, value in self.data.items():
             new_data[key] = value

edsl/data/SQLiteDict.py

@@ -1,9 +1,7 @@
 from __future__ import annotations
 import json
-from sqlalchemy import create_engine
-from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.orm import sessionmaker
 from typing import Any, Generator, Optional, Union
+
 from edsl.config import CONFIG
 from edsl.data.CacheEntry import CacheEntry
 from edsl.data.orm import Base, Data
@@ -25,10 +23,16 @@ class SQLiteDict:
         >>> import os; os.unlink(temp_db_path)  # Clean up the temp file after the test
 
         """
+        from sqlalchemy.exc import SQLAlchemyError
+        from sqlalchemy.orm import sessionmaker
+        from sqlalchemy import create_engine
+
         self.db_path = db_path or CONFIG.get("EDSL_DATABASE_PATH")
         if not self.db_path.startswith("sqlite:///"):
             self.db_path = f"sqlite:///{self.db_path}"
         try:
+            from edsl.data.orm import Base, Data
+
             self.engine = create_engine(self.db_path, echo=False, future=True)
             Base.metadata.create_all(self.engine)
             self.Session = sessionmaker(bind=self.engine)
@@ -55,6 +59,8 @@ class SQLiteDict:
         if not isinstance(value, CacheEntry):
             raise ValueError(f"Value must be a CacheEntry object (got {type(value)}).")
         with self.Session() as db:
+            from edsl.data.orm import Base, Data
+
             db.merge(Data(key=key, value=json.dumps(value.to_dict())))
             db.commit()
 
@@ -69,6 +75,8 @@ class SQLiteDict:
         True
         """
         with self.Session() as db:
+            from edsl.data.orm import Base, Data
+
             value = db.query(Data).filter_by(key=key).first()
             if not value:
                 raise KeyError(f"Key '{key}' not found.")

edsl/jobs/Answers.py

@@ -8,7 +8,15 @@ class Answers(UserDict):
     """Helper class to hold the answers to a survey."""
 
     def add_answer(self, response, question) -> None:
-        """Add a response to the answers dictionary."""
+        """Add a response to the answers dictionary.
+
+        >>> from edsl import QuestionFreeText
+        >>> q = QuestionFreeText.example()
+        >>> answers = Answers()
+        >>> answers.add_answer({"answer": "yes"}, q)
+        >>> answers[q.question_name]
+        'yes'
+        """
         answer = response.get("answer")
         comment = response.pop("comment", None)
         # record the answer
@@ -42,3 +50,9 @@ class Answers(UserDict):
             table.add_row(attr_name, repr(attr_value))
 
         return table
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

edsl/jobs/Jobs.py

@@ -1,30 +1,15 @@
 # """The Jobs class is a collection of agents, scenarios and models and one survey."""
 from __future__ import annotations
-import os
 import warnings
 from itertools import product
 from typing import Optional, Union, Sequence, Generator
-from edsl import Model
-from edsl.agents import Agent
-from edsl.agents.AgentList import AgentList
+
 from edsl.Base import Base
-from edsl.data.Cache import Cache
-from edsl.data.CacheHandler import CacheHandler
-from edsl.results.Dataset import Dataset
 
-from edsl.exceptions.jobs import MissingRemoteInferenceError
 from edsl.exceptions import MissingAPIKeyError
 from edsl.jobs.buckets.BucketCollection import BucketCollection
 from edsl.jobs.interviews.Interview import Interview
-from edsl.language_models import LanguageModel
-from edsl.results import Results
-from edsl.scenarios import Scenario
-from edsl import ScenarioList
-from edsl.surveys import Survey
 from edsl.jobs.runners.JobsRunnerAsyncio import JobsRunnerAsyncio
-
-from edsl.language_models.ModelList import ModelList
-
 from edsl.utilities.decorators import add_edsl_version, remove_edsl_version
 
 
@@ -37,10 +22,10 @@ class Jobs(Base):
 
     def __init__(
         self,
-        survey: Survey,
-        agents: Optional[list[Agent]] = None,
-        models: Optional[list[LanguageModel]] = None,
-        scenarios: Optional[list[Scenario]] = None,
+        survey: "Survey",
+        agents: Optional[list["Agent"]] = None,
+        models: Optional[list["LanguageModel"]] = None,
+        scenarios: Optional[list["Scenario"]] = None,
     ):
         """Initialize a Jobs instance.
 
@@ -50,8 +35,8 @@ class Jobs(Base):
         :param scenarios: a list of scenarios
         """
         self.survey = survey
-        self.agents: AgentList = agents
-        self.scenarios: ScenarioList = scenarios
+        self.agents: "AgentList" = agents
+        self.scenarios: "ScenarioList" = scenarios
         self.models = models
 
         self.__bucket_collection = None
@@ -62,6 +47,8 @@ class Jobs(Base):
 
     @models.setter
     def models(self, value):
+        from edsl import ModelList
+
         if value:
             if not isinstance(value, ModelList):
                 self._models = ModelList(value)
@@ -76,6 +63,8 @@ class Jobs(Base):
 
     @agents.setter
     def agents(self, value):
+        from edsl import AgentList
+
         if value:
             if not isinstance(value, AgentList):
                 self._agents = AgentList(value)
@@ -90,6 +79,8 @@ class Jobs(Base):
 
     @scenarios.setter
     def scenarios(self, value):
+        from edsl import ScenarioList
+
         if value:
             if not isinstance(value, ScenarioList):
                 self._scenarios = ScenarioList(value)
@@ -101,10 +92,10 @@ class Jobs(Base):
     def by(
         self,
         *args: Union[
-            Agent,
-            Scenario,
-            LanguageModel,
-            Sequence[Union[Agent, Scenario, LanguageModel]],
+            "Agent",
+            "Scenario",
+            "LanguageModel",
+            Sequence[Union["Agent", "Scenario", "LanguageModel"]],
         ],
     ) -> Jobs:
         """
@@ -144,7 +135,7 @@ class Jobs(Base):
         setattr(self, objects_key, new_objects)  # update the job
         return self
 
-    def prompts(self) -> Dataset:
+    def prompts(self) -> "Dataset":
         """Return a Dataset of prompts that will be used.
 
 
@@ -160,6 +151,7 @@ class Jobs(Base):
         user_prompts = []
         system_prompts = []
         scenario_indices = []
+        from edsl.results.Dataset import Dataset
 
         for interview_index, interview in enumerate(interviews):
             invigilators = list(interview._build_invigilators(debug=False))
@@ -182,7 +174,10 @@ class Jobs(Base):
 
     @staticmethod
     def _get_container_class(object):
-        from edsl import AgentList
+        from edsl.agents.AgentList import AgentList
+        from edsl.agents.Agent import Agent
+        from edsl.scenarios.Scenario import Scenario
+        from edsl.scenarios.ScenarioList import ScenarioList
 
         if isinstance(object, Agent):
             return AgentList
@@ -218,6 +213,10 @@ class Jobs(Base):
     def _get_current_objects_of_this_type(
         self, object: Union[Agent, Scenario, LanguageModel]
     ) -> tuple[list, str]:
+        from edsl.agents.Agent import Agent
+        from edsl.scenarios.Scenario import Scenario
+        from edsl.language_models.LanguageModel import LanguageModel
+
         """Return the current objects of the same type as the first argument.
 
         >>> from edsl.jobs import Jobs
@@ -246,6 +245,9 @@ class Jobs(Base):
     @staticmethod
     def _get_empty_container_object(object):
         from edsl import AgentList
+        from edsl import Agent
+        from edsl import Scenario
+        from edsl import ScenarioList
 
         if isinstance(object, Agent):
             return AgentList([])
@@ -310,6 +312,10 @@ class Jobs(Base):
         with us filling in defaults.
         """
         # if no agents, models, or scenarios are set, set them to defaults
+        from edsl.agents.Agent import Agent
+        from edsl.language_models.registry import Model
+        from edsl.scenarios.Scenario import Scenario
+
         self.agents = self.agents or [Agent()]
         self.models = self.models or [Model()]
         self.scenarios = self.scenarios or [Scenario()]
@@ -325,6 +331,7 @@ class Jobs(Base):
         These buckets are used to track API calls and token usage.
 
         >>> from edsl.jobs import Jobs
+        >>> from edsl import Model
         >>> j = Jobs.example().by(Model(temperature = 1), Model(temperature = 0.5))
         >>> bc = j.create_bucket_collection()
         >>> bc
@@ -368,6 +375,8 @@ class Jobs(Base):
         """Check if the parameters in the survey and scenarios are consistent.
 
         >>> from edsl import QuestionFreeText
+        >>> from edsl import Survey
+        >>> from edsl import Scenario
         >>> q = QuestionFreeText(question_text = "{{poo}}", question_name = "ugly_question")
         >>> j = Jobs(survey = Survey(questions=[q]))
         >>> with warnings.catch_warnings(record=True) as w:
@@ -464,8 +473,25 @@ class Jobs(Base):
                 status="queued",
             )
             self._output("Job sent!")
-            self._output(remote_job_data)
-            return remote_job_data
+            # Create mock results object to store job data
+            results = Results(
+                survey=Survey(),
+                data=[
+                    Result(
+                        agent=Agent.example(),
+                        scenario=Scenario.example(),
+                        model=Model(),
+                        iteration=1,
+                        answer={"info": "Remote job details"},
+                    )
+                ],
+            )
+            results.add_columns_from_dict([remote_job_data])
+            if self.verbose:
+                results.select(["info", "uuid", "status", "version"]).print(
+                    format="rich"
+                )
+            return results
         else:
             if check_api_keys:
                 for model in self.models + [Model()]:
@@ -477,8 +503,12 @@ class Jobs(Base):
 
         # handle cache
         if cache is None:
+            from edsl.data.CacheHandler import CacheHandler
+
             cache = CacheHandler().get_cache()
         if cache is False:
+            from edsl.data.Cache import Cache
+
             cache = Cache()
 
         if not remote_cache:
@@ -630,6 +660,11 @@ class Jobs(Base):
     @remove_edsl_version
     def from_dict(cls, data: dict) -> Jobs:
         """Creates a Jobs instance from a dictionary."""
+        from edsl import Survey
+        from edsl.agents.Agent import Agent
+        from edsl.language_models.LanguageModel import LanguageModel
+        from edsl.scenarios.Scenario import Scenario
+
         return cls(
             survey=Survey.from_dict(data["survey"]),
             agents=[Agent.from_dict(agent) for agent in data["agents"]],
@@ -656,7 +691,8 @@ class Jobs(Base):
         """
         import random
         from edsl.questions import QuestionMultipleChoice
-        from edsl import Agent
+        from edsl.agents.Agent import Agent
+        from edsl.scenarios.Scenario import Scenario
 
         # (status, question, period)
         agent_answers = {
@@ -695,6 +731,8 @@ class Jobs(Base):
             question_options=["Good", "Great", "OK", "Terrible"],
             question_name="how_feeling_yesterday",
         )
+        from edsl import Survey, ScenarioList
+
         base_survey = Survey(questions=[q1, q2])
 
         scenario_list = ScenarioList(

edsl/jobs/buckets/ModelBuckets.py

@@ -1,4 +1,4 @@
-from edsl.jobs.buckets.TokenBucket import TokenBucket
+# from edsl.jobs.buckets.TokenBucket import TokenBucket
 
 
 class ModelBuckets:
@@ -8,7 +8,7 @@ class ModelBuckets:
     A request is one call to the service. The number of tokens required for a request depends on parameters.
     """
 
-    def __init__(self, requests_bucket: TokenBucket, tokens_bucket: TokenBucket):
+    def __init__(self, requests_bucket: "TokenBucket", tokens_bucket: "TokenBucket"):
         """Initialize the model buckets.
 
         The requests bucket captures requests per unit of time.
@@ -28,6 +28,8 @@ class ModelBuckets:
     @classmethod
     def infinity_bucket(cls, model_name: str = "not_specified") -> "ModelBuckets":
         """Create a bucket with infinite capacity and refill rate."""
+        from edsl.jobs.buckets.TokenBucket import TokenBucket
+
         return cls(
             requests_bucket=TokenBucket(
                 bucket_name=model_name,

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

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"))
 
@@ -150,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)
 

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.
@@ -169,6 +157,8 @@ class JobsRunnerAsyncio(JobsRunnerStatusMixin):
                 question_name + "_raw_model_response"
             ] = result["raw_model_response"]
 
+        from edsl.results.Result import Result
+
         result = Result(
             agent=interview.agent,
             scenario=interview.scenario,
@@ -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