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

edsl/language_models/language_model.py

@@ -49,6 +49,7 @@ from ..data_transfer_models import (
 )
 
 if TYPE_CHECKING:
+    from .price_manager import ResponseCost
     from ..caching import Cache
     from ..scenarios import FileStore
     from ..questions import QuestionBase
@@ -365,6 +366,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 +733,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,
@@ -726,13 +783,18 @@ class LanguageModel(
         # Calculate cost for the response
         cost = self.cost(response)
         # Return a structured response with metadata
-        return ModelResponse(
+        response = ModelResponse(
             response=response,
             cache_used=cache_used,
             cache_key=cache_key,
             cached_response=cached_response,
-            cost=cost,
+            input_tokens=cost.input_tokens,
+            output_tokens=cost.output_tokens,
+            input_price_per_million_tokens=cost.input_price_per_million_tokens,
+            output_price_per_million_tokens=cost.output_price_per_million_tokens,
+            total_cost=cost.total_cost,
         )
+        return response
 
     _get_intended_model_call_outcome = sync_wrapper(
         _async_get_intended_model_call_outcome
@@ -825,7 +887,7 @@ class LanguageModel(
 
     get_response = sync_wrapper(async_get_response)
 
-    def cost(self, raw_response: dict[str, Any]) -> Union[float, str]:
+    def cost(self, raw_response: dict[str, Any]) -> ResponseCost:
         """Calculate the monetary cost of a model API call.
 
         This method extracts token usage information from the response and
@@ -836,7 +898,7 @@ class LanguageModel(
             raw_response: The complete response dictionary from the model API
 
         Returns:
-            Union[float, str]: The calculated cost in dollars, or an error message
+            ResponseCost: Object containing token counts and total cost
         """
         # Extract token usage data from the response
         usage = self.get_usage_dict(raw_response)
@@ -844,7 +906,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 +935,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 +981,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 +1085,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
@@ -1067,13 +1153,25 @@ class LanguageModel(
             }
             cached_response, cache_key = cache.fetch(**cache_call_params)
             response = json.loads(cached_response)
-            cost = 0
+
+            try:
+                usage = self.get_usage_dict(response)
+                input_tokens = int(usage[self.input_token_name])
+                output_tokens = int(usage[self.output_token_name])
+            except Exception as e:
+                print(f"Could not fetch tokens from model response: {e}")
+                input_tokens = None
+                output_tokens = None
             return ModelResponse(
                 response=response,
                 cache_used=True,
                 cache_key=cache_key,
                 cached_response=cached_response,
-                cost=cost,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                input_price_per_million_tokens=0,
+                output_price_per_million_tokens=0,
+                total_cost=0,
             )
 
         # Bind the new method to the copied instance

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

@@ -1,4 +1,22 @@
-from typing import Dict, Tuple, Union
+from dataclasses import dataclass
+from typing import Dict, Literal, Tuple, Union
+from collections import namedtuple
+
+
+@dataclass
+class ResponseCost:
+    """
+    Class for storing the cost and token usage of a language model response.
+
+    If an error occurs when computing the cost, the total_cost will contain a string with the error message.
+    All other fields will be None.
+    """
+
+    input_tokens: Union[int, None] = None
+    output_tokens: Union[int, None] = None
+    input_price_per_million_tokens: Union[float, None] = None
+    output_price_per_million_tokens: Union[float, None] = None
+    total_cost: Union[float, str, None] = None
 
 
 class PriceManager:
@@ -8,20 +26,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,29 +71,14 @@ 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:
@@ -68,73 +93,95 @@ class PriceManager:
         Returns:
             Dict: Price information
         """
+        PriceEntry = namedtuple("PriceEntry", ["tokens_per_usd", "price_info"])
+
         service_prices = [
             prices
             for (service, _), prices in self._price_lookup.items()
             if service == inference_service
         ]
 
-        input_tokens_per_usd = [
-            float(p["input"]["one_usd_buys"]) for p in service_prices if "input" in p
+        default_price_info = {
+            "one_usd_buys": 1_000_000,
+            "service_stated_token_qty": 1_000_000,
+            "service_stated_token_price": 1.0,
+        }
+
+        # Find the most expensive price entries (lowest tokens per USD)
+        input_price_info = default_price_info
+        output_price_info = default_price_info
+
+        input_prices = [
+            PriceEntry(float(p["input"]["one_usd_buys"]), p["input"])
+            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
+        if input_prices:
+            input_price_info = min(
+                input_prices, key=lambda price: price.tokens_per_usd
+            ).price_info
 
-        output_tokens_per_usd = [
-            float(p["output"]["one_usd_buys"]) for p in service_prices if "output" in p
+        output_prices = [
+            PriceEntry(float(p["output"]["one_usd_buys"]), p["output"])
+            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
+        if output_prices:
+            output_price_info = min(
+                output_prices, key=lambda price: price.tokens_per_usd
+            ).price_info
 
         return {
-            "input": {"one_usd_buys": min_input_tokens},
-            "output": {"one_usd_buys": min_output_tokens},
+            "input": input_price_info,
+            "output": output_price_info,
         }
 
-    def calculate_cost(
+    def get_price_per_million_tokens(
         self,
-        inference_service: str,
-        model: str,
-        usage: Dict[str, Union[str, int]],
-        input_token_name: str,
-        output_token_name: str,
-    ) -> Union[float, str]:
+        relevant_prices: Dict,
+        token_type: Literal["input", "output"],
+    ) -> Dict:
         """
-        Calculate the total cost for a model usage based on input and output tokens.
+        Get the price per million tokens for a specific service, model, and token type.
+        """
+        service_price = relevant_prices[token_type]["service_stated_token_price"]
+        service_qty = relevant_prices[token_type]["service_stated_token_qty"]
 
-        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
+        if service_qty == 1_000_000:
+            price_per_million_tokens = service_price
+        elif service_qty == 1_000:
+            price_per_million_tokens = service_price * 1_000
+        else:
+            price_per_token = service_price / service_qty
+            price_per_million_tokens = round(price_per_token * 1_000_000, 10)
+        return price_per_million_tokens
 
-        Returns:
-            Union[float, str]: Total cost if calculation successful, error message string if not
+    def _calculate_total_cost(
+        self,
+        relevant_prices: Dict,
+        input_tokens: int,
+        output_tokens: int,
+    ) -> float:
         """
-        relevant_prices = self.get_price(inference_service, model)
-
-        # Extract token counts
-        try:
-            input_tokens = int(usage[input_token_name])
-            output_tokens = int(usage[output_token_name])
-        except Exception as e:
-            return f"Could not fetch tokens from model response: {e}"
+        Calculate the total cost for a model usage based on input and output tokens.
 
+        Returns:
+            float: Total cost
+        """
         # Extract price information
         try:
             inverse_output_price = relevant_prices["output"]["one_usd_buys"]
             inverse_input_price = relevant_prices["input"]["one_usd_buys"]
         except Exception as e:
             if "output" not in relevant_prices:
-                return f"Could not fetch prices from {relevant_prices} - {e}; Missing 'output' key."
+                raise KeyError(
+                    f"Could not fetch prices from {relevant_prices} - {e}; Missing 'output' key."
+                )
             if "input" not in relevant_prices:
-                return f"Could not fetch prices from {relevant_prices} - {e}; Missing 'input' key."
-            return f"Could not fetch prices from {relevant_prices} - {e}"
+                raise KeyError(
+                    f"Could not fetch prices from {relevant_prices} - {e}; Missing 'input' key."
+                )
+            raise Exception(f"Could not fetch prices from {relevant_prices} - {e}")
 
         # Calculate input cost
         if inverse_input_price == "infinity":
@@ -143,7 +190,7 @@ class PriceManager:
             try:
                 input_cost = input_tokens / float(inverse_input_price)
             except Exception as e:
-                return f"Could not compute input price - {e}."
+                raise Exception(f"Could not compute input price - {e}")
 
         # Calculate output cost
         if inverse_output_price == "infinity":
@@ -152,16 +199,74 @@ class PriceManager:
             try:
                 output_cost = output_tokens / float(inverse_output_price)
             except Exception as e:
-                return f"Could not compute output price - {e}"
+                raise Exception(f"Could not compute output price - {e}")
 
         return input_cost + output_cost
 
-    @property
-    def is_initialized(self) -> bool:
+    def calculate_cost(
+        self,
+        inference_service: str,
+        model: str,
+        usage: Dict[str, Union[str, int]],
+        input_token_name: str,
+        output_token_name: str,
+    ) -> ResponseCost:
         """
-        Check if the PriceManager has been initialized.
+        Calculate the cost and token usage for a model response.
+
+        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:
-            bool: True if initialized, False otherwise
+            ResponseCost: Object containing token counts and total cost
         """
+        try:
+            input_tokens = int(usage[input_token_name])
+            output_tokens = int(usage[output_token_name])
+        except Exception as e:
+            return ResponseCost(
+                total_cost=f"Could not fetch tokens from model response: {e}",
+            )
+
+        try:
+            relevant_prices = self.get_price(inference_service, model)
+        except Exception as e:
+            return ResponseCost(
+                total_cost=f"Could not fetch prices from {inference_service} - {model}: {e}",
+            )
+
+        try:
+            input_price_per_million_tokens = self.get_price_per_million_tokens(
+                relevant_prices, "input"
+            )
+            output_price_per_million_tokens = self.get_price_per_million_tokens(
+                relevant_prices, "output"
+            )
+        except Exception as e:
+            return ResponseCost(
+                total_cost=f"Could not compute price per million tokens: {e}",
+            )
+
+        try:
+            total_cost = self._calculate_total_cost(
+                relevant_prices, input_tokens, output_tokens
+            )
+        except Exception as e:
+            return ResponseCost(total_cost=f"{e}")
+
+        return ResponseCost(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            input_price_per_million_tokens=input_price_per_million_tokens,
+            output_price_per_million_tokens=output_price_per_million_tokens,
+            total_cost=total_cost,
+        )
+
+    @property
+    def is_initialized(self) -> bool:
+        """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."""