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

edsl/inference_services/InferenceServicesCollection.py

@@ -1,122 +1,97 @@
-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")
-
-
-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 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":
-        service = self.resolver.resolve_model(model_name, service_name)
-        return service.create_model(model_name)
-
-
-if __name__ == "__main__":
-    import doctest
-
-    doctest.testmod()
+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")

edsl/inference_services/MistralAIService.py

@@ -1,120 +1,123 @@
-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
+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

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] = []