edsl 0.1.47__py3-none-any.whl → 0.1.48__py3-none-any.whl

edsl/coop/{ExpectedParrotKeyHandler.py → ep_key_handling.py}

@@ -17,16 +17,55 @@ def get_input_with_timeout(prompt, timeout=5, default="y"):
 
 
 class ExpectedParrotKeyHandler:
+    """
+    Manages Expected Parrot API keys for user authentication.
+    
+    This class handles the storage, retrieval, and management of Expected Parrot API keys.
+    It provides functionality to securely store API keys in platform-specific user
+    configuration directories and retrieve them when needed. It also handles key
+    preference management (e.g., environment variables vs. stored keys).
+    
+    The key handler follows a priority order when retrieving keys:
+    1. Environment variables (EXPECTED_PARROT_API_KEY)
+    2. Platform-specific user config directory
+    
+    Attributes:
+        asked_to_store_file_name (str): Filename for tracking if user was asked about storage
+        ep_key_file_name (str): Filename for the stored API key
+        application_name (str): Application name for the config directory
+    """
     asked_to_store_file_name = "asked_to_store.txt"
     ep_key_file_name = "ep_api_key.txt"
     application_name = "edsl"
 
     @property
-    def config_dir(self):
+    def config_dir(self) -> str:
+        """
+        Get the platform-specific user configuration directory for the application.
+        
+        This property uses the platformdirs library to determine the appropriate
+        location for storing configuration files based on the user's operating system.
+        
+        Returns:
+            str: Path to the user configuration directory
+            
+        Notes:
+            - On Windows, typically: C:\\Users\\<username>\\AppData\\Local\\edsl
+            - On macOS, typically: ~/Library/Application Support/edsl
+            - On Linux, typically: ~/.config/edsl
+        """
         return platformdirs.user_config_dir(self.application_name)
 
     def _ep_key_file_exists(self) -> bool:
-        """Check if the Expected Parrot key file exists."""
+        """
+        Check if the Expected Parrot key file exists in the config directory.
+        
+        Returns:
+            bool: True if the key file exists, False otherwise
+            
+        Notes:
+            - Does not check the validity of the stored key
+        """
         return Path(self.config_dir).joinpath(self.ep_key_file_name).exists()
 
     def ok_to_ask_to_store(self):
@@ -73,37 +112,59 @@ class ExpectedParrotKeyHandler:
                     f.write("Yes")
         return False
 
-    def get_ep_api_key(self):
-        # check if the key is stored in the config_dir
+    def get_ep_api_key(self) -> str:
+        """
+        Retrieve the Expected Parrot API key from available sources.
+        
+        This method checks multiple sources for the API key, with the following priority:
+        1. Environment variable (EXPECTED_PARROT_API_KEY)
+        2. Stored key in the user's config directory
+        
+        If keys are found in multiple sources and they differ, the environment
+        variable takes precedence. A warning is issued in this case.
+        
+        Returns:
+            str: The Expected Parrot API key, or None if not found
+            
+        Notes:
+            - If a key is found, it will attempt to store it persistently if appropriate
+            - Warnings are issued if conflicting keys are found in different sources
+            - Environment variables always take precedence over stored keys
+        """
+        # Initialize variables
         api_key = None
         api_key_from_cache = None
         api_key_from_os = None
 
+        # Try to get key from config directory
         if self._ep_key_file_exists():
             with open(Path(self.config_dir).joinpath(self.ep_key_file_name), "r") as f:
                 api_key_from_cache = f.read().strip()
 
+        # Try to get key from environment variable
         api_key_from_os = os.getenv("EXPECTED_PARROT_API_KEY")
 
+        # Handle the case where both sources have keys
         if api_key_from_os and api_key_from_cache:
             if api_key_from_os != api_key_from_cache:
                 import warnings
-
                 warnings.warn(
                     "WARNING: The Expected Parrot API key from the environment variable "
                     "differs from the one stored in the config directory. Using the one "
                     "from the environment variable."
                 )
             api_key = api_key_from_os
-
-        if api_key_from_os and not api_key_from_cache:
+        # Handle the case where only OS environment has key
+        elif api_key_from_os:
             api_key = api_key_from_os
-
-        if not api_key_from_os and api_key_from_cache:
+        # Handle the case where only cached key exists
+        elif api_key_from_cache:
             api_key = api_key_from_cache
 
+        # If a key was found, ask to store it persistently
         if api_key is not None:
             _ = self.ask_to_store(api_key)
+            
         return api_key
 
     def delete_ep_api_key(self):
@@ -112,7 +173,22 @@ class ExpectedParrotKeyHandler:
             os.remove(key_path)
             print("Deleted Expected Parrot API key at ", key_path)
 
-    def store_ep_api_key(self, api_key):
+    def store_ep_api_key(self, api_key: str) -> None:
+        """
+        Store the Expected Parrot API key in the user's config directory.
+        
+        This method saves the provided API key to a file in the platform-specific
+        user configuration directory, creating the directory if it doesn't exist.
+        
+        Parameters:
+            api_key (str): The API key to store
+            
+        Notes:
+            - The key is stored in plain text in the user's config directory
+            - The directory is created if it doesn't exist
+            - Any existing key file will be overwritten
+            - The location of the config directory is platform-specific
+        """
         # Create the directory if it doesn't exist
         os.makedirs(self.config_dir, exist_ok=True)
 

edsl/coop/exceptions.py

@@ -0,0 +1,62 @@
+"""
+Exceptions specific to the Expected Parrot cloud service integration.
+
+This module defines exception classes for various error conditions that can
+occur when interacting with Expected Parrot's cloud services through the
+Coop module. These exceptions help with specific error handling and reporting.
+"""
+
+from ..base import BaseException
+
+
+class CoopErrors(BaseException):
+    """
+    Base class for all Coop-related exceptions.
+    
+    This is the parent class for all exceptions raised by the Coop module.
+    It inherits from EDSL's BaseException to maintain consistency with
+    the library's exception hierarchy.
+    """
+    pass
+
+
+class CoopInvalidURLError(CoopErrors):
+    """
+    Exception raised when an invalid URL format is provided.
+    
+    This exception is raised when a URL provided to the Coop client
+    does not match the expected format for object or resource URLs.
+    
+    Example:
+        When a URL doesn't follow the pattern:
+        - https://expectedparrot.com/content/{uuid}
+        - https://expectedparrot.com/content/{username}/{alias}
+    """
+    pass
+
+
+class CoopNoUUIDError(CoopErrors):
+    """
+    Exception raised when a required UUID or identifier is missing.
+    
+    This exception is raised when an operation requires a UUID or other
+    identifier, but none was provided.
+    
+    Example:
+        When calling get() without providing a UUID or URL
+    """
+    pass
+
+
+class CoopServerResponseError(CoopErrors):
+    """
+    Exception raised when the server returns an error response.
+    
+    This exception is raised when the Expected Parrot API returns an error
+    response, such as authentication failures, rate limits, or server errors.
+    The exception message typically includes the error details from the server.
+    
+    Example:
+        When the server returns a 401 Unauthorized response due to an invalid API key
+    """
+    pass

edsl/coop/price_fetcher.py

@@ -0,0 +1,126 @@
+"""
+Module for retrieving and caching language model pricing information.
+
+This module provides functionality to fetch current pricing information for various
+language models from the Expected Parrot API. It uses a singleton pattern to ensure
+that price information is only fetched once and then cached for efficiency.
+"""
+
+import requests
+import csv
+from io import StringIO
+from typing import Dict, Tuple, Optional, Any
+
+
+class PriceFetcher:
+    """
+    A singleton class for fetching and caching language model pricing information.
+    
+    This class retrieves the current pricing for language models from the Expected
+    Parrot API and caches it to avoid unnecessary network requests. It implements
+    a singleton pattern to ensure that only one instance exists throughout the
+    application.
+    
+    Attributes:
+        _instance (PriceFetcher): The singleton instance of the class
+        _cached_prices (Dict[Tuple[str, str], Dict]): Cached pricing information
+            mapping (service, model) tuples to their pricing details
+    """
+    _instance = None
+    
+    def __new__(cls):
+        """
+        Create or return the singleton instance of PriceFetcher.
+        
+        This method ensures that only one instance of PriceFetcher exists.
+        When called multiple times, it returns the same instance.
+        
+        Returns:
+            PriceFetcher: The singleton instance
+        """
+        if cls._instance is None:
+            cls._instance = super(PriceFetcher, cls).__new__(cls)
+            cls._instance._cached_prices = None
+        return cls._instance
+
+    def fetch_prices(self) -> Dict[Tuple[str, str], Dict[str, Dict[str, Any]]]:
+        """
+        Fetch current pricing information for language models.
+        
+        This method retrieves the latest pricing information from the Expected Parrot API
+        for all supported language models. It caches the results to avoid redundant API calls.
+        
+        Returns:
+            Dict[Tuple[str, str], Dict[str, Dict[str, Any]]]: A dictionary mapping 
+            (service, model) tuples to their pricing information for different token types.
+            Structure example:
+            {
+                ('openai', 'gpt-4'): {
+                    'input': {
+                        'service': 'openai',
+                        'model': 'gpt-4',
+                        'token_type': 'input',
+                        'usd_per_1M_tokens': 30.0,
+                        ...
+                    },
+                    'output': {
+                        'service': 'openai',
+                        'model': 'gpt-4',
+                        'token_type': 'output',
+                        'usd_per_1M_tokens': 60.0,
+                        ...
+                    }
+                },
+                ...
+            }
+            
+        Notes:
+            - If the request fails, returns an empty dictionary
+            - Uses caching to avoid redundant API calls
+            - Pricing data includes both input and output token costs
+            - This method is automatically called by the Coop.fetch_prices() method
+        """
+        # Return cached prices if available
+        if self._cached_prices is not None:
+            return self._cached_prices
+
+        import os
+        import requests
+        from ..config import CONFIG
+
+        try:
+            # Fetch the pricing data
+            url = f"{CONFIG.EXPECTED_PARROT_URL}/api/v0/prices"
+            api_key = os.getenv("EXPECTED_PARROT_API_KEY")
+            headers = {}
+            if api_key:
+                headers["Authorization"] = f"Bearer {api_key}"
+            else:
+                headers["Authorization"] = f"Bearer None"
+
+            response = requests.get(url, headers=headers, timeout=20)
+            response.raise_for_status()  # Raise an exception for bad responses
+
+            # Parse the data
+            data = response.json()
+
+            # Organize pricing data by (service, model) and token type
+            price_lookup = {}
+            for entry in data:
+                service = entry.get("service", None)
+                model = entry.get("model", None)
+                if service and model:
+                    token_type = entry.get("token_type", None)
+                    if (service, model) in price_lookup:
+                        price_lookup[(service, model)].update({token_type: entry})
+                    else:
+                        price_lookup[(service, model)] = {token_type: entry}
+                        
+            # Cache and return the results
+            self._cached_prices = price_lookup
+            return self._cached_prices
+
+        except requests.RequestException as e:
+            # Silently handle errors and return empty dict
+            # print(f"An error occurred: {e}")
+            return {}

edsl/coop/utils.py

@@ -1,18 +1,16 @@
 from typing import Literal, Optional, Type, Union
 
-from edsl.agents.Agent import Agent
-from edsl.agents.AgentList import AgentList
-from edsl.data.Cache import Cache
-from edsl.language_models.ModelList import ModelList
-from edsl.notebooks.Notebook import Notebook
-from edsl.results.Results import Results
-from edsl.scenarios.Scenario import Scenario
-from edsl.scenarios.ScenarioList import ScenarioList
-from edsl.surveys.Survey import Survey
-from edsl.study.Study import Study
-
-from edsl.language_models.LanguageModel import LanguageModel
-from edsl.questions.QuestionBase import QuestionBase
+from ..agents import Agent, AgentList
+from ..caching import Cache
+from ..language_models import ModelList
+from ..notebooks import Notebook
+from ..results import Results
+from ..scenarios import Scenario, ScenarioList
+from ..surveys import Survey
+from ..study import Study
+
+from ..language_models import LanguageModel
+from ..questions import QuestionBase
 
 EDSLObject = Union[
     Agent,
@@ -61,7 +59,23 @@ VisibilityType = Literal[
 
 class ObjectRegistry:
     """
-    Utility class to map object types to database models.
+    Registry that maps between EDSL class types and their cloud storage object types.
+    
+    This utility class maintains a bidirectional mapping between EDSL Python classes
+    (like Survey, Agent, Results) and their corresponding object type identifiers
+    used in the cloud storage system. It enables the proper serialization,
+    deserialization, and type checking for objects stored in Expected Parrot's
+    cloud services.
+    
+    The registry is used by the Coop client to:
+    1. Determine the correct object type when uploading EDSL objects
+    2. Instantiate the correct class when downloading objects
+    3. Validate that retrieved objects match expected types
+    
+    Attributes:
+        objects (list): List of mappings between object types and EDSL classes
+        object_type_to_edsl_class (dict): Maps object type strings to EDSL classes
+        edsl_class_to_object_type (dict): Maps EDSL class names to object type strings
     """
 
     objects = [
@@ -78,6 +92,8 @@ class ObjectRegistry:
         {"object_type": "survey", "edsl_class": Survey},
         {"object_type": "study", "edsl_class": Study},
     ]
+    
+    # Create mappings for efficient lookups
     object_type_to_edsl_class = {o["object_type"]: o["edsl_class"] for o in objects}
     edsl_class_to_object_type = {
         o["edsl_class"].__name__: o["object_type"] for o in objects
@@ -85,12 +101,36 @@ class ObjectRegistry:
 
     @classmethod
     def get_object_type_by_edsl_class(cls, edsl_object: EDSLObject) -> ObjectType:
+        """
+        Get the object type identifier for an EDSL class or instance.
+        
+        This method determines the appropriate object type string for a given EDSL class
+        or instance, which is needed when storing the object in the cloud.
+        
+        Parameters:
+            edsl_object (EDSLObject): An EDSL class (type) or instance
+            
+        Returns:
+            ObjectType: The corresponding object type string (e.g., "survey", "agent")
+            
+        Raises:
+            ValueError: If no mapping exists for the provided object
+            
+        Notes:
+            - Special handling for Question classes, which all map to "question"
+            - Works with both class types and instances
+        """
+        # Handle both class objects and instances
         if isinstance(edsl_object, type):
             edsl_class_name = edsl_object.__name__
         else:
             edsl_class_name = type(edsl_object).__name__
+            
+        # Special handling for question classes
         if edsl_class_name.startswith("Question"):
             edsl_class_name = "QuestionBase"
+            
+        # Look up the object type
         object_type = cls.edsl_class_to_object_type.get(edsl_class_name)
         if object_type is None:
             raise ValueError(f"Object type not found for {edsl_object=}")
@@ -98,10 +138,25 @@ class ObjectRegistry:
 
     @classmethod
     def get_edsl_class_by_object_type(cls, object_type: ObjectType) -> EDSLObject:
-        EDSL_object = cls.object_type_to_edsl_class.get(object_type)
-        if EDSL_object is None:
+        """
+        Get the EDSL class for a given object type identifier.
+        
+        This method returns the appropriate EDSL class for a given object type string,
+        which is needed when retrieving objects from the cloud.
+        
+        Parameters:
+            object_type (ObjectType): The object type string (e.g., "survey", "agent")
+            
+        Returns:
+            EDSLObject: The corresponding EDSL class
+            
+        Raises:
+            ValueError: If no mapping exists for the provided object type
+        """
+        EDSL_class = cls.object_type_to_edsl_class.get(object_type)
+        if EDSL_class is None:
             raise ValueError(f"EDSL class not found for {object_type=}")
-        return EDSL_object
+        return EDSL_class
 
     @classmethod
     def get_registry(
@@ -110,14 +165,24 @@ class ObjectRegistry:
         exclude_classes: Optional[list] = None,
     ) -> dict:
         """
-        Return the registry of objects.
-
-        Exclude objects that are already registered in subclass_registry.
-        This allows the user to isolate Coop-only objects.
-
-        Also exclude objects if their class name is in the exclude_classes list.
+        Get a filtered registry of EDSL classes.
+        
+        This method returns a dictionary of EDSL classes, optionally excluding
+        classes that are already registered elsewhere or explicitly excluded.
+        
+        Parameters:
+            subclass_registry (dict, optional): Dictionary of classes to exclude
+                because they are already registered elsewhere
+            exclude_classes (list, optional): List of class names to explicitly exclude
+            
+        Returns:
+            dict: Dictionary mapping class names to EDSL classes
+            
+        Notes:
+            - This method is useful for building registries of classes that
+              can be serialized and stored in the cloud
+            - It helps avoid duplicate registrations of classes
         """
-
         if subclass_registry is None:
             subclass_registry = {}
         if exclude_classes is None:

edsl/data_transfer_models.py

@@ -1,74 +1,7 @@
-from typing import NamedTuple, Dict, List, Optional, Any
-from dataclasses import dataclass, fields
+"""Data transfer models module wrapper for backward compatibility.
 
+This module re-exports everything from the edsl.base.data_transfer_models module.
+"""
 
-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
-
-
-class EDSLOutput(NamedTuple):
-    "This is the edsl dictionary that is returned by the model"
-    answer: Any
-    generated_tokens: str
-    comment: Optional[str] = None
-
-
-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
-
-
-class AgentResponseDict(NamedTuple):
-    edsl_dict: EDSLOutput
-    model_inputs: ModelInputs
-    model_outputs: ModelResponse
-
-
-class EDSLResultObjectInput(NamedTuple):
-    generated_tokens: str
-    question_name: str
-    prompts: dict
-    cached_response: str
-    raw_model_response: str
-    cache_used: bool
-    cache_key: str
-    answer: Any
-    comment: str
-    validated: bool = False
-    exception_occurred: Exception = None
-    cost: Optional[float] = None
-
-
-@dataclass
-class ImageInfo:
-    file_path: str
-    file_name: str
-    image_format: str
-    file_size: int
-    encoded_image: str
-
-    def __repr__(self):
-        import reprlib
-
-        reprlib_instance = reprlib.Repr()
-        reprlib_instance.maxstring = 30  # Limit the string length for the encoded image
-
-        # Get all fields except encoded_image
-        field_reprs = [
-            f"{f.name}={getattr(self, f.name)!r}"
-            for f in fields(self)
-            if f.name != "encoded_image"
-        ]
-
-        # Add the reprlib-restricted encoded_image field
-        field_reprs.append(f"encoded_image={reprlib_instance.repr(self.encoded_image)}")
-
-        # Join everything to create the repr
-        return f"{self.__class__.__name__}({', '.join(field_reprs)})"
+# Re-export everything from edsl.base.data_transfer_models module
+from edsl.base.data_transfer_models import *

edsl/dataset/__init__.py

@@ -0,0 +1,10 @@
+from .dataset import Dataset
+
+from .dataset_operations_mixin import AgentListOperationsMixin
+from .dataset_operations_mixin import ScenarioListOperationsMixin
+from .dataset_operations_mixin import DatasetOperationsMixin
+from .dataset_operations_mixin import ResultsOperationsMixin
+
+__all__ = [
+    "Dataset",
+]