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

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,11 +1,35 @@
 from abc import abstractmethod, ABC
 from typing import Any
 import re
+from edsl.config import CONFIG
 
 
 class InferenceServiceABC(ABC):
     """Abstract class for inference services."""
 
+    # check if child class has cls attribute "key_sequence"
+    def __init_subclass__(cls):
+        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."
+            )
+
+    def get_tpm(cls):
+        key = f"EDSL_SERVICE_TPM_{cls._inference_service_.upper()}"
+        if key not in CONFIG:
+            key = "EDSL_SERVICE_TPM_BASELINE"
+        return int(CONFIG.get(key))
+
+    def get_rpm(cls):
+        key = f"EDSL_SERVICE_RPM_{cls._inference_service_.upper()}"
+        if key not in CONFIG:
+            key = "EDSL_SERVICE_RPM_BASELINE"
+        return int(CONFIG.get(key))
+
     @abstractmethod
     def available() -> list[str]:
         pass

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,10 +186,6 @@ 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()
                 params = {
                     "model": self.model,
@@ -195,24 +204,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

edsl/inference_services/TestService.py

@@ -0,0 +1,71 @@
+from typing import Any, List
+import os
+import asyncio
+from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
+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.enums import InferenceServiceType
+
+
+class TestService(InferenceServiceABC):
+    """OpenAI service class."""
+
+    key_sequence = None
+    model_exclude_list = []
+    _inference_service_ = "test"
+    input_token_name = "prompt_tokens"
+    output_token_name = "completion_tokens"
+
+    @classmethod
+    def available(cls) -> list[str]:
+        return ["test"]
+
+    @classmethod
+    def create_model(cls, model_name, model_class_name=None) -> LanguageModel:
+        throw_exception = False
+
+        class TestServiceLanguageModel(LanguageModel):
+            _model_ = "test"
+            _parameters_ = {"temperature": 0.5}
+            _inference_service_ = InferenceServiceType.TEST.value
+            usage_sequence = ["usage"]
+            key_sequence = ["message", 0, "text"]
+            input_token_name = cls.input_token_name
+            output_token_name = cls.output_token_name
+            _rpm = 1000
+            _tpm = 100000
+
+            @property
+            def _canned_response(self):
+                if hasattr(self, "canned_response"):
+                    return self.canned_response
+                else:
+                    return "Hello, world"
+
+            async def async_execute_model_call(
+                self, user_prompt: str, system_prompt: str
+            ) -> dict[str, Any]:
+                await asyncio.sleep(0.1)
+                # return {"message": """{"answer": "Hello, world"}"""}
+                if hasattr(self, "throw_exception") and self.throw_exception:
+                    raise Exception("This is a test error")
+                return {
+                    "message": [{"text": f"{self._canned_response}"}],
+                    "usage": {"prompt_tokens": 1, "completion_tokens": 1},
+                }
+
+        return TestServiceLanguageModel
+
+    # _inference_service_ = "openai"
+    # _env_key_name_ = "OPENAI_API_KEY"
+    # _base_url_ = None
+
+    # _sync_client_ = openai.OpenAI
+    # _async_client_ = openai.AsyncOpenAI
+
+    # _sync_client_instance = None
+    # _async_client_instance = None
+
+    # key_sequence = ["choices", 0, "message", "content"]

edsl/inference_services/models_available_cache.py

@@ -70,12 +70,6 @@ models_available = {
         "amazon.titan-tg1-large",
         "amazon.titan-text-lite-v1",
         "amazon.titan-text-express-v1",
-        "ai21.j2-grande-instruct",
-        "ai21.j2-jumbo-instruct",
-        "ai21.j2-mid",
-        "ai21.j2-mid-v1",
-        "ai21.j2-ultra",
-        "ai21.j2-ultra-v1",
         "anthropic.claude-instant-v1",
         "anthropic.claude-v2:1",
         "anthropic.claude-v2",

edsl/inference_services/registry.py

@@ -10,6 +10,8 @@ from edsl.inference_services.GroqService import GroqService
 from edsl.inference_services.AwsBedrock import AwsBedrockService
 from edsl.inference_services.AzureAI import AzureAIService
 from edsl.inference_services.OllamaService import OllamaService
+from edsl.inference_services.TestService import TestService
+from edsl.inference_services.MistralAIService import MistralAIService
 
 default = InferenceServicesCollection(
     [
@@ -21,5 +23,7 @@ default = InferenceServicesCollection(
         AwsBedrockService,
         AzureAIService,
         OllamaService,
+        TestService,
+        MistralAIService,
     ]
 )

edsl/jobs/Answers.py

@@ -2,24 +2,22 @@
 
 from collections import UserDict
 from rich.table import Table
+from edsl.data_transfer_models import EDSLResultObjectInput
 
 
 class Answers(UserDict):
     """Helper class to hold the answers to a survey."""
 
-    def add_answer(self, response, question) -> None:
-        """Add a response to the answers dictionary.
-
-        >>> from edsl import QuestionFreeText
-        >>> q = QuestionFreeText.example()
-        >>> answers = Answers()
-        >>> answers.add_answer({"answer": "yes"}, q)
-        >>> answers[q.question_name]
-        'yes'
-        """
-        answer = response.get("answer")
-        comment = response.pop("comment", None)
+    def add_answer(
+        self, response: EDSLResultObjectInput, question: "QuestionBase"
+    ) -> None:
+        """Add a response to the answers dictionary."""
+        answer = response.answer
+        comment = response.comment
+        generated_tokens = response.generated_tokens
         # record the answer
+        if generated_tokens:
+            self[question.question_name + "_generated_tokens"] = generated_tokens
         self[question.question_name] = answer
         if comment:
             self[question.question_name + "_comment"] = comment

edsl/jobs/FailedQuestion.py

@@ -0,0 +1,78 @@
+from edsl.questions import QuestionBase
+from edsl import Question, Scenario, Model, Agent
+
+from edsl.language_models.LanguageModel import LanguageModel
+
+
+class FailedQuestion:
+    # tests/jobs/test_Interview.py::test_handle_model_exceptions
+
+    # (Pdb) dir(self.exception.__traceback__)
+    # ['tb_frame', 'tb_lasti', 'tb_lineno', 'tb_next']
+
+    def __init__(
+        self, question, scenario, model, agent, raw_model_response, exception, prompts
+    ):
+        self.question = question
+        self.scenario = scenario
+        self.model = model
+        self.agent = agent
+        self.raw_model_response = raw_model_response  # JSON
+        self.exception = exception
+        self.prompts = prompts
+
+    def to_dict(self):
+        return {
+            "question": self.question._to_dict(),
+            "scenario": self.scenario._to_dict(),
+            "model": self.model._to_dict(),
+            "agent": self.agent._to_dict(),
+            "raw_model_response": self.raw_model_response,
+            "exception": self.exception.__class__.__name__,  # self.exception,
+            "prompts": self.prompts,
+        }
+
+    @classmethod
+    def from_dict(cls, data):
+        question = QuestionBase.from_dict(data["question"])
+        scenario = Scenario.from_dict(data["scenario"])
+        model = LanguageModel.from_dict(data["model"])
+        agent = Agent.from_dict(data["agent"])
+        raw_model_response = data["raw_model_response"]
+        exception = data["exception"]
+        prompts = data["prompts"]
+        return cls(
+            question, scenario, model, agent, raw_model_response, exception, prompts
+        )
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(question={repr(self.question)}, scenario={repr(self.scenario)}, model={repr(self.model)}, agent={repr(self.agent)}, raw_model_response={repr(self.raw_model_response)}, exception={repr(self.exception)})"
+
+    @property
+    def jobs(self):
+        return self.question.by(self.scenario).by(self.agent).by(self.model)
+
+    def rerun(self):
+        results = self.jobs.run()
+        return results
+
+    def help(self):
+        pass
+
+    @classmethod
+    def example(cls):
+        from edsl.language_models.utilities import create_language_model
+        from edsl.language_models.utilities import create_survey
+
+        survey = create_survey(2, chained=False, take_scenario=False)
+        fail_at_number = 1
+        model = create_language_model(ValueError, fail_at_number)()
+        from edsl import Survey
+
+        results = survey.by(model).run()
+        return results.failed_questions[0][0]
+
+
+if __name__ == "__main__":
+    fq = FailedQuestion.example()
+    new_fq = FailedQuestion.from_dict(fq.to_dict())