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

edsl/base/data_transfer_models.py

@@ -1,10 +1,11 @@
 from collections import UserDict
-from typing import NamedTuple, Dict, Optional, Any
+from typing import NamedTuple, Dict, Optional, Any, Union
 from dataclasses import dataclass, fields
 
 
 class ModelInputs(NamedTuple):
     "This is what was send by the agent to the model"
+
     user_prompt: str
     system_prompt: str
     encoded_image: Optional[str] = None
@@ -12,6 +13,7 @@ class ModelInputs(NamedTuple):
 
 class EDSLOutput(NamedTuple):
     "This is the edsl dictionary that is returned by the model"
+
     answer: Any
     generated_tokens: str
     comment: Optional[str] = None
@@ -19,11 +21,16 @@ class EDSLOutput(NamedTuple):
 
 class ModelResponse(NamedTuple):
     "This is the metadata that is returned by the model and includes info about the cache"
+
     response: dict
     cache_used: bool
     cache_key: str
     cached_response: Optional[Dict[str, Any]] = None
-    cost: Optional[float] = None
+    input_tokens: Optional[int] = None
+    output_tokens: Optional[int] = None
+    input_price_per_million_tokens: Optional[float] = None
+    output_price_per_million_tokens: Optional[float] = None
+    total_cost: Optional[Union[float, str]] = None
 
 
 class AgentResponseDict(NamedTuple):
@@ -44,7 +51,11 @@ class EDSLResultObjectInput(NamedTuple):
     comment: str
     validated: bool = False
     exception_occurred: Exception = None
-    cost: Optional[float] = None
+    input_tokens: Optional[int] = None
+    output_tokens: Optional[int] = None
+    input_price_per_million_tokens: Optional[float] = None
+    output_price_per_million_tokens: Optional[float] = None
+    total_cost: Optional[Union[float, str]] = None
 
 
 @dataclass
@@ -111,4 +122,4 @@ class Answers(UserDict):
 if __name__ == "__main__":
     import doctest
 
-    doctest.testmod()
+    doctest.testmod()

edsl/buckets/__init__.py

@@ -18,9 +18,6 @@ client (token_bucket_client) for distributed rate limiting scenarios where
 multiple processes or machines need to share rate limits.
 """
 
-from .bucket_collection import BucketCollection
-from .model_buckets import ModelBuckets
-from .token_bucket import TokenBucket
 from .exceptions import (
     BucketError,
     TokenLimitError,
@@ -28,10 +25,18 @@ from .exceptions import (
     BucketConfigurationError,
 )
 
+from .token_bucket import TokenBucket
+from .model_buckets import ModelBuckets
+from .token_bucket_client import TokenBucketClient  # Add explicit import for TokenBucketClient
+
+# Import BucketCollection last to avoid circular import issues
+from .bucket_collection import BucketCollection
+
 __all__ = [
     "BucketCollection", 
     "ModelBuckets", 
     "TokenBucket",
+    "TokenBucketClient",
     "BucketError",
     "TokenLimitError",
     "TokenBucketClientError",

edsl/buckets/bucket_collection.py

@@ -10,11 +10,14 @@ share the same rate limit buckets.
 from typing import TYPE_CHECKING, Dict, List, Tuple
 from collections import UserDict
 from threading import RLock
-from matplotlib.figure import Figure
+import functools
+import inspect
+
+# Import the synchronized_class decorator directly
+from ..jobs.decorators import synchronized_class
 
 from .token_bucket import TokenBucket
 from .model_buckets import ModelBuckets
-from ..jobs.decorators import synchronized_class
 
 if TYPE_CHECKING:
     from ..language_models import LanguageModel
@@ -262,7 +265,7 @@ class BucketCollection(UserDict):
                     )
                     self.services_to_buckets[service].tokens_bucket = new_tokens_bucket
 
-    def visualize(self) -> Dict["LanguageModel", Tuple[Figure, Figure]]:
+    def visualize(self) -> Dict["LanguageModel", Tuple["Figure", "Figure"]]:
         """
         Visualize the token and request buckets for all models.
         
@@ -279,6 +282,9 @@ class BucketCollection(UserDict):
             >>> plots = bucket_collection.visualize()
             >>> # Now you can display or save these plots
         """
+        # Import Figure only for type checking when the function is called
+        from matplotlib.figure import Figure
+        
         plots = {}
         for model in self:
             plots[model] = self[model].visualize()

edsl/buckets/model_buckets.py

@@ -7,7 +7,6 @@ instance contains two TokenBucket instances - one for requests and one for token
 """
 
 from typing import TYPE_CHECKING, Tuple
-from matplotlib.figure import Figure
 
 if TYPE_CHECKING:
     from .token_bucket import TokenBucket
@@ -147,7 +146,7 @@ class ModelBuckets:
             ),
         )
 
-    def visualize(self) -> Tuple[Figure, Figure]:
+    def visualize(self) -> Tuple["Figure", "Figure"]:
         """
         Create visualizations of token usage over time for both buckets.
         
@@ -163,6 +162,9 @@ class ModelBuckets:
             >>> ## request_plot, token_plot = buckets.visualize()
             >>> ## Now you can display or save these plots
         """
+        # Import Figure only for type checking when the function is called
+        from matplotlib.figure import Figure
+        
         plot1 = self.requests_bucket.visualize()
         plot2 = self.tokens_bucket.visualize()
         return plot1, plot2

edsl/buckets/token_bucket.py

@@ -80,8 +80,8 @@ class TokenBucket:
             'local-rate-limit'
         """
         if remote_url is not None:
-            # Import here to avoid circular imports
-            from ..buckets import TokenBucketClient
+            # Import the client directly from its module to avoid circular imports
+            from .token_bucket_client import TokenBucketClient
 
             return TokenBucketClient(
                 bucket_name=bucket_name,

edsl/buckets/token_bucket_client.py

@@ -11,8 +11,6 @@ from typing import Union, Optional, Dict, Any
 import asyncio
 import time
 import aiohttp
-from matplotlib import pyplot as plt
-from matplotlib.figure import Figure
 
 from .exceptions import BucketError, TokenBucketClientError
 
@@ -417,7 +415,7 @@ class TokenBucketClient:
     #     """Server-side wait time calculation (future implementation)"""
     #     return 0  # TODO - Need to implement this on the server side
 
-    def visualize(self) -> Figure:
+    def visualize(self) -> "Figure":
         """
         Visualize the token bucket usage over time as a matplotlib figure.
         
@@ -442,6 +440,10 @@ class TokenBucketClient:
         start_time = times[0]
         times = [t - start_time for t in times]
 
+        # Import here to avoid loading matplotlib until needed
+        from matplotlib import pyplot as plt
+        from matplotlib.figure import Figure
+        
         # Create the plot
         fig = plt.figure(figsize=(10, 6))
         plt.plot(times, tokens, label="Tokens Available")

edsl/caching/cache.py

@@ -41,30 +41,31 @@ from .sql_dict import SQLiteDict
 if TYPE_CHECKING:
     from .cache_entry import CacheEntry
 
+
 class Cache(Base):
     """Cache for storing and retrieving language model responses.
-    
+
     The Cache class manages a collection of CacheEntry objects, providing methods for
     storing, retrieving, and persisting language model responses. It serves as the core
     component of EDSL's caching infrastructure, helping to reduce redundant API calls,
     save costs, and ensure reproducibility.
-    
+
     Cache can use different storage backends:
     - In-memory dictionary (default)
     - SQLite database via SQLiteDict
     - JSON lines file (.jsonl)
-    
+
     The cache operates by generating deterministic keys based on the model, parameters,
     prompts, and iteration number. This allows for efficient lookup of cached responses
     when identical requests are made.
-    
+
     Attributes:
         data (dict or SQLiteDict): The primary storage for cache entries
         new_entries (dict): Entries added in the current session
         fetched_data (dict): Entries retrieved in the current session
         filename (str, optional): Path for persistence if provided
         immediate_write (bool): Whether to update data immediately (True) or defer (False)
-        
+
     Technical Notes:
         - Can be used as a context manager to automatically persist changes on exit
         - Supports serialization/deserialization via to_dict/from_dict methods
@@ -86,25 +87,25 @@ class Cache(Base):
         verbose=False,
     ):
         """Initialize a new Cache instance.
-        
-        Creates a new cache for storing language model responses. The cache can be initialized 
+
+        Creates a new cache for storing language model responses. The cache can be initialized
         with existing data or connected to a persistent storage file.
-        
+
         Args:
             filename: Path to a persistent storage file (.jsonl or .db). If provided, the cache
                      will be initialized from this file and changes will be written back to it.
                      Cannot be used together with data parameter.
-            data: Initial cache data as a dictionary or SQLiteDict. Cannot be used together 
+            data: Initial cache data as a dictionary or SQLiteDict. Cannot be used together
                   with filename parameter.
             immediate_write: If True, new entries are immediately added to the main data store.
                             If False, they're kept separate until explicitly written.
             method: Deprecated. Legacy parameter for backward compatibility.
             verbose: If True, prints diagnostic information about cache hits and misses.
-            
+
         Raises:
-            CacheError: If both filename and data are provided, or if the filename has an 
+            CacheError: If both filename and data are provided, or if the filename has an
                        invalid extension.
-                       
+
         Implementation Notes:
             - The cache maintains separate dictionaries for tracking:
               * data: The main persistent storage
@@ -153,12 +154,12 @@ class Cache(Base):
 
     def keys(self):
         """Return a list of all cache keys.
-        
+
         Retrieves all cache keys, which are the unique identifiers for each cache entry.
-        
+
         Returns:
             list: A list of string keys in the cache
-            
+
         Examples:
             >>> from edsl import Cache
             >>> Cache.example().keys()
@@ -168,12 +169,12 @@ class Cache(Base):
 
     def values(self):
         """Return a list of all cache entry values.
-        
+
         Retrieves all CacheEntry objects stored in the cache.
-        
+
         Returns:
             list: A list of CacheEntry objects
-            
+
         Examples:
             >>> from edsl import Cache
             >>> entries = Cache.example().values()
@@ -186,10 +187,10 @@ class Cache(Base):
 
     def items(self):
         """Return an iterator of (key, value) pairs in the cache.
-        
+
         Similar to dict.items(), provides an iterator over all key-value pairs
         in the cache for easy iteration.
-        
+
         Returns:
             zip: An iterator of (key, CacheEntry) tuples
         """
@@ -219,34 +220,35 @@ class Cache(Base):
         system_prompt: str,
         user_prompt: str,
         iteration: int,
+        validated: bool = False,
     ) -> tuple(Union[None, str], str):
         """Retrieve a cached language model response if available.
-        
+
         This method attempts to find a cached response matching the exact input parameters.
         The combination of model, parameters, prompts, and iteration creates a unique key
         that identifies a specific language model request.
-        
+
         Args:
             model: Language model identifier (e.g., "gpt-3.5-turbo")
             parameters: Model configuration parameters (e.g., temperature, max_tokens)
             system_prompt: The system instructions given to the model
             user_prompt: The user query/prompt given to the model
             iteration: The iteration number for this specific request
-            
+
         Returns:
             tuple: (response, key) where:
                 - response: The cached model output as a string, or None if not found
                 - key: The cache key string generated for this request
-                
+
         Technical Notes:
             - Uses CacheEntry.gen_key() to generate a consistent hash-based key
             - Updates self.fetched_data when a hit occurs to track cache usage
             - Optionally logs cache hit/miss when verbose=True
             - The response is returned as a JSON string for consistency
-            
+
         Examples:
             >>> c = Cache()
-            >>> c.fetch(model="gpt-3", parameters="default", system_prompt="Hello", 
+            >>> c.fetch(model="gpt-3", parameters="default", system_prompt="Hello",
             ...         user_prompt="Hi", iteration=1)[0] is None
             True
         """
@@ -278,12 +280,13 @@ class Cache(Base):
         response: dict,
         iteration: int,
         service: str,
+        validated: bool = False,
     ) -> str:
         """Store a new language model response in the cache.
-        
+
         Creates a new CacheEntry from the provided parameters and response, then
         adds it to the cache using a deterministic key derived from the input parameters.
-        
+
         Args:
             model: Language model identifier (e.g., "gpt-3.5-turbo")
             parameters: Model configuration parameters (e.g., temperature, max_tokens)
@@ -292,29 +295,30 @@ class Cache(Base):
             response: The model's response as a dictionary
             iteration: The iteration number for this specific request
             service: The service provider (e.g., "openai", "anthropic")
-            
+            validated: Whether the response has been validated (default: False)
+
         Returns:
             str: The cache key generated for this entry
-            
+
         Technical Notes:
             - Creates a new CacheEntry object to encapsulate the response and metadata
             - Adds the entry to self.new_entries to track entries added in this session
             - Adds the entry to the main data store if immediate_write=True
             - Otherwise, stores in new_entries_to_write_later for deferred writing
             - The response is stored as a JSON string for consistency and compatibility
-            
+
         Storage Behavior:
             The method's behavior depends on the immediate_write setting:
             - If True: Immediately writes to the main data store (self.data)
             - If False: Stores in a separate dict for writing later (e.g., at context exit)
-            
+
         Examples:
             >>> from edsl import Cache, Model, Question
-            >>> m = Model("test") 
+            >>> m = Model("test")
             >>> c = Cache()
             >>> len(c)
             0
-            >>> results = Question.example("free_text").by(m).run(cache=c, 
+            >>> results = Question.example("free_text").by(m).run(cache=c,
             ...         disable_remote_cache=True, disable_remote_inference=True)
             >>> len(c)
             1
@@ -329,6 +333,7 @@ class Cache(Base):
             output=json.dumps(response),
             iteration=iteration,
             service=service,
+            validated=validated,
         )
         key = entry.key
         self.new_entries[key] = entry
@@ -486,20 +491,20 @@ class Cache(Base):
 
     def __floordiv__(self, other: "Cache") -> "Cache":
         """Subtract one cache from another, returning entries unique to this cache.
-        
-        This operator implements set difference between two caches, returning a new cache 
+
+        This operator implements set difference between two caches, returning a new cache
         containing only entries that exist in this cache but not in the other cache.
         The floor division operator (//) is used as an intuitive alternative to subtraction.
-        
+
         Args:
             other: Another Cache object to subtract from this one
-            
+
         Returns:
             Cache: A new Cache containing only entries unique to this cache
-            
+
         Raises:
             CacheError: If the provided object is not a Cache instance
-            
+
         Examples:
             >>> from edsl.caching import CacheEntry
             >>> ce1 = CacheEntry.example(randomize=True)
@@ -511,7 +516,7 @@ class Cache(Base):
             1
             >>> c3.data[ce2.key] == ce2
             True
-            
+
         Technical Notes:
             - Comparison is based on cache keys, not the full entry contents
             - Returns a new Cache instance with the same immediate_write setting
@@ -534,14 +539,14 @@ class Cache(Base):
 
     def __enter__(self):
         """Set up the cache when used as a context manager.
-        
+
         Enables usage of Cache in a with statement, e.g.:
         ```python
         with Cache(filename="my_cache.db") as cache:
             # Use cache...
         # Changes automatically saved when exiting the context
         ```
-        
+
         Returns:
             Cache: The cache instance itself
         """
@@ -549,21 +554,22 @@ class Cache(Base):
 
     def __exit__(self, exc_type, exc_value, traceback):
         """Clean up and persist cache when exiting the context.
-        
+
         This method is called automatically when exiting a with block.
         It performs two key operations:
         1. Writes any deferred entries to the main data store
         2. Persists the cache to disk if a filename was provided
-        
+
         Args:
             exc_type: Exception type if an exception was raised in the with block
             exc_value: Exception value if an exception was raised
             traceback: Traceback if an exception was raised
-            
+
         Technical Notes:
             - Deferred entries (new_entries_to_write_later) are written to the main data store
             - If a filename was provided at initialization, cache is persisted to that file
             - Persistence format is determined by the filename extension (.jsonl or .db)
+            - SQLAlchemy resources are properly disposed when the context is exited
         """
         # Write any deferred entries to the main data store
         for key, entry in self.new_entries_to_write_later.items():
@@ -573,6 +579,9 @@ class Cache(Base):
         if self.filename:
             self.write(self.filename)
 
+        # Clean up SQLAlchemy resources
+        self.close()
+
     def __hash__(self):
         """Return the hash of the Cache."""
 
@@ -580,15 +589,15 @@ class Cache(Base):
 
     def to_dict(self, add_edsl_version=True) -> dict:
         """Serialize the cache to a dictionary for storage or transmission.
-        
+
         Converts the Cache object into a plain dictionary format that can be
         easily serialized to JSON or other formats. Each CacheEntry is also
         converted to a dictionary using its to_dict method.
-        
+
         Args:
             add_edsl_version: If True, includes the EDSL version and class name
                               in the serialized output for compatibility tracking
-                              
+
         Returns:
             dict: A dictionary representation of the cache with the structure:
                 {
@@ -598,7 +607,7 @@ class Cache(Base):
                     "edsl_version": "x.x.x",  # if add_edsl_version=True
                     "edsl_class_name": "Cache"  # if add_edsl_version=True
                 }
-                
+
         Technical Notes:
             - Used by from_dict for deserialization
             - Used by __hash__ for cache comparison
@@ -635,6 +644,46 @@ class Cache(Base):
     def to_dataset(self):
         return self.to_scenario_list().to_dataset()
 
+    def _repr_html_(self):
+        """Generate an HTML representation for Jupyter notebooks.
+
+        This method is automatically called by Jupyter to render the object
+        as HTML in notebook cells. It handles empty caches gracefully.
+
+        Returns:
+            str: HTML representation of the object
+        """
+        # Get class name and documentation link
+        class_name = self.__class__.__name__
+        docs = getattr(self, "__documentation__", "")
+
+        # Create header with link to documentation
+        header = f"<a href='{docs}'>{class_name}</a>"
+
+        # Add summary if available
+        if hasattr(self, "_summary"):
+            summary_dict = self._summary()
+            summary_line = "".join([f" {k}: {v};" for k, v in summary_dict.items()])
+            header = f"<p>{header}{summary_line}</p>"
+        else:
+            header = f"<p>{header}</p>"
+
+        # Handle empty cache
+        if len(self.data) == 0:
+            return f"{header}<p><em>Empty cache</em></p>"
+
+        # For non-empty caches, render the table as usual
+        from edsl.dataset.display.table_display import TableDisplay
+
+        try:
+            return header + self.table()._repr_html_()
+        except Exception:
+            # Fallback if table() fails - display as dictionary
+            display_dict = {"entries": len(self.data)}
+            return (
+                header + TableDisplay.from_dictionary_wide(display_dict)._repr_html_()
+            )
+
     @classmethod
     @remove_edsl_version
     def from_dict(cls, data) -> Cache:
@@ -661,26 +710,26 @@ class Cache(Base):
 
     def __add__(self, other: "Cache"):
         """Combine this cache with another, updating in-place.
-        
+
         This operator implements a set union operation between two caches, adding all
         entries from the other cache into this one. The operation modifies this cache
         in-place rather than creating a new one.
-        
+
         Args:
             other: Another Cache object to merge into this one
-            
+
         Returns:
             Cache: Self, with entries from other added
-            
+
         Raises:
             CacheError: If the provided object is not a Cache instance
-            
+
         Technical Notes:
             - Modifies this cache in-place (unlike __floordiv__ which returns a new cache)
             - If both caches have the same key, this cache's entry will be overwritten
             - Useful for merging caches from different sources
             - No special handling for conflicting entries - last one wins
-            
+
         Examples:
             >>> from edsl.caching import CacheEntry
             >>> ce1 = CacheEntry.example(randomize=True)
@@ -697,6 +746,26 @@ class Cache(Base):
         self.data.update(other.data)
         return self
 
+    def close(self):
+        """Explicitly close and clean up resources.
+
+        This method properly disposes of any SQLAlchemy engines and
+        connections to prevent memory leaks.
+        """
+        # Clean up SQLiteDict resources if present
+        if not isinstance(self.data, dict):
+            # Handle SQLiteDict or other database-backed storage
+            if hasattr(self.data, "engine") and self.data.engine:
+                self.data.engine.dispose()
+
+    def __del__(self):
+        """Destructor for proper resource cleanup.
+
+        Ensures SQLAlchemy connections are properly closed when the Cache
+        object is garbage collected.
+        """
+        self.close()
+
     def __repr__(self):
         """
         Return a string representation of the Cache object.
@@ -773,23 +842,23 @@ class Cache(Base):
     @classmethod
     def example(cls, randomize: bool = False) -> Cache:
         """Create an example Cache instance for testing and demonstration.
-        
+
         Creates a Cache object pre-populated with example CacheEntry objects.
         This method is useful for documentation, testing, and demonstration purposes.
-        
+
         Args:
             randomize: If True, creates CacheEntry objects with randomized content
                       for uniqueness. If False, uses consistent example entries.
-                      
+
         Returns:
             Cache: A new Cache object containing example CacheEntry objects
-            
+
         Technical Notes:
             - Uses CacheEntry.example() to create sample entries
             - When randomize=True, generates unique keys for each call
             - When randomize=False, produces consistent examples for doctests
             - Creates an in-memory cache (no persistent file)
-            
+
         Examples:
             >>> cache = Cache.example()
             >>> len(cache) > 0
@@ -797,7 +866,7 @@ class Cache(Base):
             >>> from edsl.caching.cache_entry import CacheEntry
             >>> all(isinstance(entry, CacheEntry) for entry in cache.values())
             True
-            
+
             >>> # Create examples with randomized content
             >>> cache1 = Cache.example(randomize=True)
             >>> cache2 = Cache.example(randomize=True)