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

edsl/jobs/remote_inference.py

@@ -1,6 +1,11 @@
 from typing import Optional, Union, Literal, TYPE_CHECKING, NewType, Callable, Any
-
 from dataclasses import dataclass
+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
+from .exceptions import RemoteInferenceError
 
 
 Seconds = NewType("Seconds", float)
@@ -9,14 +14,7 @@ JobUUID = NewType("JobUUID", str)
 if TYPE_CHECKING:
     from ..results import Results
     from .jobs import Jobs
-    from .jobs_remote_inference_logger 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."""
@@ -53,7 +51,6 @@ class JobsRemoteInferenceHandler:
     def _create_logger(self) -> JobLogger:
         from ..utilities import is_notebook
         from .jobs_remote_inference_logger import (
-            JupyterJobLogger,
             StdOutJobLogger,
         )
         from .html_table_job_logger import HTMLTableJobLogger
@@ -75,7 +72,7 @@ class JobsRemoteInferenceHandler:
                 return user_edsl_settings.get("remote_inference", False)
             except requests.ConnectionError:
                 pass
-            except CoopServerResponseError as e:
+            except CoopServerResponseError:
                 pass
 
         return False
@@ -84,10 +81,9 @@ 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 ..config import CONFIG
         from ..coop import Coop
 
         logger = self._create_logger()
@@ -137,7 +133,7 @@ class JobsRemoteInferenceHandler:
     @staticmethod
     def check_status(
         job_uuid: JobUUID,
-    ) -> 'RemoteInferenceResponse':
+    ) -> "RemoteInferenceResponse":
         from ..coop import Coop
 
         coop = Coop()
@@ -182,6 +178,13 @@ class JobsRemoteInferenceHandler:
     ) -> None:
         "Handles a failed job by logging the error and updating the job status."
         latest_error_report_url = remote_job_data.get("latest_error_report_url")
+
+        reason = remote_job_data.get("reason")
+
+        if reason == "insufficient funds":
+            latest_error_report_url = "Error: Insufficient balance to start the job"
+            print("❌ Error: Insufficient balance to start the job")
+
         if latest_error_report_url:
             job_info.logger.add_info("error_report_url", latest_error_report_url)
 
@@ -235,13 +238,13 @@ class JobsRemoteInferenceHandler:
         """Makes one attempt to fetch and process a remote job's status and results."""
         remote_job_data = remote_job_data_fetcher(job_info.job_uuid)
         status = remote_job_data.get("status")
-
+        reason = remote_job_data.get("reason")
         if status == "cancelled":
             self._handle_cancelled_job(job_info)
-            return None
+            return None, reason
 
-        elif status == "failed" or status == "completed":
-            if status == "failed":
+        elif status == "failed" or status == "completed" or status == "partial_failed":
+            if status == "failed" or status == "partial_failed":
                 self._handle_failed_job(job_info, remote_job_data)
 
             results_uuid = remote_job_data.get("results_uuid")
@@ -252,13 +255,13 @@ class JobsRemoteInferenceHandler:
                     remote_job_data=remote_job_data,
                     object_fetcher=object_fetcher,
                 )
-                return results
+                return results, reason
             else:
-                return None
+                return None, reason
 
         else:
             self._sleep_for_a_bit(job_info, status)
-            return "continue"
+            return "continue", reason
 
     def poll_remote_inference_job(
         self,
@@ -274,11 +277,11 @@ class JobsRemoteInferenceHandler:
 
         job_in_queue = True
         while job_in_queue:
-            result = self._attempt_fetch_job(
+            result, reason = self._attempt_fetch_job(
                 job_info, remote_job_data_fetcher, object_fetcher
             )
             if result != "continue":
-                return result
+                return result, reason
 
     async def create_and_poll_remote_job(
         self,
@@ -310,7 +313,7 @@ class JobsRemoteInferenceHandler:
             ),
         )
         if job_info is None:
-            raise ValueError("Remote job creation failed.")
+            raise RemoteInferenceError("Remote job creation failed.")
 
         return await loop.run_in_executor(
             None,

edsl/jobs/results_exceptions_handler.py

@@ -1,8 +1,10 @@
-from typing import Optional, TYPE_CHECKING, Protocol
+from typing import Protocol
 import sys
 #from edsl.scenarios.FileStore import HTMLFileStore
 from edsl.config import CONFIG
 from edsl.coop.coop import Coop
+from ..scenarios import FileStore
+from .exceptions import JobsErrors
 
 
 class ResultsProtocol(Protocol):
@@ -41,7 +43,8 @@ class ResultsExceptionsHandler:
         self.parameters = parameters
 
         self.open_in_browser = self._get_browser_setting()
-        self.remote_logging = self._get_remote_logging_setting()
+        #self.remote_logging = self._get_remote_logging_setting()
+        self.remote_logging = False
 
     def _get_browser_setting(self) -> bool:
         """Determine if exceptions should be opened in browser based on config."""
@@ -51,7 +54,7 @@ class ResultsExceptionsHandler:
         elif setting == "False":
             return False
         else:
-            raise Exception(
+            raise JobsErrors(
                 "EDSL_OPEN_EXCEPTION_REPORT_URL must be either True or False"
             )
 
@@ -60,13 +63,13 @@ class ResultsExceptionsHandler:
         try:
             coop = Coop()
             return coop.edsl_settings["remote_logging"]
-        except Exception as e:
+        except Exception:
             # print(e)
             return False
 
     def _generate_error_message(self, indices) -> str:
         """Generate appropriate error message based on number of exceptions."""
-        msg = f"Exceptions were raised.\n" 
+        msg = "Exceptions were raised.\n" 
         return msg
 
     def handle_exceptions(self) -> None:

edsl/key_management/__init__.py

@@ -25,4 +25,6 @@ openai_key = keys['openai'].api_token
 
 from .key_lookup import KeyLookup
 from .key_lookup_collection import KeyLookupCollection
-from .key_lookup_builder import KeyLookupBuilder
+from .key_lookup_builder import KeyLookupBuilder
+
+__all__ = ["KeyLookup", "KeyLookupCollection", "KeyLookupBuilder"]

edsl/key_management/exceptions.py

@@ -0,0 +1,62 @@
+"""
+Exceptions specific to the key_management module.
+
+This module defines custom exception classes for all key management-related errors
+in the EDSL framework, ensuring consistent error handling and user feedback.
+"""
+
+from ..base import BaseException
+
+
+class KeyManagementError(BaseException):
+    """
+    Base exception class for all key management-related errors.
+    
+    This is the parent class for all exceptions related to API key management,
+    including key lookup, validation, and configuration.
+    
+    Examples:
+        ```python
+        # Usually not raised directly, but through subclasses
+        # For example, when attempting to use an invalid key
+        key_lookup.add_key(service="invalid_service", key="api_key")  # Would raise KeyManagementServiceError
+        ```
+    """
+    relevant_doc = "https://docs.expectedparrot.com/en/latest/api_keys.html"
+
+
+class KeyManagementValueError(KeyManagementError):
+    """
+    Exception raised when invalid values are provided to key management operations.
+    
+    This exception occurs when:
+    - Invalid fetch orders or configuration parameters are provided
+    - API key format is invalid
+    - Service names don't match expected values
+    
+    Examples:
+        ```python
+        # Invalid fetch order
+        key_lookup_builder.set_fetch_order("invalid_order")  # Raises KeyManagementValueError
+        ```
+    """
+    relevant_doc = "https://docs.expectedparrot.com/en/latest/api_keys.html"
+
+
+class KeyManagementDuplicateError(KeyManagementError):
+    """
+    Exception raised when duplicate keys or services are encountered.
+    
+    This exception occurs when:
+    - Attempting to add a duplicate service ID
+    - Adding a key that already exists for a service
+    - Registering multiple handlers for the same service
+    
+    Examples:
+        ```python
+        # Adding a duplicate service
+        builder.register_env_var(service_id="openai", env_var="OPENAI_API_KEY")
+        builder.register_env_var(service_id="openai", env_var="ANOTHER_KEY")  # Raises KeyManagementDuplicateError
+        ```
+    """
+    relevant_doc = "https://docs.expectedparrot.com/en/latest/api_keys.html"

edsl/key_management/key_lookup.py

@@ -19,7 +19,7 @@ class KeyLookup(UserDict):
     - Creating example instances for testing
     
     Typical usage:
-        >>> from edsl.key_management.models import LanguageModelInput
+        >>> from .models import LanguageModelInput  # Import for doctest
         >>> lookup = KeyLookup()
         >>> lookup['openai'] = LanguageModelInput(api_token='sk-key123', rpm=60, tpm=100000)
         >>> openai_config = lookup['openai']

edsl/key_management/key_lookup_builder.py

@@ -1,9 +1,25 @@
-from typing import Optional, List
+from typing import Optional, TYPE_CHECKING
 import os
 from functools import lru_cache
+import textwrap
+
+if TYPE_CHECKING:
+    from ..coop import Coop
 
 from ..enums import service_to_api_keyname
-from ..exceptions.general import MissingAPIKeyError
+from ..base import BaseException
+
+class MissingAPIKeyError(BaseException):
+    def __init__(self, full_message=None, model_name=None, inference_service=None, silent=False):
+        if model_name and inference_service:
+            full_message = textwrap.dedent(
+                f"""
+                An API Key for model `{model_name}` is missing from the .env file.
+                This key is associated with the inference service `{inference_service}`.
+                Please see https://docs.expectedparrot.com/en/latest/api_keys.html for more information.
+                """
+            )
+        super().__init__(full_message, show_docs=False, silent=silent)
 
 from .key_lookup import KeyLookup
 from .models import (
@@ -31,6 +47,8 @@ api_id_to_service = {"AWS_ACCESS_KEY_ID": "bedrock"}
 class KeyLookupBuilder:
     """Factory class for building KeyLookup objects by gathering credentials from multiple sources.
     
+    >>> from edsl.key_management.exceptions import KeyManagementValueError
+    
     KeyLookupBuilder is responsible for discovering, organizing, and consolidating API keys
     and rate limits from various sources. It can pull credentials from:
     
@@ -67,9 +85,9 @@ class KeyLookupBuilder:
     Validation examples:
         >>> try:
         ...     KeyLookupBuilder(fetch_order=["config", "env"])  # Should be tuple
-        ... except ValueError as e:
-        ...     str(e)
-        'fetch_order must be a tuple'
+        ... except KeyManagementValueError as e:
+        ...     "fetch_order must be a tuple" in str(e)
+        True
         
         >>> builder = KeyLookupBuilder()
         >>> builder.extract_service("EDSL_SERVICE_RPM_OPENAI")
@@ -93,7 +111,8 @@ class KeyLookupBuilder:
         fetch_order: Optional[tuple[str]] = None,
         coop: Optional["Coop"] = None,
     ):
-        from ..coop import Coop
+        # Import here to avoid circular import issues
+        from ..coop import Coop  # Import Coop type for type hinting
 
         # Fetch order goes from lowest priority to highest priority
         if fetch_order is None:
@@ -102,7 +121,8 @@ class KeyLookupBuilder:
             self.fetch_order = fetch_order
 
         if not isinstance(self.fetch_order, tuple):
-            raise ValueError("fetch_order must be a tuple")
+            from edsl.key_management.exceptions import KeyManagementValueError
+            raise KeyManagementValueError("fetch_order must be a tuple")
 
         if coop is None:
             self.coop = Coop()
@@ -141,8 +161,8 @@ class KeyLookupBuilder:
             >>> lookup = builder.build()
             >>> isinstance(lookup, KeyLookup)
             True
-            >>> lookup['test'].api_token  # Test service should always exist
-            'test'
+            >>> lookup['test'].api_token == 'test'  # Test service should always exist
+            True
             
         Technical Notes:
             - Skips services with missing API keys
@@ -194,7 +214,7 @@ class KeyLookupBuilder:
             - Supports services that require both API key and API ID
         """
         if (key_entries := self.key_data.get(service)) is None:
-            raise MissingAPIKeyError(f"No key found for service '{service}'")
+            raise MissingAPIKeyError(f"No key found for service '{service}'", silent=True)
 
         if len(key_entries) == 1:
             api_key_entry = key_entries[0]
@@ -286,6 +306,8 @@ class KeyLookupBuilder:
 
     def _add_id(self, key: str, value: str, source: str) -> None:
         """Add an API ID to the id_data dictionary.
+        
+        >>> from edsl.key_management.exceptions import KeyManagementDuplicateError
 
         >>> builder = KeyLookupBuilder()
         >>> builder._add_id("AWS_ACCESS_KEY_ID", "AKIA1234", "env")
@@ -293,9 +315,9 @@ class KeyLookupBuilder:
         'AKIA1234'
         >>> try:
         ...     builder._add_id("AWS_ACCESS_KEY_ID", "AKIA5678", "env")
-        ... except ValueError as e:
-        ...     str(e)
-        'Duplicate ID for service bedrock'
+        ... except KeyManagementDuplicateError as e:
+        ...     "Duplicate ID for service bedrock" in str(e)
+        True
         """
         service = api_id_to_service[key]
         if service not in self.id_data:
@@ -303,7 +325,8 @@ class KeyLookupBuilder:
                 service=service, name=key, value=value, source=source
             )
         else:
-            raise ValueError(f"Duplicate ID for service {service}")
+            from edsl.key_management.exceptions import KeyManagementDuplicateError
+            raise KeyManagementDuplicateError(f"Duplicate ID for service {service}")
 
     def _add_limit(self, key: str, value: str, source: str) -> None:
         """Add a rate limit entry to the limit_data dictionary.

edsl/key_management/key_lookup_collection.py

@@ -1,5 +1,6 @@
 from collections import UserDict
 
+# Import for doctest and type hints
 from .key_lookup_builder import KeyLookupBuilder
 from .key_lookup import KeyLookup
 
@@ -25,6 +26,7 @@ class KeyLookupCollection(UserDict):
         True
         
     Basic usage:
+        >>> from edsl.key_management import KeyLookup
         >>> collection = KeyLookupCollection()
         >>> collection.add_key_lookup(("config", "env"))
         >>> lookup = collection[("config", "env")]  # Get the stored KeyLookup

edsl/language_models/__init__.py

@@ -4,4 +4,4 @@ from .model_list import ModelList
 
 from .exceptions import LanguageModelBadResponseError
 
-__all__ = ["Model", "ModelList", "LanguageModelBadResponseError"]
+__all__ = ["Model", "ModelList", "LanguageModelBadResponseError", "LanguageModel"]