edsl 0.1.39.dev1__py3-none-any.whl → 0.1.39.dev2__py3-none-any.whl

edsl/inference_services/InferenceServiceABC.py

@@ -1,5 +1,4 @@
 from abc import abstractmethod, ABC
-import os
 import re
 from datetime import datetime, timedelta
 from edsl.config import CONFIG
@@ -8,31 +7,32 @@ from edsl.config import CONFIG
 class InferenceServiceABC(ABC):
     """
     Abstract class for inference services.
-    Anthropic: https://docs.anthropic.com/en/api/rate-limits
     """
 
     _coop_config_vars = None
 
-    default_levels = {
-        "google": {"tpm": 2_000_000, "rpm": 15},
-        "openai": {"tpm": 2_000_000, "rpm": 10_000},
-        "anthropic": {"tpm": 2_000_000, "rpm": 500},
-    }
-
     def __init_subclass__(cls):
         """
         Check that the subclass has the required attributes.
         - `key_sequence` attribute determines...
         - `model_exclude_list` attribute determines...
         """
-        if not hasattr(cls, "key_sequence"):
-            raise NotImplementedError(
-                f"Class {cls.__name__} must have a 'key_sequence' attribute."
-            )
-        if not hasattr(cls, "model_exclude_list"):
-            raise NotImplementedError(
-                f"Class {cls.__name__} must have a 'model_exclude_list' attribute."
-            )
+        must_have_attributes = [
+            "key_sequence",
+            "model_exclude_list",
+            "usage_sequence",
+            "input_token_name",
+            "output_token_name",
+        ]
+        for attr in must_have_attributes:
+            if not hasattr(cls, attr):
+                raise NotImplementedError(
+                    f"Class {cls.__name__} must have a '{attr}' attribute."
+                )
+
+    @property
+    def service_name(self):
+        return self._inference_service_
 
     @classmethod
     def _should_refresh_coop_config_vars(cls):
@@ -44,44 +44,6 @@ class InferenceServiceABC(ABC):
             return True
         return (datetime.now() - cls._last_config_fetch) > timedelta(hours=24)
 
-    @classmethod
-    def _get_limt(cls, limit_type: str) -> int:
-        key = f"EDSL_SERVICE_{limit_type.upper()}_{cls._inference_service_.upper()}"
-        if key in os.environ:
-            return int(os.getenv(key))
-
-        if cls._coop_config_vars is None or cls._should_refresh_coop_config_vars():
-            try:
-                from edsl import Coop
-
-                c = Coop()
-                cls._coop_config_vars = c.fetch_rate_limit_config_vars()
-                cls._last_config_fetch = datetime.now()
-                if key in cls._coop_config_vars:
-                    return cls._coop_config_vars[key]
-            except Exception:
-                cls._coop_config_vars = None
-        else:
-            if key in cls._coop_config_vars:
-                return cls._coop_config_vars[key]
-
-        if cls._inference_service_ in cls.default_levels:
-            return int(cls.default_levels[cls._inference_service_][limit_type])
-
-        return int(CONFIG.get(f"EDSL_SERVICE_{limit_type.upper()}_BASELINE"))
-
-    def get_tpm(cls) -> int:
-        """
-        Returns the TPM for the service. If the service is not defined in the environment variables, it will return the baseline TPM.
-        """
-        return cls._get_limt(limit_type="tpm")
-
-    def get_rpm(cls):
-        """
-        Returns the RPM for the service. If the service is not defined in the environment variables, it will return the baseline RPM.
-        """
-        return cls._get_limt(limit_type="rpm")
-
     @abstractmethod
     def available() -> list[str]:
         """
@@ -113,35 +75,6 @@ class InferenceServiceABC(ABC):
 
 
 if __name__ == "__main__":
-    pass
-    # deep_infra_service = DeepInfraService("deep_infra", "DEEP_INFRA_API_KEY")
-    # deep_infra_service.available()
-    # m = deep_infra_service.create_model("microsoft/WizardLM-2-7B")
-    # response = m().hello()
-    # print(response)
-
-    # anthropic_service = AnthropicService("anthropic", "ANTHROPIC_API_KEY")
-    # anthropic_service.available()
-    # m = anthropic_service.create_model("claude-3-opus-20240229")
-    # response = m().hello()
-    # print(response)
-    # factory = OpenAIService("openai", "OPENAI_API")
-    # factory.available()
-    # m = factory.create_model("gpt-3.5-turbo")
-    # response = m().hello()
-
-    # from edsl import QuestionFreeText
-    # results = QuestionFreeText.example().by(m()).run()
-
-    # collection = InferenceServicesCollection([
-    #     OpenAIService,
-    #     AnthropicService,
-    #     DeepInfraService
-    # ])
+    import doctest
 
-    # available = collection.available()
-    # factory = collection.create_model_factory(*available[0])
-    # m = factory()
-    # from edsl import QuestionFreeText
-    # results = QuestionFreeText.example().by(m).run()
-    # print(results)
+    doctest.testmod()

edsl/inference_services/InferenceServicesCollection.py

@@ -1,97 +1,122 @@
+from functools import lru_cache
+from collections import defaultdict
+from typing import Optional, Protocol, Dict, List, Tuple, TYPE_CHECKING, Literal
+
 from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
-import warnings
+from edsl.inference_services.AvailableModelFetcher import AvailableModelFetcher
+from edsl.exceptions.inference_services import InferenceServiceError
+
+if TYPE_CHECKING:
+    from edsl.language_models.LanguageModel import LanguageModel
+    from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
+
+
+class ModelCreator(Protocol):
+    def create_model(self, model_name: str) -> "LanguageModel":
+        ...
+
+
+from edsl.enums import InferenceServiceLiteral
+
+
+class ModelResolver:
+    def __init__(
+        self,
+        services: List[InferenceServiceLiteral],
+        models_to_services: Dict[InferenceServiceLiteral, InferenceServiceABC],
+        availability_fetcher: "AvailableModelFetcher",
+    ):
+        """
+        Class for determining which service to use for a given model.
+        """
+        self.services = services
+        self._models_to_services = models_to_services
+        self.availability_fetcher = availability_fetcher
+        self._service_names_to_classes = {
+            service._inference_service_: service for service in services
+        }
+
+    def resolve_model(
+        self, model_name: str, service_name: Optional[InferenceServiceLiteral] = None
+    ) -> InferenceServiceABC:
+        """Returns an InferenceServiceABC object for the given model name.
+
+        :param model_name: The name of the model to resolve. E.g., 'gpt-4o'
+        :param service_name: The name of the service to use. E.g., 'openai'
+        :return: An InferenceServiceABC object
+
+        """
+        if model_name == "test":
+            from edsl.inference_services.TestService import TestService
+
+            return TestService()
+
+        if service_name is not None:
+            service: InferenceServiceABC = self._service_names_to_classes.get(
+                service_name
+            )
+            if not service:
+                raise InferenceServiceError(f"Service {service_name} not found")
+            return service
+
+        if model_name in self._models_to_services:  # maybe we've seen it before!
+            return self._models_to_services[model_name]
+
+        for service in self.services:
+            (
+                available_models,
+                service_name,
+            ) = self.availability_fetcher.get_available_models_by_service(service)
+            if model_name in available_models:
+                self._models_to_services[model_name] = service
+                return service
+
+        raise InferenceServiceError(f"Model {model_name} not found in any services")
 
 
 class InferenceServicesCollection:
-    added_models = {}
+    added_models = defaultdict(list)  # Moved back to class level
 
-    def __init__(self, services: list[InferenceServiceABC] = None):
+    def __init__(self, services: Optional[List[InferenceServiceABC]] = None):
         self.services = services or []
+        self._models_to_services: Dict[str, InferenceServiceABC] = {}
+
+        self.availability_fetcher = AvailableModelFetcher(
+            self.services, self.added_models
+        )
+        self.resolver = ModelResolver(
+            self.services, self._models_to_services, self.availability_fetcher
+        )
 
     @classmethod
-    def add_model(cls, service_name, model_name):
+    def add_model(cls, service_name: str, model_name: str) -> None:
         if service_name not in cls.added_models:
-            cls.added_models[service_name] = []
-        cls.added_models[service_name].append(model_name)
-
-    @staticmethod
-    def _get_service_available(service, warn: bool = False) -> list[str]:
-        try:
-            service_models = service.available()
-        except Exception:
-            if warn:
-                warnings.warn(
-                    f"""Error getting models for {service._inference_service_}. 
-                    Check that you have properly stored your Expected Parrot API key and activated remote inference, or stored your own API keys for the language models that you want to use.
-                    See https://docs.expectedparrot.com/en/latest/api_keys.html for instructions on storing API keys.
-                    Relying on Coop.""",
-                    UserWarning,
-                )
-
-            # Use the list of models on Coop as a fallback
-            try:
-                from edsl import Coop
-
-                c = Coop()
-                models_from_coop = c.fetch_models()
-                service_models = models_from_coop.get(service._inference_service_, [])
-
-                # cache results
-                service._models_list_cache = service_models
-
-            # Finally, use the available models cache from the Python file
-            except Exception:
-                if warn:
-                    warnings.warn(
-                        f"""Error getting models for {service._inference_service_}. 
-                        Relying on EDSL cache.""",
-                        UserWarning,
-                    )
-
-                from edsl.inference_services.models_available_cache import (
-                    models_available,
-                )
-
-                service_models = models_available.get(service._inference_service_, [])
-
-                # cache results
-                service._models_list_cache = service_models
-
-        return service_models
-
-    def available(self):
-        total_models = []
-        for service in self.services:
-            service_models = self._get_service_available(service)
-            for model in service_models:
-                total_models.append([model, service._inference_service_, -1])
+            cls.added_models[service_name].append(model_name)
 
-            for model in self.added_models.get(service._inference_service_, []):
-                total_models.append([model, service._inference_service_, -1])
+    def available(
+        self,
+        service: Optional[str] = None,
+    ) -> List[Tuple[str, str, int]]:
+        return self.availability_fetcher.available(service)
 
-        sorted_models = sorted(total_models)
-        for i, model in enumerate(sorted_models):
-            model[2] = i
-            model = tuple(model)
-        return sorted_models
+    def reset_cache(self) -> None:
+        self.availability_fetcher.reset_cache()
 
-    def register(self, service):
-        self.services.append(service)
+    @property
+    def num_cache_entries(self) -> int:
+        return self.availability_fetcher.num_cache_entries
 
-    def create_model_factory(self, model_name: str, service_name=None, index=None):
-        from edsl.inference_services.TestService import TestService
+    def register(self, service: InferenceServiceABC) -> None:
+        self.services.append(service)
 
-        if model_name == "test":
-            return TestService.create_model(model_name)
+    def create_model_factory(
+        self, model_name: str, service_name: Optional[InferenceServiceLiteral] = None
+    ) -> "LanguageModel":
+        service = self.resolver.resolve_model(model_name, service_name)
+        return service.create_model(model_name)
 
-        if service_name:
-            for service in self.services:
-                if service_name == service._inference_service_:
-                    return service.create_model(model_name)
 
-        for service in self.services:
-            if model_name in self._get_service_available(service):
-                if service_name is None or service_name == service._inference_service_:
-                    return service.create_model(model_name)
+if __name__ == "__main__":
+    import doctest
 
-        raise Exception(f"Model {model_name} not found in any of the services")
+    doctest.testmod()

edsl/inference_services/MistralAIService.py

@@ -85,9 +85,6 @@ class MistralAIService(InferenceServiceABC):
                 "top_p": 0.9,
             }
 
-            _tpm = cls.get_tpm(cls)
-            _rpm = cls.get_rpm(cls)
-
             def sync_client(self):
                 return cls.sync_client()
 

edsl/inference_services/OpenAIService.py

@@ -5,7 +5,7 @@ import os
 import openai
 
 from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
-from edsl.language_models import LanguageModel
+from edsl.language_models.LanguageModel import LanguageModel
 from edsl.inference_services.rate_limits_cache import rate_limits
 from edsl.utilities.utilities import fix_partial_correct_response
 
@@ -107,9 +107,6 @@ class OpenAIService(InferenceServiceABC):
             input_token_name = cls.input_token_name
             output_token_name = cls.output_token_name
 
-            _rpm = cls.get_rpm(cls)
-            _tpm = cls.get_tpm(cls)
-
             _inference_service_ = cls._inference_service_
             _model_ = model_name
             _parameters_ = {

edsl/inference_services/PerplexityService.py

@@ -51,9 +51,6 @@ class PerplexityService(OpenAIService):
             input_token_name = cls.input_token_name
             output_token_name = cls.output_token_name
 
-            _rpm = cls.get_rpm(cls)
-            _tpm = cls.get_tpm(cls)
-
             _inference_service_ = cls._inference_service_
             _model_ = model_name
 

edsl/inference_services/ServiceAvailability.py

@@ -0,0 +1,135 @@
+from enum import Enum
+from typing import List, Optional, TYPE_CHECKING
+from functools import partial
+import warnings
+
+from edsl.inference_services.data_structures import AvailableModels, ModelNamesList
+
+if TYPE_CHECKING:
+    from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
+
+
+class ModelSource(Enum):
+    LOCAL = "local"
+    COOP = "coop"
+    CACHE = "cache"
+
+
+class ServiceAvailability:
+    """This class is responsible for fetching the available models from different sources."""
+
+    _coop_model_list = None
+
+    def __init__(self, source_order: Optional[List[ModelSource]] = None):
+        """
+        Initialize with custom source order.
+        Default order is LOCAL -> COOP -> CACHE
+        """
+        self.source_order = source_order or [
+            ModelSource.LOCAL,
+            ModelSource.COOP,
+            ModelSource.CACHE,
+        ]
+
+        # Map sources to their fetch functions
+        self._source_fetchers = {
+            ModelSource.LOCAL: self._fetch_from_local_service,
+            ModelSource.COOP: self._fetch_from_coop,
+            ModelSource.CACHE: self._fetch_from_cache,
+        }
+
+    @classmethod
+    def models_from_coop(cls) -> AvailableModels:
+        if not cls._coop_model_list:
+            from edsl.coop.coop import Coop
+
+            c = Coop()
+            coop_model_list = c.fetch_models()
+            cls._coop_model_list = coop_model_list
+        return cls._coop_model_list
+
+    def get_service_available(
+        self, service: "InferenceServiceABC", warn: bool = False
+    ) -> ModelNamesList:
+        """
+        Try to fetch available models from sources in specified order.
+        Returns first successful result.
+        """
+        last_error = None
+
+        for source in self.source_order:
+            try:
+                fetch_func = partial(self._source_fetchers[source], service)
+                result = fetch_func()
+
+                # Cache successful result
+                service._models_list_cache = result
+                return result
+
+            except Exception as e:
+                last_error = e
+                if warn:
+                    self._warn_source_failed(service, source)
+                continue
+
+        # If we get here, all sources failed
+        raise RuntimeError(
+            f"All sources failed to fetch models. Last error: {last_error}"
+        )
+
+    @staticmethod
+    def _fetch_from_local_service(service: "InferenceServiceABC") -> ModelNamesList:
+        """Attempt to fetch models directly from the service."""
+        return service.available()
+
+    @classmethod
+    def _fetch_from_coop(cls, service: "InferenceServiceABC") -> ModelNamesList:
+        """Fetch models from Coop."""
+        models_from_coop = cls.models_from_coop()
+        return models_from_coop.get(service._inference_service_, [])
+
+    @staticmethod
+    def _fetch_from_cache(service: "InferenceServiceABC") -> ModelNamesList:
+        """Fetch models from local cache."""
+        from edsl.inference_services.models_available_cache import models_available
+
+        return models_available.get(service._inference_service_, [])
+
+    def _warn_source_failed(self, service: "InferenceServiceABC", source: ModelSource):
+        """Display appropriate warning message based on failed source."""
+        messages = {
+            ModelSource.LOCAL: f"""Error getting models for {service._inference_service_}. 
+                Check that you have properly stored your Expected Parrot API key and activated remote inference, 
+                or stored your own API keys for the language models that you want to use.
+                See https://docs.expectedparrot.com/en/latest/api_keys.html for instructions on storing API keys.
+                Trying next source.""",
+            ModelSource.COOP: f"Error getting models from Coop for {service._inference_service_}. Trying next source.",
+            ModelSource.CACHE: f"Error getting models from cache for {service._inference_service_}.",
+        }
+        warnings.warn(messages[source], UserWarning)
+
+
+if __name__ == "__main__":
+    # sa = ServiceAvailability()
+    # models_from_coop = sa.models_from_coop()
+    # print(models_from_coop)
+    from edsl.inference_services.OpenAIService import OpenAIService
+
+    openai_models = ServiceAvailability._fetch_from_local_service(OpenAIService())
+    print(openai_models)
+
+# Example usage:
+"""
+# Default order (LOCAL -> COOP -> CACHE)
+availability = ServiceAvailability()
+
+# Custom order (COOP -> LOCAL -> CACHE)
+availability_coop_first = ServiceAvailability([
+    ModelSource.COOP,
+    ModelSource.LOCAL,
+    ModelSource.CACHE
+])
+
+# Get available models using custom order
+models = availability_coop_first.get_service_available(service, warn=True)
+"""

edsl/inference_services/TestService.py

@@ -2,7 +2,7 @@ from typing import Any, List, Optional
 import os
 import asyncio
 from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
-from edsl.language_models import LanguageModel
+from edsl.language_models.LanguageModel import LanguageModel
 from edsl.inference_services.rate_limits_cache import rate_limits
 from edsl.utilities.utilities import fix_partial_correct_response
 
@@ -65,13 +65,7 @@ class TestService(InferenceServiceABC):
                 await asyncio.sleep(0.1)
                 # return {"message": """{"answer": "Hello, world"}"""}
 
-                if hasattr(self, "func"):
-                    return {
-                        "message": [
-                            {"text": self.func(user_prompt, system_prompt, files_list)}
-                        ],
-                        "usage": {"prompt_tokens": 1, "completion_tokens": 1},
-                    }
+                # breakpoint()
 
                 if hasattr(self, "throw_exception") and self.throw_exception:
                     if hasattr(self, "exception_probability"):
@@ -81,6 +75,15 @@ class TestService(InferenceServiceABC):
 
                     if random.random() < p:
                         raise Exception("This is a test error")
+
+                if hasattr(self, "func"):
+                    return {
+                        "message": [
+                            {"text": self.func(user_prompt, system_prompt, files_list)}
+                        ],
+                        "usage": {"prompt_tokens": 1, "completion_tokens": 1},
+                    }
+
                 return {
                     "message": [{"text": f"{self._canned_response}"}],
                     "usage": {"prompt_tokens": 1, "completion_tokens": 1},

edsl/inference_services/data_structures.py

@@ -0,0 +1,62 @@
+from collections import UserDict, defaultdict, UserList
+from typing import Union
+from edsl.enums import InferenceServiceLiteral
+from dataclasses import dataclass
+
+
+@dataclass
+class LanguageModelInfo:
+    model_name: str
+    service_name: str
+
+    def __getitem__(self, key: int) -> str:
+        import warnings
+
+        warnings.warn(
+            "Accessing LanguageModelInfo via index is deprecated. "
+            "Please use .model_name, .service_name, or .index attributes instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+        if key == 0:
+            return self.model_name
+        elif key == 1:
+            return self.service_name
+        else:
+            raise IndexError("Index out of range")
+
+
+class ModelNamesList(UserList):
+    pass
+
+
+class AvailableModels(UserList):
+    def __init__(self, data: list) -> None:
+        super().__init__(data)
+
+    def __contains__(self, model_name: str) -> bool:
+        for model_entry in self:
+            if model_entry.model_name == model_name:
+                return True
+        return False
+
+
+class ServiceToModelsMapping(UserDict):
+    def __init__(self, data: dict) -> None:
+        super().__init__(data)
+
+    @property
+    def service_names(self) -> list[str]:
+        return list(self.data.keys())
+
+    def _validate_service_names(self):
+        for service in self.service_names:
+            if service not in InferenceServiceLiteral:
+                raise ValueError(f"Invalid service name: {service}")
+
+    def model_to_services(self) -> dict:
+        self._model_to_service = defaultdict(list)
+        for service, models in self.data.items():
+            for model in models:
+                self._model_to_service[model].append(service)