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

edsl/jobs/remote_inference.py

@@ -1,3 +1,4 @@
+import re
 from typing import Optional, Union, Literal, TYPE_CHECKING, NewType, Callable, Any
 from dataclasses import dataclass
 from ..coop import CoopServerResponseError
@@ -112,13 +113,18 @@ class JobsRemoteInferenceHandler:
         )
         logger.add_info("job_uuid", job_uuid)
 
+        remote_inference_url = self.remote_inference_url
+        if "localhost" in remote_inference_url:
+            remote_inference_url = remote_inference_url.replace("8000", "1234")
         logger.update(
-            f"Job details are available at your Coop account. [Go to Remote Inference page]({self.remote_inference_url})",
+            f"Job details are available at your Coop account. [Go to Remote Inference page]({remote_inference_url})",
             status=JobsStatus.RUNNING,
         )
         progress_bar_url = (
             f"{self.expected_parrot_url}/home/remote-job-progress/{job_uuid}"
         )
+        if "localhost" in progress_bar_url:
+            progress_bar_url = progress_bar_url.replace("8000", "1234")
         logger.add_info("progress_bar_url", progress_bar_url)
         logger.update(
             f"View job progress [here]({progress_bar_url})", status=JobsStatus.RUNNING
@@ -200,10 +206,35 @@ class JobsRemoteInferenceHandler:
             status=JobsStatus.FAILED,
         )
 
+    def _handle_partially_failed_job_interview_details(
+        self, job_info: RemoteJobInfo, remote_job_data: RemoteInferenceResponse
+    ) -> None:
+        "Extracts the interview details from the remote job data."
+        try:
+            # Job details is a string of the form "64 out of 1,758 interviews failed"
+            job_details = remote_job_data.get("latest_failure_description")
+
+            text_without_commas = job_details.replace(",", "")
+
+            # Find all numbers in the text
+            numbers = [int(num) for num in re.findall(r"\d+", text_without_commas)]
+
+            failed = numbers[0]
+            total = numbers[1]
+            completed = total - failed
+
+            job_info.logger.add_info("completed_interviews", completed)
+            job_info.logger.add_info("failed_interviews", failed)
+        # This is mainly helpful metadata, and any errors here should not stop the code
+        except:
+            pass
+
     def _handle_partially_failed_job(
         self, job_info: RemoteJobInfo, remote_job_data: RemoteInferenceResponse
     ) -> None:
         "Handles a partially failed job by logging the error and updating the job status."
+        self._handle_partially_failed_job_interview_details(job_info, remote_job_data)
+
         latest_error_report_url = remote_job_data.get("latest_error_report_url")
 
         if latest_error_report_url:
@@ -244,6 +275,8 @@ class JobsRemoteInferenceHandler:
         job_info.logger.add_info("results_uuid", results_uuid)
         results = object_fetcher(results_uuid, expected_object_type="results")
         results_url = remote_job_data.get("results_url")
+        if "localhost" in results_url:
+            results_url = results_url.replace("8000", "1234")
         job_info.logger.add_info("results_url", results_url)
 
         if job_status == "completed":
@@ -256,6 +289,7 @@ class JobsRemoteInferenceHandler:
                 f"View partial results [here]({results_url})",
                 status=JobsStatus.PARTIALLY_FAILED,
             )
+
         results.job_uuid = job_info.job_uuid
         results.results_uuid = results_uuid
         return results

edsl/key_management/key_lookup_builder.py

@@ -2,6 +2,7 @@ from typing import Optional, TYPE_CHECKING
 import os
 from functools import lru_cache
 import textwrap
+import requests
 
 if TYPE_CHECKING:
     from ..coop import Coop
@@ -255,7 +256,11 @@ class KeyLookupBuilder:
         return dict(list(os.environ.items()))
 
     def _coop_key_value_pairs(self):
-        return dict(list(self.coop.fetch_rate_limit_config_vars().items()))
+        try:
+            return dict(list(self.coop.fetch_rate_limit_config_vars().items()))
+        except requests.ConnectionError:
+            # If connection fails, return empty dict instead of raising error
+            return {}
 
     def _config_key_value_pairs(self):
         from ..config import CONFIG

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.
 
@@ -679,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,
@@ -844,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_,
@@ -873,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_,
         }
 
@@ -913,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:
@@ -999,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

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/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):