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

edsl/jobs/jobs_runner_asyncio.py

@@ -0,0 +1,282 @@
+"""
+Asynchronous execution engine for EDSL jobs.
+
+This module provides the core functionality for running interviews asynchronously,
+which is essential for efficient execution of large jobs. It handles the complex
+process of coordinating multiple concurrent interviews, managing progress tracking,
+and gracefully handling cancellations and errors.
+
+Key components:
+- JobsRunnerAsyncio: The main class that orchestrates async execution
+- Progress bar integration with remote status tracking
+- Error handling and graceful cancellation
+- Result collection and organization
+
+This module is primarily used internally by the Jobs class and is typically not
+accessed directly by end users, though advanced users may need to understand its
+behavior when customizing job execution.
+"""
+from __future__ import annotations
+import time
+import asyncio
+import threading
+import warnings
+from typing import TYPE_CHECKING, List, Generator, Tuple, Optional, Any
+
+if TYPE_CHECKING:
+    from ..results import Results
+
+from ..results import Results
+from ..tasks import TaskHistory
+from ..utilities.decorators import jupyter_nb_handler
+
+from .jobs_runner_status import JobsRunnerStatus
+from .async_interview_runner import AsyncInterviewRunner
+from .data_structures import RunEnvironment, RunParameters, RunConfig
+
+if TYPE_CHECKING:
+    from ..jobs import Jobs
+    from ..interviews import Interview
+
+
+class JobsRunnerAsyncio:
+    """
+    Executes a collection of interviews asynchronously with progress tracking.
+    
+    This class is the main execution engine for EDSL jobs. It manages the asynchronous
+    running of interviews, handles progress tracking, and organizes results. It is
+    instantiated by a Jobs object and handles the complex execution logic that makes
+    parallel interview processing efficient.
+    
+    Key responsibilities:
+    1. Coordinating asynchronous execution of interviews
+    2. Tracking and reporting progress
+    3. Handling errors and cancellations
+    4. Collecting and organizing results
+    
+    This class supports two main execution modes:
+    - run(): For synchronous contexts (returns after completion)
+    - run_async(): For asynchronous contexts (can be awaited)
+    """
+
+    def __init__(self, jobs: "Jobs", environment: RunEnvironment):
+        """
+        Initialize a JobsRunnerAsyncio instance.
+        
+        Parameters:
+            jobs (Jobs): The Jobs instance containing the interviews to run
+            environment (RunEnvironment): The environment configuration containing
+                resources like cache, key_lookup, and bucket_collection
+                
+        Notes:
+            - The Jobs instance provides the interviews to be executed
+            - The environment contains resources like caches and API keys
+            - Additional runtime state like completion status is initialized when run() is called
+        """
+        self.jobs = jobs
+        self.environment = environment
+        # These will be set when run() is called
+        self.start_time = None
+        self.completed = None
+
+    def __len__(self):
+        return len(self.jobs)
+
+    async def run_async(self, parameters: RunParameters) -> 'Results':
+        """
+        Execute interviews asynchronously without progress tracking.
+        
+        This method provides a simplified version of the run method, primarily used
+        by other modules that need direct access to asynchronous execution without
+        the full feature set of the main run() method. This is a lower-level interface
+        that doesn't include progress bars or advanced error handling.
+        
+        Parameters:
+            parameters (RunParameters): Configuration parameters for the run
+            
+        Returns:
+            Results: A Results object containing all responses and metadata
+            
+        Notes:
+            - This method doesn't support progress bars or interactive cancellation
+            - It doesn't handle keyboard interrupts specially
+            - It's primarily meant for internal use by other EDSL components
+            - For most use cases, the main run() method is preferred
+        """
+        # Initialize a simple status tracker (no progress bar)
+        self.environment.jobs_runner_status = JobsRunnerStatus(self, n=parameters.n)
+        data = []
+        task_history = TaskHistory(include_traceback=False)
+
+        run_config = RunConfig(parameters=parameters, environment=self.environment)
+        result_generator = AsyncInterviewRunner(self.jobs, run_config)
+
+        # Process results as they come in
+        async for result, interview in result_generator.run():
+            data.append(result)
+            task_history.add_interview(interview)
+
+        # Create the results object
+        results = Results(survey=self.jobs.survey, task_history=task_history, data=data)
+
+        # Extract only the relevant cache entries
+        relevant_cache = results.relevant_cache(self.environment.cache)
+
+        return Results(
+            survey=self.jobs.survey,
+            task_history=task_history,
+            data=data,
+            cache=relevant_cache,
+        )
+
+    def simple_run(self, parameters: Optional[RunParameters] = None) -> Results:
+        """
+        Run interviews synchronously with minimal configuration.
+        
+        This is a convenience method that provides a very simple synchronous interface
+        for running jobs. It's primarily used for quick tests or debugging, not for 
+        production use.
+        
+        Parameters:
+            parameters (RunParameters, optional): Configuration parameters for the run.
+                If not provided, default parameters will be used.
+                
+        Returns:
+            Results: A Results object containing all responses and metadata
+            
+        Notes:
+            - This method is synchronous (blocks until completion)
+            - It doesn't include progress tracking or advanced error handling
+            - For production use, use the main run() method instead
+        """
+        if parameters is None:
+            parameters = RunParameters()
+            
+        data = asyncio.run(self.run_async(parameters))
+        return Results(survey=self.jobs.survey, data=data)
+
+    @jupyter_nb_handler
+    async def run(self, parameters: RunParameters) -> Results:
+        """
+        Execute interviews asynchronously with full feature support.
+        
+        This is the main method for running jobs with full feature support, including
+        progress tracking, error handling, and graceful cancellation. It's decorated
+        with @jupyter_nb_handler to ensure proper handling in notebook environments.
+        
+        Parameters:
+            parameters (RunParameters): Configuration parameters for the run
+            
+        Returns:
+            Results: A Results object containing all responses and metadata
+            
+        Raises:
+            Exception: Any unhandled exception from interviews if stop_on_exception=True
+            KeyboardInterrupt: If the user interrupts execution and it can't be handled gracefully
+            
+        Notes:
+            - Supports progress bars with remote tracking via Coop
+            - Handles keyboard interrupts gracefully
+            - Manages concurrent execution of multiple interviews
+            - Collects and consolidates results from all completed interviews
+            - Can be used in both async and sync contexts due to the @jupyter_nb_handler decorator
+        """
+
+        run_config = RunConfig(parameters=parameters, environment=self.environment)
+
+        self.start_time = time.monotonic()
+        self.completed = False
+
+        from edsl.coop import Coop
+
+        coop = Coop()
+        endpoint_url = coop.get_progress_bar_url()
+
+        def set_up_jobs_runner_status(jobs_runner_status):
+            if jobs_runner_status is not None:
+                return jobs_runner_status(
+                    self,
+                    n=parameters.n,
+                    endpoint_url=endpoint_url,
+                    job_uuid=parameters.job_uuid,
+                )
+            else:
+                return JobsRunnerStatus(
+                    self,
+                    n=parameters.n,
+                    endpoint_url=endpoint_url,
+                    job_uuid=parameters.job_uuid,
+                )
+
+        run_config.environment.jobs_runner_status = set_up_jobs_runner_status(
+            self.environment.jobs_runner_status
+        )
+
+        async def get_results(results) -> None:
+            """Conducted the interviews and append to the results list."""
+            result_generator = AsyncInterviewRunner(self.jobs, run_config)
+            async for result, interview in result_generator.run():
+                results.append(result)
+                results.task_history.add_interview(interview)
+
+            self.completed = True
+
+        def run_progress_bar(stop_event, jobs_runner_status) -> None:
+            """Runs the progress bar in a separate thread."""
+            jobs_runner_status.update_progress(stop_event)
+
+        def set_up_progress_bar(progress_bar: bool, jobs_runner_status):
+            progress_thread = None
+            if progress_bar and jobs_runner_status.has_ep_api_key():
+                jobs_runner_status.setup()
+                progress_thread = threading.Thread(
+                    target=run_progress_bar, args=(stop_event, jobs_runner_status)
+                )
+                progress_thread.start()
+            elif progress_bar:
+                warnings.warn(
+                    "You need an Expected Parrot API key to view job progress bars."
+                )
+            return progress_thread
+
+        results = Results(
+            survey=self.jobs.survey,
+            data=[],
+            task_history=TaskHistory(),
+            #           cache=self.environment.cache.new_entries_cache(),
+        )
+
+        stop_event = threading.Event()
+        progress_thread = set_up_progress_bar(
+            parameters.progress_bar, run_config.environment.jobs_runner_status
+        )
+
+        exception_to_raise = None
+        try:
+            await get_results(results)
+        except KeyboardInterrupt:
+            print("Keyboard interrupt received. Stopping gracefully...")
+            stop_event.set()
+        except Exception as e:
+            if parameters.stop_on_exception:
+                exception_to_raise = e
+            stop_event.set()
+        finally:
+            stop_event.set()
+            if progress_thread is not None:
+                progress_thread.join()
+
+            if exception_to_raise:
+                raise exception_to_raise
+
+            relevant_cache = results.relevant_cache(self.environment.cache)
+            results.cache = relevant_cache
+            # breakpoint()
+            results.bucket_collection = self.environment.bucket_collection
+
+            from edsl.jobs.results_exceptions_handler import ResultsExceptionsHandler
+
+            results_exceptions_handler = ResultsExceptionsHandler(results, parameters)
+
+            results_exceptions_handler.handle_exceptions()
+            return results

edsl/jobs/{JobsRemoteInferenceHandler.py → remote_inference.py}

@@ -6,20 +6,17 @@ from dataclasses import dataclass
 Seconds = NewType("Seconds", float)
 JobUUID = NewType("JobUUID", str)
 
-from edsl.exceptions.coop import CoopServerResponseError
-
 if TYPE_CHECKING:
-    from edsl.results.Results import Results
-    from edsl.jobs.Jobs import Jobs
-    from edsl.coop.coop import RemoteInferenceResponse, RemoteInferenceCreationInfo
-    from edsl.jobs.JobsRemoteInferenceLogger import JobLogger
-
-from edsl.coop.coop import RemoteInferenceResponse, RemoteInferenceCreationInfo
+    from ..results import Results
+    from .jobs import Jobs
+    from .jobs_remote_inference_logger import JobLogger
 
-from edsl.jobs.jobs_status_enums import JobsStatus
-from edsl.coop.utils import VisibilityType
-from edsl.jobs.JobsRemoteInferenceLogger import JobLogger
+from ..coop import CoopServerResponseError
+from ..coop.utils import VisibilityType
+from ..coop.coop import RemoteInferenceResponse, RemoteInferenceCreationInfo
 
+from .jobs_status_enums import JobsStatus
+from .jobs_remote_inference_logger import JobLogger
 
 class RemoteJobConstants:
     """Constants for remote job handling."""
@@ -54,12 +51,12 @@ class JobsRemoteInferenceHandler:
         self.remote_inference_url = f"{self.expected_parrot_url}/home/remote-inference"
 
     def _create_logger(self) -> JobLogger:
-        from edsl.utilities.is_notebook import is_notebook
-        from edsl.jobs.JobsRemoteInferenceLogger import (
+        from ..utilities import is_notebook
+        from .jobs_remote_inference_logger import (
             JupyterJobLogger,
             StdOutJobLogger,
         )
-        from edsl.jobs.loggers.HTMLTableJobLogger import HTMLTableJobLogger
+        from .html_table_job_logger import HTMLTableJobLogger
 
         if is_notebook():
             return HTMLTableJobLogger(verbose=self.verbose)
@@ -72,7 +69,7 @@ class JobsRemoteInferenceHandler:
             return False
         if not disable_remote_inference:
             try:
-                from edsl.coop.coop import Coop
+                from ..coop import Coop
 
                 user_edsl_settings = Coop().edsl_settings
                 return user_edsl_settings.get("remote_inference", False)
@@ -87,11 +84,11 @@ class JobsRemoteInferenceHandler:
         self,
         iterations: int = 1,
         remote_inference_description: Optional[str] = None,
-        remote_inference_results_visibility: Optional[VisibilityType] = "unlisted",
+        remote_inference_results_visibility: Optional['VisibilityType'] = "unlisted",
         fresh: Optional[bool] = False,
     ) -> RemoteJobInfo:
-        from edsl.config import CONFIG
-        from edsl.coop.coop import Coop
+        from ..config import CONFIG
+        from ..coop import Coop
 
         logger = self._create_logger()
 
@@ -140,8 +137,8 @@ class JobsRemoteInferenceHandler:
     @staticmethod
     def check_status(
         job_uuid: JobUUID,
-    ) -> RemoteInferenceResponse:
-        from edsl.coop.coop import Coop
+    ) -> 'RemoteInferenceResponse':
+        from ..coop import Coop
 
         coop = Coop()
         return coop.remote_inference_get(job_uuid)
@@ -152,7 +149,7 @@ class JobsRemoteInferenceHandler:
         if testing_simulated_response is not None:
             return lambda job_uuid: testing_simulated_response
         else:
-            from edsl.coop.coop import Coop
+            from ..coop import Coop
 
             coop = Coop()
             return coop.remote_inference_get
@@ -164,7 +161,7 @@ class JobsRemoteInferenceHandler:
         if testing_simulated_response is not None:
             return lambda results_uuid, expected_object_type: Results.example()
         else:
-            from edsl.coop.coop import Coop
+            from ..coop import Coop
 
             coop = Coop()
             return coop.get

edsl/jobs/results_exceptions_handler.py

@@ -1,6 +1,6 @@
 from typing import Optional, TYPE_CHECKING, Protocol
 import sys
-from edsl.scenarios.FileStore import HTMLFileStore
+#from edsl.scenarios.FileStore import HTMLFileStore
 from edsl.config import CONFIG
 from edsl.coop.coop import Coop
 
@@ -88,6 +88,6 @@ class ResultsExceptionsHandler:
 
         # Handle remote logging if enabled
         if self.remote_logging:
-            filestore = HTMLFileStore(filepath)
+            filestore = FileStore(filepath)
             coop_details = filestore.push(description="Exceptions Report")
             print(coop_details)

edsl/key_management/__init__.py

@@ -0,0 +1,28 @@
+"""Key management system for API tokens and rate limits.
+
+The key_management module provides a flexible system for managing API keys, credentials, and 
+rate limits for various language model services. It handles discovery, storage, and retrieval
+of API keys from multiple sources including environment variables, configuration files, and
+remote services.
+
+Key components:
+- KeyLookup: Dictionary-like container for service credentials and rate limits
+- KeyLookupBuilder: Factory that builds KeyLookup objects by gathering credentials
+- KeyLookupCollection: Singleton collection to avoid rebuilding KeyLookup objects
+- Data models: Structured representations of API keys, rate limits, and credentials
+
+This module supports multiple credential sources with configurable priority, allowing
+EDSL to use different API keys in different environments while maintaining a consistent
+interface for the rest of the system.
+
+Typical usage:
+```python
+from edsl.key_management import KeyLookupBuilder
+keys = KeyLookupBuilder().build()
+openai_key = keys['openai'].api_token
+```
+"""
+
+from .key_lookup import KeyLookup
+from .key_lookup_collection import KeyLookupCollection
+from .key_lookup_builder import KeyLookupBuilder

edsl/key_management/key_lookup.py

@@ -0,0 +1,161 @@
+from collections import UserDict
+from dataclasses import asdict
+
+from ..enums import service_to_api_keyname
+
+from .models import LanguageModelInput
+
+class KeyLookup(UserDict):
+    """Dictionary-like container for storing and accessing language model service credentials.
+    
+    KeyLookup provides a centralized store for API keys, rate limits, and other configuration
+    needed to authenticate with various language model services. It inherits from UserDict,
+    using service names as keys (e.g., 'openai', 'anthropic') and LanguageModelInput objects
+    as values.
+    
+    The class provides convenient methods for:
+    - Serializing to and from dictionaries for storage
+    - Generating .env files for environment configuration
+    - Creating example instances for testing
+    
+    Typical usage:
+        >>> from edsl.key_management.models import LanguageModelInput
+        >>> lookup = KeyLookup()
+        >>> lookup['openai'] = LanguageModelInput(api_token='sk-key123', rpm=60, tpm=100000)
+        >>> openai_config = lookup['openai']
+        >>> openai_config.api_token
+        'sk-key123'
+    
+    Serialization example:
+        >>> lookup = KeyLookup()
+        >>> lm_input = LanguageModelInput.example()
+        >>> lookup['test'] = lm_input
+        >>> lookup.to_dict()['test']['api_token']
+        'sk-abcd123'
+        >>> restored = KeyLookup.from_dict(lookup.to_dict())
+        >>> restored['test'].api_token
+        'sk-abcd123'
+        
+    Technical Notes:
+        - Uses LanguageModelInput dataclass for structured storage
+        - Preserves source information for debugging and transparency
+        - Supports conversion to environment variable format
+    """
+
+    def to_dict(self):
+        """Convert the KeyLookup to a serializable dictionary.
+        
+        Converts each LanguageModelInput value to a dictionary using dataclasses.asdict,
+        producing a nested dictionary structure suitable for JSON serialization.
+        
+        Returns:
+            dict: A dictionary with service names as keys and serialized LanguageModelInput
+                 objects as values
+        
+        Examples:
+            >>> kl = KeyLookup.example()
+            >>> serialized = kl.to_dict()
+            >>> 'test' in serialized
+            True
+            >>> 'api_token' in serialized['test']
+            True
+            
+            >>> kl2 = KeyLookup.from_dict(kl.to_dict())
+            >>> kl2 == kl  # Equal content
+            True
+            >>> kl2 is kl  # But different objects
+            False
+        """
+        return {k: asdict(v) for k, v in self.data.items()}
+
+    @classmethod
+    def from_dict(cls, d):
+        """Create a KeyLookup instance from a dictionary representation.
+        
+        Converts a dictionary produced by to_dict() back into a KeyLookup instance,
+        reconstructing LanguageModelInput objects from their serialized form.
+        
+        Args:
+            d (dict): Dictionary with service names as keys and serialized 
+                     LanguageModelInput objects as values
+                     
+        Returns:
+            KeyLookup: A new KeyLookup instance populated with the deserialized data
+            
+        Examples:
+            >>> data = {
+            ...     'openai': {
+            ...         'api_token': 'sk-test', 
+            ...         'rpm': 60, 
+            ...         'tpm': 100000
+            ...     }
+            ... }
+            >>> lookup = KeyLookup.from_dict(data)
+            >>> lookup['openai'].api_token
+            'sk-test'
+        """
+        return cls({k: LanguageModelInput(**v) for k, v in d.items()})
+
+    @classmethod
+    def example(cls):
+        """Create an example KeyLookup instance for testing and documentation.
+        
+        Returns:
+            KeyLookup: A new KeyLookup instance with example services and credentials
+            
+        Examples:
+            >>> example = KeyLookup.example()
+            >>> 'test' in example
+            True
+            >>> 'openai' in example
+            True
+        """
+        return cls(
+            {
+                "test": LanguageModelInput.example(),
+                "openai": LanguageModelInput.example(),
+            }
+        )
+
+    def to_dot_env(self):
+        """Generate environment variable definitions for a .env file.
+        
+        Creates a string with environment variable definitions suitable for a .env file,
+        containing service API keys and rate limits in the standard format expected
+        by the key_management system.
+        
+        Returns:
+            str: A string with newline-separated environment variable definitions
+            
+        Examples:
+            >>> lookup = KeyLookup({
+            ...     'test': LanguageModelInput(api_token='test', rpm=10, tpm=20000),
+            ...     'openai': LanguageModelInput(api_token='sk-1234', rpm=60, tpm=100000)
+            ... })
+            >>> env_str = lookup.to_dot_env()
+            >>> 'EDSL_SERVICE_RPM_OPENAI=60' in env_str
+            True
+            >>> 'OPENAI_API_KEY=sk-1234' in env_str
+            True
+            
+        Technical Notes:
+            - Skips the 'test' service which is for internal testing
+            - Handles special cases for service names that don't match their API key names
+            - Includes API IDs for services that require them (e.g., AWS Bedrock)
+        """
+        lines = []
+        for service, lm_input in self.items():
+            if service != "test":
+                lines.append(f"EDSL_SERVICE_RPM_{service.upper()}={lm_input.rpm}")
+                lines.append(f"EDSL_SERVICE_TPM_{service.upper()}={lm_input.tpm}")
+                key_name = service_to_api_keyname.get(service, service)
+                lines.append(f"{key_name.upper()}={lm_input.api_token}")
+                if lm_input.api_id is not None:
+                    lines.append(f"{service.upper()}_API_ID={lm_input.api_id}")
+        return "\n".join([f"{line}" for line in lines])
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)