edsl 0.1.32__py3-none-any.whl → 0.1.33__py3-none-any.whl

edsl/inference_services/AnthropicService.py

@@ -11,6 +11,11 @@ class AnthropicService(InferenceServiceABC):
 
     _inference_service_ = "anthropic"
     _env_key_name_ = "ANTHROPIC_API_KEY"
+    key_sequence = ["content", 0, "text"]  # ["content"][0]["text"]
+    usage_sequence = ["usage"]
+    input_token_name = "input_tokens"
+    output_token_name = "output_tokens"
+    model_exclude_list = []
 
     @classmethod
     def available(cls):
@@ -34,6 +39,11 @@ class AnthropicService(InferenceServiceABC):
             Child class of LanguageModel for interacting with OpenAI 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_ = {
@@ -46,6 +56,9 @@ class AnthropicService(InferenceServiceABC):
                 "top_logprobs": 3,
             }
 
+            _tpm = cls.get_tpm(cls)
+            _rpm = cls.get_rpm(cls)
+
             async def async_execute_model_call(
                 self, user_prompt: str, system_prompt: str = ""
             ) -> dict[str, Any]:
@@ -66,17 +79,6 @@ class AnthropicService(InferenceServiceABC):
                 )
                 return response.model_dump()
 
-            @staticmethod
-            def parse_response(raw_response: dict[str, Any]) -> str:
-                """Parses the API response and returns the response text."""
-                response = raw_response["content"][0]["text"]
-                pattern = r"^```json(?:\\n|\n)(.+?)(?:\\n|\n)```$"
-                match = re.match(pattern, response, re.DOTALL)
-                if match:
-                    return match.group(1)
-                else:
-                    return response
-
         LLM.__name__ = model_class_name
 
         return LLM

edsl/inference_services/AwsBedrock.py

@@ -16,6 +16,18 @@ class AwsBedrockService(InferenceServiceABC):
     _env_key_name_ = (
         "AWS_ACCESS_KEY_ID"  # or any other environment key for AWS credentials
     )
+    key_sequence = ["output", "message", "content", 0, "text"]
+    input_token_name = "inputTokens"
+    output_token_name = "outputTokens"
+    usage_sequence = ["usage"]
+    model_exclude_list = [
+        "ai21.j2-grande-instruct",
+        "ai21.j2-jumbo-instruct",
+        "ai21.j2-mid",
+        "ai21.j2-mid-v1",
+        "ai21.j2-ultra",
+        "ai21.j2-ultra-v1",
+    ]
 
     @classmethod
     def available(cls):
@@ -28,7 +40,7 @@ class AwsBedrockService(InferenceServiceABC):
         else:
             all_models_ids = cls._models_list_cache
 
-        return all_models_ids
+        return [m for m in all_models_ids if m not in cls.model_exclude_list]
 
     @classmethod
     def create_model(
@@ -42,6 +54,8 @@ class AwsBedrockService(InferenceServiceABC):
             Child class of LanguageModel for interacting with AWS Bedrock models.
             """
 
+            key_sequence = cls.key_sequence
+            usage_sequence = cls.usage_sequence
             _inference_service_ = cls._inference_service_
             _model_ = model_name
             _parameters_ = {
@@ -49,6 +63,10 @@ class AwsBedrockService(InferenceServiceABC):
                 "max_tokens": 512,
                 "top_p": 0.9,
             }
+            input_token_name = cls.input_token_name
+            output_token_name = cls.output_token_name
+            _rpm = cls.get_rpm(cls)
+            _tpm = cls.get_tpm(cls)
 
             async def async_execute_model_call(
                 self, user_prompt: str, system_prompt: str = ""
@@ -89,22 +107,6 @@ class AwsBedrockService(InferenceServiceABC):
                     print(e)
                     return {"error": str(e)}
 
-            @staticmethod
-            def parse_response(raw_response: dict[str, Any]) -> str:
-                """Parses the API response and returns the response text."""
-                if "output" in raw_response and "message" in raw_response["output"]:
-                    response = raw_response["output"]["message"]["content"][0]["text"]
-                    pattern = r"^```json(?:\\n|\n)(.+?)(?:\\n|\n)```$"
-                    match = re.match(pattern, response, re.DOTALL)
-                    if match:
-                        return match.group(1)
-                    else:
-                        out = fix_partial_correct_response(response)
-                        if "error" not in out:
-                            response = out["extracted_json"]
-                        return response
-                return "Error parsing response"
-
         LLM.__name__ = model_class_name
 
         return LLM

edsl/inference_services/AzureAI.py

@@ -25,11 +25,22 @@ def json_handle_none(value: Any) -> Any:
 class AzureAIService(InferenceServiceABC):
     """Azure AI service class."""
 
+    # key_sequence = ["content", 0, "text"]  # ["content"][0]["text"]
+    key_sequence = ["choices", 0, "message", "content"]
+    usage_sequence = ["usage"]
+    input_token_name = "prompt_tokens"
+    output_token_name = "completion_tokens"
+
     _inference_service_ = "azure"
     _env_key_name_ = (
         "AZURE_ENDPOINT_URL_AND_KEY"  # Environment variable for Azure API key
     )
     _model_id_to_endpoint_and_key = {}
+    model_exclude_list = [
+        "Cohere-command-r-plus-xncmg",
+        "Mistral-Nemo-klfsi",
+        "Mistral-large-2407-ojfld",
+    ]
 
     @classmethod
     def available(cls):
@@ -82,7 +93,7 @@ class AzureAIService(InferenceServiceABC):
 
             except Exception as e:
                 raise e
-        return out
+        return [m for m in out if m not in cls.model_exclude_list]
 
     @classmethod
     def create_model(
@@ -96,6 +107,10 @@ class AzureAIService(InferenceServiceABC):
             Child class of LanguageModel for interacting with Azure OpenAI 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_ = {
@@ -103,6 +118,8 @@ class AzureAIService(InferenceServiceABC):
                 "max_tokens": 512,
                 "top_p": 0.9,
             }
+            _rpm = cls.get_rpm(cls)
+            _tpm = cls.get_tpm(cls)
 
             async def async_execute_model_call(
                 self, user_prompt: str, system_prompt: str = ""
@@ -172,25 +189,25 @@ class AzureAIService(InferenceServiceABC):
                     )
                     return response.model_dump()
 
-            @staticmethod
-            def parse_response(raw_response: dict[str, Any]) -> str:
-                """Parses the API response and returns the response text."""
-                if (
-                    raw_response
-                    and "choices" in raw_response
-                    and raw_response["choices"]
-                ):
-                    response = raw_response["choices"][0]["message"]["content"]
-                    pattern = r"^```json(?:\\n|\n)(.+?)(?:\\n|\n)```$"
-                    match = re.match(pattern, response, re.DOTALL)
-                    if match:
-                        return match.group(1)
-                    else:
-                        out = fix_partial_correct_response(response)
-                        if "error" not in out:
-                            response = out["extracted_json"]
-                        return response
-                return "Error parsing response"
+            # @staticmethod
+            # def parse_response(raw_response: dict[str, Any]) -> str:
+            #     """Parses the API response and returns the response text."""
+            #     if (
+            #         raw_response
+            #         and "choices" in raw_response
+            #         and raw_response["choices"]
+            #     ):
+            #         response = raw_response["choices"][0]["message"]["content"]
+            #         pattern = r"^```json(?:\\n|\n)(.+?)(?:\\n|\n)```$"
+            #         match = re.match(pattern, response, re.DOTALL)
+            #         if match:
+            #             return match.group(1)
+            #         else:
+            #             out = fix_partial_correct_response(response)
+            #             if "error" not in out:
+            #                 response = out["extracted_json"]
+            #             return response
+            #     return "Error parsing response"
 
         LLM.__name__ = model_class_name
 

edsl/inference_services/GoogleService.py

@@ -10,10 +10,16 @@ from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
 
 class GoogleService(InferenceServiceABC):
     _inference_service_ = "google"
+    key_sequence = ["candidates", 0, "content", "parts", 0, "text"]
+    usage_sequence = ["usageMetadata"]
+    input_token_name = "promptTokenCount"
+    output_token_name = "candidatesTokenCount"
+
+    model_exclude_list = []
 
     @classmethod
     def available(cls):
-        return ["gemini-pro"]
+        return ["gemini-pro", "gemini-1.5-pro", "gemini-1.5-flash", "gemini-1.0-pro"]
 
     @classmethod
     def create_model(
@@ -24,7 +30,15 @@ class GoogleService(InferenceServiceABC):
 
         class LLM(LanguageModel):
             _model_ = model_name
+            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_
+
+            _tpm = cls.get_tpm(cls)
+            _rpm = cls.get_rpm(cls)
+
             _parameters_ = {
                 "temperature": 0.5,
                 "topP": 1,
@@ -50,7 +64,7 @@ class GoogleService(InferenceServiceABC):
                         "stopSequences": self.stopSequences,
                     },
                 }
-
+                # print(combined_prompt)
                 async with aiohttp.ClientSession() as session:
                     async with session.post(
                         url, headers=headers, data=json.dumps(data)
@@ -58,16 +72,6 @@ class GoogleService(InferenceServiceABC):
                         raw_response_text = await response.text()
                         return json.loads(raw_response_text)
 
-            def parse_response(self, raw_response: dict[str, Any]) -> str:
-                data = raw_response
-                try:
-                    return data["candidates"][0]["content"]["parts"][0]["text"]
-                except KeyError as e:
-                    print(
-                        f"The data return was {data}, which was missing the key 'candidates'"
-                    )
-                    raise e
-
         LLM.__name__ = model_name
 
         return LLM

edsl/inference_services/GroqService.py

@@ -13,6 +13,8 @@ class GroqService(OpenAIService):
     _sync_client_ = groq.Groq
     _async_client_ = groq.AsyncGroq
 
+    model_exclude_list = ["whisper-large-v3", "distil-whisper-large-v3-en"]
+
     # _base_url_ = "https://api.deepinfra.com/v1/openai"
     _base_url_ = None
     _models_list_cache: List[str] = []

edsl/inference_services/InferenceServiceABC.py

@@ -1,22 +1,77 @@
 from abc import abstractmethod, ABC
-from typing import Any
+import os
 import re
+from edsl.config import CONFIG
 
 
 class InferenceServiceABC(ABC):
-    """Abstract class for inference services."""
+    """
+    Abstract class for inference services.
+    Anthropic: https://docs.anthropic.com/en/api/rate-limits
+    """
+
+    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."
+            )
+
+    @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._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]:
+        """
+        Returns a list of available models for the service.
+        """
         pass
 
     @abstractmethod
     def create_model():
+        """
+        Returns a LanguageModel object.
+        """
         pass
 
     @staticmethod
     def to_class_name(s):
-        """Convert a string to a valid class name.
+        """
+        Converts a string to a valid class name.
 
         >>> InferenceServiceABC.to_class_name("hello world")
         'HelloWorld'

edsl/inference_services/MistralAIService.py

@@ -0,0 +1,120 @@
+import os
+from typing import Any, List
+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 = ""
+            ) -> 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/OpenAIService.py

@@ -1,8 +1,7 @@
-from typing import Any, List
-import re
+from __future__ import annotations
+from typing import Any, List, Optional
 import os
 
-# from openai import AsyncOpenAI
 import openai
 
 from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
@@ -10,6 +9,8 @@ from edsl.language_models import LanguageModel
 from edsl.inference_services.rate_limits_cache import rate_limits
 from edsl.utilities.utilities import fix_partial_correct_response
 
+from edsl.config import CONFIG
+
 
 class OpenAIService(InferenceServiceABC):
     """OpenAI service class."""
@@ -21,19 +22,36 @@ class OpenAIService(InferenceServiceABC):
     _sync_client_ = openai.OpenAI
     _async_client_ = openai.AsyncOpenAI
 
+    _sync_client_instance = None
+    _async_client_instance = None
+
+    key_sequence = ["choices", 0, "message", "content"]
+    usage_sequence = ["usage"]
+    input_token_name = "prompt_tokens"
+    output_token_name = "completion_tokens"
+
+    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):
-        return cls._sync_client_(
-            api_key=os.getenv(cls._env_key_name_), base_url=cls._base_url_
-        )
+        if cls._sync_client_instance is None:
+            cls._sync_client_instance = cls._sync_client_(
+                api_key=os.getenv(cls._env_key_name_), base_url=cls._base_url_
+            )
+        return cls._sync_client_instance
 
     @classmethod
     def async_client(cls):
-        return cls._async_client_(
-            api_key=os.getenv(cls._env_key_name_), base_url=cls._base_url_
-        )
+        if cls._async_client_instance is None:
+            cls._async_client_instance = cls._async_client_(
+                api_key=os.getenv(cls._env_key_name_), base_url=cls._base_url_
+            )
+        return cls._async_client_instance
 
-    # TODO: Make this a coop call
     model_exclude_list = [
         "whisper-1",
         "davinci-002",
@@ -48,6 +66,8 @@ class OpenAIService(InferenceServiceABC):
         "text-embedding-3-small",
         "text-embedding-ada-002",
         "ft:davinci-002:mit-horton-lab::8OfuHgoo",
+        "gpt-3.5-turbo-instruct-0914",
+        "gpt-3.5-turbo-instruct",
     ]
     _models_list_cache: List[str] = []
 
@@ -61,11 +81,8 @@ class OpenAIService(InferenceServiceABC):
 
     @classmethod
     def available(cls) -> List[str]:
-        # from openai import OpenAI
-
         if not cls._models_list_cache:
             try:
-                # client = OpenAI(api_key = os.getenv(cls._env_key_name_), base_url = cls._base_url_)
                 cls._models_list_cache = [
                     m.id
                     for m in cls.get_model_list()
@@ -73,15 +90,6 @@ class OpenAIService(InferenceServiceABC):
                 ]
             except Exception as e:
                 raise
-                # print(
-                #     f"""Error retrieving models: {e}.
-                #     See instructions about storing your API keys: https://docs.expectedparrot.com/en/latest/api_keys.html"""
-                # )
-                # cls._models_list_cache = [
-                #     "gpt-3.5-turbo",
-                #     "gpt-4-1106-preview",
-                #     "gpt-4",
-                # ]  # Fallback list
         return cls._models_list_cache
 
     @classmethod
@@ -94,6 +102,14 @@ class OpenAIService(InferenceServiceABC):
             Child class of LanguageModel for interacting with OpenAI 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
+
+            _rpm = cls.get_rpm(cls)
+            _tpm = cls.get_tpm(cls)
+
             _inference_service_ = cls._inference_service_
             _model_ = model_name
             _parameters_ = {
@@ -114,15 +130,9 @@ class OpenAIService(InferenceServiceABC):
 
             @classmethod
             def available(cls) -> list[str]:
-                # import openai
-                # client = openai.OpenAI(api_key = os.getenv(cls._env_key_name_), base_url = cls._base_url_)
-                # return client.models.list()
                 return cls.sync_client().models.list()
 
             def get_headers(self) -> dict[str, Any]:
-                # from openai import OpenAI
-
-                # client = OpenAI(api_key = os.getenv(cls._env_key_name_), base_url = cls._base_url_)
                 client = self.sync_client()
                 response = client.chat.completions.with_raw_response.create(
                     messages=[
@@ -159,6 +169,9 @@ class OpenAIService(InferenceServiceABC):
                 user_prompt: str,
                 system_prompt: str = "",
                 encoded_image=None,
+                invigilator: Optional[
+                    "InvigilatorAI"
+                ] = None,  # TBD - can eventually be used for function-calling
             ) -> dict[str, Any]:
                 """Calls the OpenAI API and returns the API response."""
                 if encoded_image:
@@ -173,17 +186,16 @@ class OpenAIService(InferenceServiceABC):
                     )
                 else:
                     content = user_prompt
-                # self.client = AsyncOpenAI(
-                #     api_key = os.getenv(cls._env_key_name_),
-                #     base_url = cls._base_url_
-                #     )
                 client = self.async_client()
+                messages = [
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": content},
+                ]
+                if system_prompt == "" and self.omit_system_prompt_if_empty:
+                    messages = messages[1:]
                 params = {
                     "model": self.model,
-                    "messages": [
-                        {"role": "system", "content": system_prompt},
-                        {"role": "user", "content": content},
-                    ],
+                    "messages": messages,
                     "temperature": self.temperature,
                     "max_tokens": self.max_tokens,
                     "top_p": self.top_p,
@@ -195,24 +207,6 @@ class OpenAIService(InferenceServiceABC):
                 response = await client.chat.completions.create(**params)
                 return response.model_dump()
 
-            @staticmethod
-            def parse_response(raw_response: dict[str, Any]) -> str:
-                """Parses the API response and returns the response text."""
-                try:
-                    response = raw_response["choices"][0]["message"]["content"]
-                except KeyError:
-                    print("Tried to parse response but failed:")
-                    print(raw_response)
-                pattern = r"^```json(?:\\n|\n)(.+?)(?:\\n|\n)```$"
-                match = re.match(pattern, response, re.DOTALL)
-                if match:
-                    return match.group(1)
-                else:
-                    out = fix_partial_correct_response(response)
-                    if "error" not in out:
-                        response = out["extracted_json"]
-                    return response
-
         LLM.__name__ = "LanguageModel"
 
         return LLM