edsl 0.1.53__py3-none-any.whl → 0.1.55__py3-none-any.whl

edsl/language_models/language_model.py

@@ -365,6 +365,59 @@ class LanguageModel(
             self._api_token = info.api_token
         return self._api_token
 
+    def copy(self) -> "LanguageModel":
+        """Create a deep copy of this language model instance.
+
+        This method creates a completely independent copy of the language model
+        by creating a new instance with the same parameters and copying relevant attributes.
+
+        Returns:
+            LanguageModel: A new language model instance that is functionally identical to this one
+
+        Examples:
+            >>> m1 = LanguageModel.example()
+            >>> m2 = m1.copy()
+            >>> m1 == m2  # Functionally equivalent
+            True
+            >>> id(m1) == id(m2)  # But different objects
+            False
+        """
+        # Create a new instance of the same class with the same parameters
+        try:
+            # For most models, we can instantiate with the saved parameters
+            new_model = self.__class__(**self.parameters)
+
+            # Copy all important instance attributes
+            for key, value in self.__dict__.items():
+                if key not in ("_api_token",) and not key.startswith("__"):
+                    setattr(new_model, key, value)
+
+            return new_model
+        except Exception:
+            # Fallback for dynamically created classes like TestServiceLanguageModel
+            from ..inference_services import default
+
+            # If this is a test model, create a new test model instance
+            if getattr(self, "_inference_service_", "") == "test":
+                service = default.get_service("test")
+                model_class = service.create_model("test")
+                new_model = model_class(**self.parameters)
+
+                # Copy attributes
+                for key, value in self.__dict__.items():
+                    if key not in ("_api_token",) and not key.startswith("__"):
+                        setattr(new_model, key, value)
+
+                return new_model
+
+            # If we can't create the model directly, just return a simple test model
+            # This is a last resort fallback
+            from ..inference_services import get_service
+
+            service = get_service("test")
+            model_class = service.create_model("test")
+            return model_class()
+
     def __getitem__(self, key):
         """Allow dictionary-style access to model attributes.
 
@@ -509,7 +562,9 @@ class LanguageModel(
         return self.execute_model_call(user_prompt, system_prompt)
 
     @abstractmethod
-    async def async_execute_model_call(self, user_prompt: str, system_prompt: str):
+    async def async_execute_model_call(
+        self, user_prompt: str, system_prompt: str, question_name: Optional[str] = None
+    ):
         """Execute the model call asynchronously.
 
         This abstract method must be implemented by all model subclasses
@@ -518,6 +573,7 @@ class LanguageModel(
         Args:
             user_prompt: The user message or input prompt
             system_prompt: The system message or context
+            question_name: Optional name of the question being asked (primarily used for test models)
 
         Returns:
             Coroutine that resolves to the model response
@@ -529,7 +585,7 @@ class LanguageModel(
         pass
 
     async def remote_async_execute_model_call(
-        self, user_prompt: str, system_prompt: str
+        self, user_prompt: str, system_prompt: str, question_name: Optional[str] = None
     ):
         """Execute the model call remotely through the EDSL Coop service.
 
@@ -540,6 +596,7 @@ class LanguageModel(
         Args:
             user_prompt: The user message or input prompt
             system_prompt: The system message or context
+            question_name: Optional name of the question being asked (primarily used for test models)
 
         Returns:
             Coroutine that resolves to the model response from the remote service
@@ -563,6 +620,7 @@ class LanguageModel(
         Args:
             *args: Positional arguments to pass to async_execute_model_call
             **kwargs: Keyword arguments to pass to async_execute_model_call
+                     Can include question_name for test models
 
         Returns:
             The model response
@@ -674,9 +732,12 @@ class LanguageModel(
             user_prompt_with_hashes = user_prompt
 
         # Prepare parameters for cache lookup
+        cache_parameters = self.parameters.copy()
+        if self.model == "test":
+            cache_parameters.pop("canned_response", None)
         cache_call_params = {
             "model": str(self.model),
-            "parameters": self.parameters,
+            "parameters": cache_parameters,
             "system_prompt": system_prompt,
             "user_prompt": user_prompt_with_hashes,
             "iteration": iteration,
@@ -702,7 +763,9 @@ class LanguageModel(
                 "system_prompt": system_prompt,
                 "files_list": files_list,
             }
-
+            # Add question_name parameter for test models
+            if self.model == "test" and invigilator:
+                params["question_name"] = invigilator.question.question_name
             # Get timeout from configuration
             from ..config import CONFIG
 
@@ -710,7 +773,6 @@ class LanguageModel(
 
             # Execute the model call with timeout
             response = await asyncio.wait_for(f(**params), timeout=TIMEOUT)
-
             # Store the response in the cache
             new_cache_key = cache.store(
                 **cache_call_params, response=response, service=self._inference_service_
@@ -801,7 +863,6 @@ class LanguageModel(
 
         # Create structured input record
         model_inputs = ModelInputs(user_prompt=user_prompt, system_prompt=system_prompt)
-
         # Get model response (using cache if available)
         model_outputs: ModelResponse = (
             await self._async_get_intended_model_call_outcome(**params)
@@ -839,7 +900,7 @@ class LanguageModel(
         # Use the price manager to calculate the actual cost
         from .price_manager import PriceManager
 
-        price_manager = PriceManager()
+        price_manager = PriceManager.get_instance()
 
         return price_manager.calculate_cost(
             inference_service=self._inference_service_,
@@ -868,9 +929,15 @@ class LanguageModel(
             {'model': '...', 'parameters': {'temperature': ..., 'max_tokens': ..., 'top_p': ..., 'frequency_penalty': ..., 'presence_penalty': ..., 'logprobs': False, 'top_logprobs': ...}, 'inference_service': 'openai', 'edsl_version': '...', 'edsl_class_name': 'LanguageModel'}
         """
         # Build the base dictionary with essential model information
+        parameters = self.parameters.copy()
+
+        # For test models, ensure canned_response is included in serialization
+        if self.model == "test" and hasattr(self, "canned_response"):
+            parameters["canned_response"] = self.canned_response
+
         d = {
             "model": self.model,
-            "parameters": self.parameters,
+            "parameters": parameters,
             "inference_service": self._inference_service_,
         }
 
@@ -908,7 +975,25 @@ class LanguageModel(
             data["model"], service_name=data.get("inference_service", None)
         )
 
-        # Create and return a new instance
+        # Handle canned_response in parameters for test models
+        if (
+            data["model"] == "test"
+            and "parameters" in data
+            and "canned_response" in data["parameters"]
+        ):
+            # Extract canned_response from parameters to set as a direct attribute
+            canned_response = data["parameters"]["canned_response"]
+            params_copy = data.copy()
+
+            # Direct attribute will be set during initialization
+            # Add it as a top-level parameter for model initialization
+            if isinstance(params_copy, dict) and "parameters" in params_copy:
+                params_copy["canned_response"] = canned_response
+
+            # Create the instance with canned_response as a direct parameter
+            return model_class(**params_copy)
+
+        # For non-test models or test models without canned_response
         return model_class(**data)
 
     def __repr__(self) -> str:
@@ -994,8 +1079,8 @@ class LanguageModel(
 
             Create a test model that throws exceptions:
 
-            >>> 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)
+            >>> m = LanguageModel.example(test_model=True, canned_response="WOWZA!", throw_exception=True) # doctest: +SKIP
+            >>> r = q.by(m).run(cache=False, disable_remote_cache=True, disable_remote_inference=True, print_exceptions=True) # doctest: +SKIP
             Exception report saved to ...
         """
         from ..language_models import Model
@@ -1046,7 +1131,12 @@ class LanguageModel(
         ]
 
         # Define a new async_execute_model_call that only reads from cache
-        async def async_execute_model_call(self, user_prompt: str, system_prompt: str):
+        async def async_execute_model_call(
+            self,
+            user_prompt: str,
+            system_prompt: str,
+            question_name: Optional[str] = None,
+        ):
             """Only use cached responses, never making new API calls."""
             cache_call_params = {
                 "model": str(self.model),

edsl/language_models/model.py

@@ -6,8 +6,10 @@ from ..utilities import PrettyList
 from ..config import CONFIG
 from .exceptions import LanguageModelValueError
 
+# Import only what's needed initially to avoid circular imports
 from ..inference_services import (InferenceServicesCollection, 
-                                  AvailableModels, InferenceServiceABC, InferenceServiceError, default)
+                                  AvailableModels, InferenceServiceABC, InferenceServiceError)
+# The 'default' import will be imported lazily when needed
 
 from ..enums import InferenceServiceLiteral
 
@@ -20,7 +22,10 @@ def get_model_class(
     registry: Optional[InferenceServicesCollection] = None,
     service_name: Optional[InferenceServiceLiteral] = None,
 ):
-    registry = registry or default
+    if registry is None:
+        # Import default lazily only when needed
+        from ..inference_services import default as inference_default
+        registry = inference_default
     try:
         factory = registry.create_model_factory(model_name, service_name=service_name)
         return factory
@@ -54,7 +59,9 @@ class Model(metaclass=Meta):
     def get_registry(cls) -> InferenceServicesCollection:
         """Get the current registry or initialize with default if None"""
         if cls._registry is None:
-            cls._registry = default
+            # Import default lazily only when needed
+            from ..inference_services import default as inference_default
+            cls._registry = inference_default
         return cls._registry
 
     @classmethod

edsl/language_models/price_manager.py

@@ -8,20 +8,42 @@ class PriceManager:
 
     def __new__(cls):
         if cls._instance is None:
-            cls._instance = super(PriceManager, cls).__new__(cls)
+            instance = super(PriceManager, cls).__new__(cls)
+            instance._price_lookup = {}  # Instance-specific attribute
+            instance._is_initialized = False
+            cls._instance = instance  # Store the instance directly
+            return instance
         return cls._instance
 
     def __init__(self):
-        # Only initialize once, even if __init__ is called multiple times
+        """Initialize the singleton instance only once."""
         if not self._is_initialized:
             self._is_initialized = True
             self.refresh_prices()
 
-    def refresh_prices(self) -> None:
-        """
-        Fetch fresh prices from the Coop service and update the internal price lookup.
+    @classmethod
+    def get_instance(cls):
+        """Get the singleton instance, creating it if necessary."""
+        if cls._instance is None:
+            cls()  # Create the instance if it doesn't exist
+        return cls._instance
+
+    @classmethod
+    def reset(cls):
+        """Reset the singleton instance to clean up resources."""
+        cls._instance = None
+        cls._is_initialized = False
+        cls._price_lookup = {}
 
-        """
+    def __del__(self):
+        """Ensure proper cleanup when the instance is garbage collected."""
+        try:
+            self._price_lookup = {}  # Clean up resources
+        except:
+            pass  # Ignore any cleanup errors
+
+    def refresh_prices(self) -> None:
+        """Fetch fresh prices and update the internal price lookup."""
         from edsl.coop import Coop
 
         c = Coop()
@@ -31,43 +53,18 @@ class PriceManager:
             print(f"Error fetching prices: {str(e)}")
 
     def get_price(self, inference_service: str, model: str) -> Dict:
-        """
-        Get the price information for a specific service and model combination.
-        If no specific price is found, returns a fallback price.
-
-        Args:
-            inference_service (str): The name of the inference service
-            model (str): The model identifier
-
-        Returns:
-            Dict: Price information (either actual or fallback prices)
-        """
+        """Get the price information for a specific service and model."""
         key = (inference_service, model)
         return self._price_lookup.get(key) or self._get_fallback_price(
             inference_service
         )
 
     def get_all_prices(self) -> Dict[Tuple[str, str], Dict]:
-        """
-        Get the complete price lookup dictionary.
-
-        Returns:
-            Dict[Tuple[str, str], Dict]: The complete price lookup dictionary
-        """
+        """Get the complete price lookup dictionary."""
         return self._price_lookup.copy()
 
     def _get_fallback_price(self, inference_service: str) -> Dict:
-        """
-        Get fallback prices for a service.
-        - First fallback: The highest input and output prices for that service from the price lookup.
-        - Second fallback: $1.00 per million tokens (for both input and output).
-
-        Args:
-            inference_service (str): The inference service name
-
-        Returns:
-            Dict: Price information
-        """
+        """Get fallback prices for a service."""
         service_prices = [
             prices
             for (service, _), prices in self._price_lookup.items()
@@ -77,18 +74,12 @@ class PriceManager:
         input_tokens_per_usd = [
             float(p["input"]["one_usd_buys"]) for p in service_prices if "input" in p
         ]
-        if input_tokens_per_usd:
-            min_input_tokens = min(input_tokens_per_usd)
-        else:
-            min_input_tokens = 1_000_000
+        min_input_tokens = min(input_tokens_per_usd, default=1_000_000)
 
         output_tokens_per_usd = [
             float(p["output"]["one_usd_buys"]) for p in service_prices if "output" in p
         ]
-        if output_tokens_per_usd:
-            min_output_tokens = min(output_tokens_per_usd)
-        else:
-            min_output_tokens = 1_000_000
+        min_output_tokens = min(output_tokens_per_usd, default=1_000_000)
 
         return {
             "input": {"one_usd_buys": min_input_tokens},
@@ -103,19 +94,7 @@ class PriceManager:
         input_token_name: str,
         output_token_name: str,
     ) -> Union[float, str]:
-        """
-        Calculate the total cost for a model usage based on input and output tokens.
-
-        Args:
-            inference_service (str): The inference service identifier
-            model (str): The model identifier
-            usage (Dict[str, Union[str, int]]): Dictionary containing token usage information
-            input_token_name (str): Key name for input tokens in the usage dict
-            output_token_name (str): Key name for output tokens in the usage dict
-
-        Returns:
-            Union[float, str]: Total cost if calculation successful, error message string if not
-        """
+        """Calculate the total cost for a model usage."""
         relevant_prices = self.get_price(inference_service, model)
 
         # Extract token counts
@@ -137,31 +116,22 @@ class PriceManager:
             return f"Could not fetch prices from {relevant_prices} - {e}"
 
         # Calculate input cost
-        if inverse_input_price == "infinity":
-            input_cost = 0
-        else:
-            try:
-                input_cost = input_tokens / float(inverse_input_price)
-            except Exception as e:
-                return f"Could not compute input price - {e}."
+        input_cost = (
+            0
+            if inverse_input_price == "infinity"
+            else input_tokens / float(inverse_input_price)
+        )
 
         # Calculate output cost
-        if inverse_output_price == "infinity":
-            output_cost = 0
-        else:
-            try:
-                output_cost = output_tokens / float(inverse_output_price)
-            except Exception as e:
-                return f"Could not compute output price - {e}"
+        output_cost = (
+            0
+            if inverse_output_price == "infinity"
+            else output_tokens / float(inverse_output_price)
+        )
 
         return input_cost + output_cost
 
     @property
     def is_initialized(self) -> bool:
-        """
-        Check if the PriceManager has been initialized.
-
-        Returns:
-            bool: True if initialized, False otherwise
-        """
+        """Check if the PriceManager has been initialized."""
         return self._is_initialized

edsl/language_models/registry.py

@@ -18,6 +18,11 @@ class RegisterLanguageModelsMeta(ABCMeta):
 
     _registry = {}  # Initialize the registry as a dictionary
     REQUIRED_CLASS_ATTRIBUTES = ["_model_", "_parameters_", "_inference_service_"]
+    
+    @classmethod
+    def clear_registry(cls):
+        """Clear the registry to prevent memory leaks."""
+        cls._registry = {}
 
     def __init__(cls, name, bases, dct):
         """Register the class in the registry if it has a _model_ attribute."""

edsl/language_models/utilities.py

@@ -5,6 +5,7 @@ from ..surveys import Survey
 
 from .language_model import LanguageModel
 
+
 def create_survey(num_questions: int, chained: bool = True, take_scenario=False):
     from ..questions import QuestionFreeText
 
@@ -28,7 +29,6 @@ def create_survey(num_questions: int, chained: bool = True, take_scenario=False)
 def create_language_model(
     exception: Exception, fail_at_number: int, never_ending=False
 ):
-
     class LanguageModelFromUtilities(LanguageModel):
         _model_ = "test"
         _parameters_ = {"temperature": 0.5}
@@ -45,6 +45,7 @@ def create_language_model(
             user_prompt: str,
             system_prompt: str,
             files_list: Optional[List[Any]] = None,
+            question_name: Optional[str] = None,
         ) -> dict[str, Any]:
             question_number = int(
                 user_prompt.split("XX")[1]

edsl/notebooks/notebook.py

@@ -2,6 +2,10 @@
 
 from __future__ import annotations
 import json
+import subprocess
+import tempfile
+import os
+import shutil
 from typing import Dict, List, Optional, TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -17,12 +21,56 @@ class Notebook(Base):
     """
 
     default_name = "notebook"
+    
+    @staticmethod
+    def _lint_code(code: str) -> str:
+        """
+        Lint Python code using ruff.
+        
+        :param code: The Python code to lint
+        :return: The linted code
+        """
+        try:
+            # Check if ruff is installed
+            if shutil.which("ruff") is None:
+                # If ruff is not installed, return original code
+                return code
+                
+            with tempfile.NamedTemporaryFile(mode='w+', suffix='.py', delete=False) as temp_file:
+                temp_file.write(code)
+                temp_file_path = temp_file.name
+            
+            # Run ruff to format the code
+            try:
+                result = subprocess.run(
+                    ["ruff", "format", temp_file_path], 
+                    check=True, 
+                    stdout=subprocess.PIPE, 
+                    stderr=subprocess.PIPE
+                )
+                
+                # Read the formatted code
+                with open(temp_file_path, 'r') as f:
+                    linted_code = f.read()
+                
+                return linted_code
+            except subprocess.CalledProcessError:
+                # If ruff fails, return the original code
+                return code
+            except FileNotFoundError:
+                # If ruff is not installed, return the original code
+                return code
+        finally:
+            # Clean up temporary file
+            if 'temp_file_path' in locals() and os.path.exists(temp_file_path):
+                os.unlink(temp_file_path)
 
     def __init__(
         self,
         path: Optional[str] = None,
         data: Optional[Dict] = None,
         name: Optional[str] = None,
+        lint: bool = True,
     ):
         """
         Initialize a new Notebook.
@@ -32,6 +80,7 @@ class Notebook(Base):
         :param path: A filepath from which to load the notebook.
         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.
+        :param lint: Whether to lint Python code cells using ruff. Defaults to True.
         """
         import nbformat
 
@@ -54,6 +103,16 @@ class Notebook(Base):
             raise NotebookEnvironmentError(
                 "Cannot create a notebook from within itself in this development environment"
             )
+            
+        # Store the lint parameter
+        self.lint = lint
+        
+        # Apply linting to code cells if enabled
+        if self.lint and self.data and "cells" in self.data:
+            for cell in self.data["cells"]:
+                if cell.get("cell_type") == "code" and "source" in cell:
+                    # Only lint Python code cells
+                    cell["source"] = self._lint_code(cell["source"])
 
         # TODO: perhaps add sanity check function
         # 1. could check if the notebook is a valid notebook
@@ -63,7 +122,7 @@ class Notebook(Base):
         self.name = name or self.default_name
 
     @classmethod
-    def from_script(cls, path: str, name: Optional[str] = None) -> "Notebook":
+    def from_script(cls, path: str, name: Optional[str] = None, lint: bool = True) -> "Notebook":
         import nbformat
 
         # Read the script file
@@ -78,12 +137,12 @@ class Notebook(Base):
         nb.cells.append(first_cell)
 
         # Create a Notebook instance with the notebook data
-        notebook_instance = cls(nb)
+        notebook_instance = cls(data=nb, name=name, lint=lint)
 
         return notebook_instance
 
     @classmethod
-    def from_current_script(cls) -> "Notebook":
+    def from_current_script(cls, lint: bool = True) -> "Notebook":
         import inspect
         import os
 
@@ -93,7 +152,7 @@ class Notebook(Base):
         current_file_path = os.path.abspath(caller_frame[1].filename)
 
         # Use from_script to create the notebook
-        return cls.from_script(current_file_path)
+        return cls.from_script(current_file_path, lint=lint)
 
     def __eq__(self, other):
         """
@@ -114,7 +173,7 @@ class Notebook(Base):
         """
         Serialize to a dictionary.
         """
-        d = {"name": self.name, "data": self.data}
+        d = {"name": self.name, "data": self.data, "lint": self.lint}
         if add_edsl_version:
             from .. import __version__
 
@@ -124,11 +183,17 @@ class Notebook(Base):
 
     @classmethod
     @remove_edsl_version
-    def from_dict(cls, d: Dict) -> "Notebook":
+    def from_dict(cls, d: Dict, lint: bool = None) -> "Notebook":
         """
         Convert a dictionary representation of a Notebook to a Notebook object.
+        
+        :param d: Dictionary containing notebook data and name
+        :param lint: Whether to lint Python code cells. If None, uses the value from the dictionary or defaults to True.
+        :return: A new Notebook instance
         """
-        return cls(data=d["data"], name=d["name"])
+        # Use the lint parameter from the dictionary if none is provided, otherwise default to True
+        notebook_lint = lint if lint is not None else d.get("lint", True)
+        return cls(data=d["data"], name=d["name"], lint=notebook_lint)
 
     def to_file(self, path: str):
         """
@@ -205,11 +270,13 @@ class Notebook(Base):
         return table
 
     @classmethod
-    def example(cls, randomize: bool = False) -> Notebook:
+    def example(cls, randomize: bool = False, lint: bool = True) -> Notebook:
         """
         Returns an example Notebook instance.
 
         :param randomize: If True, adds a random string one of the cells' output.
+        :param lint: Whether to lint Python code cells. Defaults to True.
+        :return: An example Notebook instance
         """
         addition = "" if not randomize else str(uuid4())
         cells = [
@@ -238,7 +305,7 @@ class Notebook(Base):
             "nbformat_minor": 4,
             "cells": cells,
         }
-        return cls(data=data)
+        return cls(data=data, lint=lint)
 
     def code(self) -> List[str]:
         """
@@ -246,7 +313,7 @@ class Notebook(Base):
         """
         lines = []
         lines.append("from edsl import Notebook")  # Keep as absolute for code generation
-        lines.append(f'nb = Notebook(data={self.data}, name="""{self.name}""")')
+        lines.append(f'nb = Notebook(data={self.data}, name="""{self.name}""", lint={self.lint})')
         return lines
 
     def to_latex(self, filename: str):