edsl 0.1.39.dev3__py3-none-any.whl → 0.1.39.dev5__py3-none-any.whl

edsl/inference_services/InferenceServicesCollection.py

@@ -1,97 +1,138 @@
-from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
-import warnings
-
-
-class InferenceServicesCollection:
-    added_models = {}
-
-    def __init__(self, services: list[InferenceServiceABC] = None):
-        self.services = services or []
-
-    @classmethod
-    def add_model(cls, service_name, model_name):
-        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])
-
-            for model in self.added_models.get(service._inference_service_, []):
-                total_models.append([model, service._inference_service_, -1])
-
-        sorted_models = sorted(total_models)
-        for i, model in enumerate(sorted_models):
-            model[2] = i
-            model = tuple(model)
-        return sorted_models
-
-    def register(self, service):
-        self.services.append(service)
-
-    def create_model_factory(self, model_name: str, service_name=None, index=None):
-        from edsl.inference_services.TestService import TestService
-
-        if model_name == "test":
-            return TestService.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)
-
-        raise Exception(f"Model {model_name} not found in any of the services")
+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
+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. 
+                                    If you know the service that has this model, use the service_name parameter directly.
+                                    E.g., Model("gpt-4o", service_name="openai")
+                                    """
+        )
+
+
+class InferenceServicesCollection:
+    added_models = defaultdict(list)  # Moved back to class level
+
+    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: str, model_name: str) -> None:
+        if service_name not in cls.added_models:
+            cls.added_models[service_name].append(model_name)
+
+    def service_names_to_classes(self) -> Dict[str, InferenceServiceABC]:
+        return {service._inference_service_: service for service in self.services}
+
+    def available(
+        self,
+        service: Optional[str] = None,
+    ) -> List[Tuple[str, str, int]]:
+        return self.availability_fetcher.available(service)
+
+    def reset_cache(self) -> None:
+        self.availability_fetcher.reset_cache()
+
+    @property
+    def num_cache_entries(self) -> int:
+        return self.availability_fetcher.num_cache_entries
+
+    def register(self, service: InferenceServiceABC) -> None:
+        self.services.append(service)
+
+    def create_model_factory(
+        self, model_name: str, service_name: Optional[InferenceServiceLiteral] = None
+    ) -> "LanguageModel":
+
+        if service_name is None:  # we try to find the right service
+            service = self.resolver.resolve_model(model_name, service_name)
+        else:  # if they passed a service, we'll use that
+            service = self.service_names_to_classes().get(service_name)
+
+        if not service:  # but if we can't find it, we'll raise an error
+            raise InferenceServiceError(f"Service {service_name} not found")
+
+        return service.create_model(model_name)
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

edsl/inference_services/MistralAIService.py

@@ -1,123 +1,120 @@
-import os
-from typing import Any, List, Optional
-from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
-from edsl.language_models.LanguageModel import LanguageModel
-import asyncio
-from mistralai import Mistral
-
-from edsl.exceptions.language_models import LanguageModelBadResponseError
-
-
-class MistralAIService(InferenceServiceABC):
-    """Mistral AI service class."""
-
-    key_sequence = ["choices", 0, "message", "content"]
-    usage_sequence = ["usage"]
-
-    _inference_service_ = "mistral"
-    _env_key_name_ = "MISTRAL_API_KEY"  # Environment variable for Mistral API key
-    input_token_name = "prompt_tokens"
-    output_token_name = "completion_tokens"
-
-    _sync_client_instance = None
-    _async_client_instance = None
-
-    _sync_client = Mistral
-    _async_client = Mistral
-
-    _models_list_cache: List[str] = []
-    model_exclude_list = []
-
-    def __init_subclass__(cls, **kwargs):
-        super().__init_subclass__(**kwargs)
-        # so subclasses have to create their own instances of the clients
-        cls._sync_client_instance = None
-        cls._async_client_instance = None
-
-    @classmethod
-    def sync_client(cls):
-        if cls._sync_client_instance is None:
-            cls._sync_client_instance = cls._sync_client(
-                api_key=os.getenv(cls._env_key_name_)
-            )
-        return cls._sync_client_instance
-
-    @classmethod
-    def async_client(cls):
-        if cls._async_client_instance is None:
-            cls._async_client_instance = cls._async_client(
-                api_key=os.getenv(cls._env_key_name_)
-            )
-        return cls._async_client_instance
-
-    @classmethod
-    def available(cls) -> list[str]:
-        if not cls._models_list_cache:
-            cls._models_list_cache = [
-                m.id for m in cls.sync_client().models.list().data
-            ]
-
-        return cls._models_list_cache
-
-    @classmethod
-    def create_model(
-        cls, model_name: str = "mistral", model_class_name=None
-    ) -> LanguageModel:
-        if model_class_name is None:
-            model_class_name = cls.to_class_name(model_name)
-
-        class LLM(LanguageModel):
-            """
-            Child class of LanguageModel for interacting with Mistral models.
-            """
-
-            key_sequence = cls.key_sequence
-            usage_sequence = cls.usage_sequence
-
-            input_token_name = cls.input_token_name
-            output_token_name = cls.output_token_name
-
-            _inference_service_ = cls._inference_service_
-            _model_ = model_name
-            _parameters_ = {
-                "temperature": 0.5,
-                "max_tokens": 512,
-                "top_p": 0.9,
-            }
-
-            _tpm = cls.get_tpm(cls)
-            _rpm = cls.get_rpm(cls)
-
-            def sync_client(self):
-                return cls.sync_client()
-
-            def async_client(self):
-                return cls.async_client()
-
-            async def async_execute_model_call(
-                self,
-                user_prompt: str,
-                system_prompt: str = "",
-                files_list: Optional[List["FileStore"]] = None,
-            ) -> dict[str, Any]:
-                """Calls the Mistral API and returns the API response."""
-                s = self.async_client()
-
-                try:
-                    res = await s.chat.complete_async(
-                        model=model_name,
-                        messages=[
-                            {
-                                "content": user_prompt,
-                                "role": "user",
-                            },
-                        ],
-                    )
-                except Exception as e:
-                    raise LanguageModelBadResponseError(f"Error with Mistral API: {e}")
-
-                return res.model_dump()
-
-        LLM.__name__ = model_class_name
-
-        return LLM
+import os
+from typing import Any, List, Optional
+from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
+from edsl.language_models.LanguageModel import LanguageModel
+import asyncio
+from mistralai import Mistral
+
+from edsl.exceptions.language_models import LanguageModelBadResponseError
+
+
+class MistralAIService(InferenceServiceABC):
+    """Mistral AI service class."""
+
+    key_sequence = ["choices", 0, "message", "content"]
+    usage_sequence = ["usage"]
+
+    _inference_service_ = "mistral"
+    _env_key_name_ = "MISTRAL_API_KEY"  # Environment variable for Mistral API key
+    input_token_name = "prompt_tokens"
+    output_token_name = "completion_tokens"
+
+    _sync_client_instance = None
+    _async_client_instance = None
+
+    _sync_client = Mistral
+    _async_client = Mistral
+
+    _models_list_cache: List[str] = []
+    model_exclude_list = []
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        # so subclasses have to create their own instances of the clients
+        cls._sync_client_instance = None
+        cls._async_client_instance = None
+
+    @classmethod
+    def sync_client(cls):
+        if cls._sync_client_instance is None:
+            cls._sync_client_instance = cls._sync_client(
+                api_key=os.getenv(cls._env_key_name_)
+            )
+        return cls._sync_client_instance
+
+    @classmethod
+    def async_client(cls):
+        if cls._async_client_instance is None:
+            cls._async_client_instance = cls._async_client(
+                api_key=os.getenv(cls._env_key_name_)
+            )
+        return cls._async_client_instance
+
+    @classmethod
+    def available(cls) -> list[str]:
+        if not cls._models_list_cache:
+            cls._models_list_cache = [
+                m.id for m in cls.sync_client().models.list().data
+            ]
+
+        return cls._models_list_cache
+
+    @classmethod
+    def create_model(
+        cls, model_name: str = "mistral", model_class_name=None
+    ) -> LanguageModel:
+        if model_class_name is None:
+            model_class_name = cls.to_class_name(model_name)
+
+        class LLM(LanguageModel):
+            """
+            Child class of LanguageModel for interacting with Mistral models.
+            """
+
+            key_sequence = cls.key_sequence
+            usage_sequence = cls.usage_sequence
+
+            input_token_name = cls.input_token_name
+            output_token_name = cls.output_token_name
+
+            _inference_service_ = cls._inference_service_
+            _model_ = model_name
+            _parameters_ = {
+                "temperature": 0.5,
+                "max_tokens": 512,
+                "top_p": 0.9,
+            }
+
+            def sync_client(self):
+                return cls.sync_client()
+
+            def async_client(self):
+                return cls.async_client()
+
+            async def async_execute_model_call(
+                self,
+                user_prompt: str,
+                system_prompt: str = "",
+                files_list: Optional[List["FileStore"]] = None,
+            ) -> dict[str, Any]:
+                """Calls the Mistral API and returns the API response."""
+                s = self.async_client()
+
+                try:
+                    res = await s.chat.complete_async(
+                        model=model_name,
+                        messages=[
+                            {
+                                "content": user_prompt,
+                                "role": "user",
+                            },
+                        ],
+                    )
+                except Exception as e:
+                    raise LanguageModelBadResponseError(f"Error with Mistral API: {e}")
+
+                return res.model_dump()
+
+        LLM.__name__ = model_class_name
+
+        return LLM

edsl/inference_services/OllamaService.py

@@ -1,18 +1,18 @@
-import aiohttp
-import json
-import requests
-from typing import Any, List
-
-# from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
-from edsl.language_models import LanguageModel
-
-from edsl.inference_services.OpenAIService import OpenAIService
-
-
-class OllamaService(OpenAIService):
-    """DeepInfra service class."""
-
-    _inference_service_ = "ollama"
-    _env_key_name_ = "DEEP_INFRA_API_KEY"
-    _base_url_ = "http://localhost:11434/v1"
-    _models_list_cache: List[str] = []
+import aiohttp
+import json
+import requests
+from typing import Any, List
+
+# from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
+from edsl.language_models import LanguageModel
+
+from edsl.inference_services.OpenAIService import OpenAIService
+
+
+class OllamaService(OpenAIService):
+    """DeepInfra service class."""
+
+    _inference_service_ = "ollama"
+    _env_key_name_ = "DEEP_INFRA_API_KEY"
+    _base_url_ = "http://localhost:11434/v1"
+    _models_list_cache: List[str] = []