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

edsl/language_models/LanguageModel.py

@@ -1,16 +1,4 @@
-"""This module contains the LanguageModel class, which is an abstract base class for all language models.
-
-Terminology:
-
-raw_response: The JSON response from the model. This has all the model meta-data about the call. 
-
-edsl_augmented_response: The JSON response from model, but augmented with EDSL-specific information, 
-such as the cache key, token usage, etc. 
-
-generated_tokens: The actual tokens generated by the model. This is the output that is used by the user.
-edsl_answer_dict: The parsed JSON response from the model either {'answer': ...} or {'answer': ..., 'comment': ...}
-
-"""
+"""This module contains the LanguageModel class, which is an abstract base class for all language models."""
 
 from __future__ import annotations
 import warnings
@@ -20,103 +8,47 @@ import json
 import time
 import os
 import hashlib
-from typing import (
-    Coroutine,
-    Any,
-    Callable,
-    Type,
-    Union,
-    List,
-    get_type_hints,
-    TypedDict,
-    Optional,
-)
+from typing import Coroutine, Any, Callable, Type, List, get_type_hints
 from abc import ABC, abstractmethod
 
-from json_repair import repair_json
 
-from edsl.data_transfer_models import (
-    ModelResponse,
-    ModelInputs,
-    EDSLOutput,
-    AgentResponseDict,
-)
+class IntendedModelCallOutcome:
+    "This is a tuple-like class that holds the response, cache_used, and cache_key."
+
+    def __init__(self, response: dict, cache_used: bool, cache_key: str):
+        self.response = response
+        self.cache_used = cache_used
+        self.cache_key = cache_key
+
+    def __iter__(self):
+        """Iterate over the class attributes.
+
+        >>> a, b, c = IntendedModelCallOutcome({'answer': "yes"}, True, 'x1289')
+        >>> a
+        {'answer': 'yes'}
+        """
+        yield self.response
+        yield self.cache_used
+        yield self.cache_key
+
+    def __len__(self):
+        return 3
+
+    def __repr__(self):
+        return f"IntendedModelCallOutcome(response = {self.response}, cache_used = {self.cache_used}, cache_key = '{self.cache_key}')"
 
 
 from edsl.config import CONFIG
+
 from edsl.utilities.decorators import sync_wrapper, jupyter_nb_handler
 from edsl.utilities.decorators import add_edsl_version, remove_edsl_version
+
 from edsl.language_models.repair import repair
 from edsl.enums import InferenceServiceType
 from edsl.Base import RichPrintingMixin, PersistenceMixin
 from edsl.enums import service_to_api_keyname
 from edsl.exceptions import MissingAPIKeyError
 from edsl.language_models.RegisterLanguageModelsMeta import RegisterLanguageModelsMeta
-from edsl.exceptions.language_models import LanguageModelBadResponseError
-
-TIMEOUT = float(CONFIG.get("EDSL_API_TIMEOUT"))
-
-
-def convert_answer(response_part):
-    import json
-
-    response_part = response_part.strip()
-
-    if response_part == "None":
-        return None
-
-    repaired = repair_json(response_part)
-    if repaired == '""':
-        # it was a literal string
-        return response_part
-
-    try:
-        return json.loads(repaired)
-    except json.JSONDecodeError as j:
-        # last resort
-        return response_part
-
-
-def extract_item_from_raw_response(data, key_sequence):
-    if isinstance(data, str):
-        try:
-            data = json.loads(data)
-        except json.JSONDecodeError as e:
-            return data
-    current_data = data
-    for i, key in enumerate(key_sequence):
-        try:
-            if isinstance(current_data, (list, tuple)):
-                if not isinstance(key, int):
-                    raise TypeError(
-                        f"Expected integer index for sequence at position {i}, got {type(key).__name__}"
-                    )
-                if key < 0 or key >= len(current_data):
-                    raise IndexError(
-                        f"Index {key} out of range for sequence of length {len(current_data)} at position {i}"
-                    )
-            elif isinstance(current_data, dict):
-                if key not in current_data:
-                    raise KeyError(
-                        f"Key '{key}' not found in dictionary at position {i}"
-                    )
-            else:
-                raise TypeError(
-                    f"Cannot index into {type(current_data).__name__} at position {i}. Full response is: {data} of type {type(data)}. Key sequence is: {key_sequence}"
-                )
-
-            current_data = current_data[key]
-        except Exception as e:
-            path = " -> ".join(map(str, key_sequence[: i + 1]))
-            if "error" in data:
-                msg = data["error"]
-            else:
-                msg = f"Error accessing path: {path}. {str(e)}. Full response is: '{data}'"
-            raise LanguageModelBadResponseError(message=msg, response_json=data)
-    if isinstance(current_data, str):
-        return current_data.strip()
-    else:
-        return current_data
 
 
 def handle_key_error(func):
@@ -160,29 +92,21 @@ class LanguageModel(
     """
 
     _model_ = None
-    key_sequence = (
-        None  # This should be something like ["choices", 0, "message", "content"]
-    )
+
     __rate_limits = None
+    __default_rate_limits = {
+        "rpm": 10_000,
+        "tpm": 2_000_000,
+    }  # TODO: Use the OpenAI Teir 1 rate limits
     _safety_factor = 0.8
 
-    def __init__(
-        self, tpm=None, rpm=None, omit_system_prompt_if_empty_string=True, **kwargs
-    ):
+    def __init__(self, **kwargs):
         """Initialize the LanguageModel."""
         self.model = getattr(self, "_model_", None)
         default_parameters = getattr(self, "_parameters_", None)
         parameters = self._overide_default_parameters(kwargs, default_parameters)
         self.parameters = parameters
         self.remote = False
-        self.omit_system_prompt_if_empty = omit_system_prompt_if_empty_string
-
-        # self._rpm / _tpm comes from the class
-        if rpm is not None:
-            self._rpm = rpm
-
-        if tpm is not None:
-            self._tpm = tpm
 
         for key, value in parameters.items():
             setattr(self, key, value)
@@ -209,6 +133,7 @@ class LanguageModel(
     def api_token(self) -> str:
         if not hasattr(self, "_api_token"):
             key_name = service_to_api_keyname.get(self._inference_service_, "NOT FOUND")
+
             if self._inference_service_ == "bedrock":
                 self._api_token = [os.getenv(key_name[0]), os.getenv(key_name[1])]
                 # Check if any of the tokens are None
@@ -217,13 +142,13 @@ class LanguageModel(
                 self._api_token = os.getenv(key_name)
                 missing_token = self._api_token is None
             if missing_token and self._inference_service_ != "test" and not self.remote:
-                print("raising error")
+                print("rainsing error")
                 raise MissingAPIKeyError(
                     f"""The key for service: `{self._inference_service_}` is not set.
                         Need a key with name {key_name} in your .env file."""
                 )
 
-        return self._api_token
+            return self._api_token
 
     def __getitem__(self, key):
         return getattr(self, key)
@@ -284,58 +209,40 @@ class LanguageModel(
         >>> m = LanguageModel.example()
         >>> m.set_rate_limits(rpm=100, tpm=1000)
         >>> m.RPM
-        100
+        80.0
         """
-        if rpm is not None:
-            self._rpm = rpm
-        if tpm is not None:
-            self._tpm = tpm
-        return None
-        # self._set_rate_limits(rpm=rpm, tpm=tpm)
-
-    # def _set_rate_limits(self, rpm=None, tpm=None) -> None:
-    #     """Set the rate limits for the model.
-
-    #     If the model does not have rate limits, use the default rate limits."""
-    #     if rpm is not None and tpm is not None:
-    #         self.__rate_limits = {"rpm": rpm, "tpm": tpm}
-    #         return
-
-    #     if self.__rate_limits is None:
-    #         if hasattr(self, "get_rate_limits"):
-    #             self.__rate_limits = self.get_rate_limits()
-    #         else:
-    #             self.__rate_limits = self.__default_rate_limits
+        self._set_rate_limits(rpm=rpm, tpm=tpm)
+
+    def _set_rate_limits(self, rpm=None, tpm=None) -> None:
+        """Set the rate limits for the model.
+
+        If the model does not have rate limits, use the default rate limits."""
+        if rpm is not None and tpm is not None:
+            self.__rate_limits = {"rpm": rpm, "tpm": tpm}
+            return
+
+        if self.__rate_limits is None:
+            if hasattr(self, "get_rate_limits"):
+                self.__rate_limits = self.get_rate_limits()
+            else:
+                self.__rate_limits = self.__default_rate_limits
 
     @property
     def RPM(self):
         """Model's requests-per-minute limit."""
-        # self._set_rate_limits()
-        # return self._safety_factor * self.__rate_limits["rpm"]
-        return self._rpm
+        self._set_rate_limits()
+        return self._safety_factor * self.__rate_limits["rpm"]
 
     @property
     def TPM(self):
-        """Model's tokens-per-minute limit."""
-        # self._set_rate_limits()
-        # return self._safety_factor * self.__rate_limits["tpm"]
-        return self._tpm
-
-    @property
-    def rpm(self):
-        return self._rpm
+        """Model's tokens-per-minute limit.
 
-    @rpm.setter
-    def rpm(self, value):
-        self._rpm = value
-
-    @property
-    def tpm(self):
-        return self._tpm
-
-    @tpm.setter
-    def tpm(self, value):
-        self._tpm = value
+        >>> m = LanguageModel.example()
+        >>> m.TPM > 0
+        True
+        """
+        self._set_rate_limits()
+        return self._safety_factor * self.__rate_limits["tpm"]
 
     @staticmethod
     def _overide_default_parameters(passed_parameter_dict, default_parameter_dict):
@@ -363,10 +270,11 @@ class LanguageModel(
         >>> m = LanguageModel.example(test_model = True)
         >>> async def test(): return await m.async_execute_model_call("Hello, model!", "You are a helpful agent.")
         >>> asyncio.run(test())
-        {'message': [{'text': 'Hello world'}], ...}
+        {'message': '{"answer": "Hello world"}'}
 
         >>> m.execute_model_call("Hello, model!", "You are a helpful agent.")
-        {'message': [{'text': 'Hello world'}], ...}
+        {'message': '{"answer": "Hello world"}'}
+
         """
         pass
 
@@ -399,40 +307,68 @@ class LanguageModel(
 
         return main()
 
-    @classmethod
-    def get_generated_token_string(cls, raw_response: dict[str, Any]) -> str:
-        """Return the generated token string from the raw response."""
-        return extract_item_from_raw_response(raw_response, cls.key_sequence)
+    @abstractmethod
+    def parse_response(raw_response: dict[str, Any]) -> str:
+        """Parse the response and returns the response text.
 
-    @classmethod
-    def get_usage_dict(cls, raw_response: dict[str, Any]) -> dict[str, Any]:
-        """Return the usage dictionary from the raw response."""
-        if not hasattr(cls, "usage_sequence"):
-            raise NotImplementedError(
-                "This inference service does not have a usage_sequence."
+        >>> m = LanguageModel.example(test_model = True)
+        >>> m
+        Model(model_name = 'test', temperature = 0.5)
+
+        What is returned by the API is model-specific and often includes meta-data that we do not need.
+        For example, here is the results from a call to GPT-4:
+        To actually track the response, we need to grab
+        data["choices[0]"]["message"]["content"].
+        """
+        raise NotImplementedError
+
+    async def _async_prepare_response(
+        self, model_call_outcome: IntendedModelCallOutcome, cache: "Cache"
+    ) -> dict:
+        """Prepare the response for return."""
+
+        model_response = {
+            "cache_used": model_call_outcome.cache_used,
+            "cache_key": model_call_outcome.cache_key,
+            "usage": model_call_outcome.response.get("usage", {}),
+            "raw_model_response": model_call_outcome.response,
+        }
+
+        answer_portion = self.parse_response(model_call_outcome.response)
+        try:
+            answer_dict = json.loads(answer_portion)
+        except json.JSONDecodeError as e:
+            # TODO: Turn into logs to generate issues
+            answer_dict, success = await repair(
+                bad_json=answer_portion, error_message=str(e), cache=cache
             )
-        return extract_item_from_raw_response(raw_response, cls.usage_sequence)
+            if not success:
+                raise Exception(
+                    f"""Even the repair failed. The error was: {e}. The response was: {answer_portion}."""
+                )
 
-    @classmethod
-    def parse_response(cls, raw_response: dict[str, Any]) -> EDSLOutput:
-        """Parses the API response and returns the response text."""
-        generated_token_string = cls.get_generated_token_string(raw_response)
-        last_newline = generated_token_string.rfind("\n")
-
-        if last_newline == -1:
-            # There is no comment
-            edsl_dict = {
-                "answer": convert_answer(generated_token_string),
-                "generated_tokens": generated_token_string,
-                "comment": None,
-            }
-        else:
-            edsl_dict = {
-                "answer": convert_answer(generated_token_string[:last_newline]),
-                "comment": generated_token_string[last_newline + 1 :].strip(),
-                "generated_tokens": generated_token_string,
-            }
-        return EDSLOutput(**edsl_dict)
+        return {**model_response, **answer_dict}
+
+    async def async_get_raw_response(
+        self,
+        user_prompt: str,
+        system_prompt: str,
+        cache: "Cache",
+        iteration: int = 0,
+        encoded_image=None,
+    ) -> IntendedModelCallOutcome:
+        import warnings
+
+        warnings.warn(
+            "This method is deprecated. Use async_get_intended_model_call_outcome."
+        )
+        return await self._async_get_intended_model_call_outcome(
+            user_prompt=user_prompt,
+            system_prompt=system_prompt,
+            cache=cache,
+            iteration=iteration,
+            encoded_image=encoded_image,
+        )
 
     async def _async_get_intended_model_call_outcome(
         self,
@@ -441,7 +377,7 @@ class LanguageModel(
         cache: "Cache",
         iteration: int = 0,
         encoded_image=None,
-    ) -> ModelResponse:
+    ) -> IntendedModelCallOutcome:
         """Handle caching of responses.
 
         :param user_prompt: The user's prompt.
@@ -460,18 +396,18 @@ class LanguageModel(
         >>> from edsl import Cache
         >>> m = LanguageModel.example(test_model = True)
         >>> m._get_intended_model_call_outcome(user_prompt = "Hello", system_prompt = "hello", cache = Cache())
-        ModelResponse(...)"""
+        IntendedModelCallOutcome(response = {'message': '{"answer": "Hello world"}'}, cache_used = False, cache_key = '24ff6ac2bc2f1729f817f261e0792577')
+        """
 
         if encoded_image:
             # the image has is appended to the user_prompt for hash-lookup purposes
             image_hash = hashlib.md5(encoded_image.encode()).hexdigest()
-            user_prompt += f" {image_hash}"
 
         cache_call_params = {
             "model": str(self.model),
             "parameters": self.parameters,
             "system_prompt": system_prompt,
-            "user_prompt": user_prompt,
+            "user_prompt": user_prompt + "" if not encoded_image else f" {image_hash}",
             "iteration": iteration,
         }
         cached_response, cache_key = cache.fetch(**cache_call_params)
@@ -489,28 +425,21 @@ class LanguageModel(
                 "system_prompt": system_prompt,
                 **({"encoded_image": encoded_image} if encoded_image else {}),
             }
-            # response = await f(**params)
-            response = await asyncio.wait_for(f(**params), timeout=TIMEOUT)
+            response = await f(**params)
             new_cache_key = cache.store(
                 **cache_call_params, response=response
             )  # store the response in the cache
             assert new_cache_key == cache_key  # should be the same
 
-        cost = self.cost(response)
-
-        return ModelResponse(
-            response=response,
-            cache_used=cache_used,
-            cache_key=cache_key,
-            cached_response=cached_response,
-            cost=cost,
+        return IntendedModelCallOutcome(
+            response=response, cache_used=cache_used, cache_key=cache_key
         )
 
     _get_intended_model_call_outcome = sync_wrapper(
         _async_get_intended_model_call_outcome
     )
 
-    # get_raw_response = sync_wrapper(async_get_raw_response)
+    get_raw_response = sync_wrapper(async_get_raw_response)
 
     def simple_ask(
         self,
@@ -549,66 +478,14 @@ class LanguageModel(
             "cache": cache,
             **({"encoded_image": encoded_image} if encoded_image else {}),
         }
-        model_inputs = ModelInputs(user_prompt=user_prompt, system_prompt=system_prompt)
-        model_outputs = await self._async_get_intended_model_call_outcome(**params)
-        edsl_dict = self.parse_response(model_outputs.response)
-        agent_response_dict = AgentResponseDict(
-            model_inputs=model_inputs,
-            model_outputs=model_outputs,
-            edsl_dict=edsl_dict,
-        )
-        return agent_response_dict
-
-        # return await self._async_prepare_response(model_call_outcome, cache=cache)
+        model_call_outcome = await self._async_get_intended_model_call_outcome(**params)
+        return await self._async_prepare_response(model_call_outcome, cache=cache)
 
     get_response = sync_wrapper(async_get_response)
 
-    def cost(self, raw_response: dict[str, Any]) -> Union[float, str]:
+    def cost(self, raw_response: dict[str, Any]) -> float:
         """Return the dollar cost of a raw response."""
-
-        usage = self.get_usage_dict(raw_response)
-        from edsl.coop import Coop
-
-        c = Coop()
-        price_lookup = c.fetch_prices()
-        key = (self._inference_service_, self.model)
-        if key not in price_lookup:
-            return f"Could not find price for model {self.model} in the price lookup."
-
-        relevant_prices = price_lookup[key]
-        try:
-            input_tokens = int(usage[self.input_token_name])
-            output_tokens = int(usage[self.output_token_name])
-        except Exception as e:
-            return f"Could not fetch tokens from model response: {e}"
-
-        try:
-            inverse_output_price = relevant_prices["output"]["one_usd_buys"]
-            inverse_input_price = relevant_prices["input"]["one_usd_buys"]
-        except Exception as e:
-            if "output" not in relevant_prices:
-                return f"Could not fetch prices from {relevant_prices} - {e}; Missing 'output' key."
-            if "input" not in relevant_prices:
-                return f"Could not fetch prices from {relevant_prices} - {e}; Missing 'input' key."
-            return f"Could not fetch prices from {relevant_prices} - {e}"
-
-        if inverse_input_price == "infinity":
-            input_cost = 0
-        else:
-            try:
-                input_cost = input_tokens / float(inverse_input_price)
-            except Exception as e:
-                return f"Could not compute input price - {e}."
-
-        if inverse_output_price == "infinity":
-            output_cost = 0
-        else:
-            try:
-                output_cost = output_tokens / float(inverse_output_price)
-            except Exception as e:
-                return f"Could not compute output price - {e}"
-
-        return input_cost + output_cost
+        raise NotImplementedError
 
     #######################
     # SERIALIZATION METHODS
@@ -622,7 +499,7 @@ class LanguageModel(
 
         >>> m = LanguageModel.example()
         >>> m.to_dict()
-        {'model': '...', 'parameters': {'temperature': ..., 'max_tokens': ..., 'top_p': ..., 'frequency_penalty': ..., 'presence_penalty': ..., 'logprobs': False, 'top_logprobs': ...}, 'edsl_version': '...', 'edsl_class_name': 'LanguageModel'}
+        {'model': 'gpt-4-1106-preview', 'parameters': {'temperature': 0.5, 'max_tokens': 1000, 'top_p': 1, 'frequency_penalty': 0, 'presence_penalty': 0, 'logprobs': False, 'top_logprobs': 3}, 'edsl_version': '...', 'edsl_class_name': 'LanguageModel'}
         """
         return self._to_dict()
 
@@ -698,8 +575,26 @@ class LanguageModel(
         """
         from edsl import Model
 
+        class TestLanguageModelGood(LanguageModel):
+            use_cache = False
+            _model_ = "test"
+            _parameters_ = {"temperature": 0.5}
+            _inference_service_ = InferenceServiceType.TEST.value
+
+            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 throw_exception:
+                    raise Exception("This is a test error")
+                return {"message": f'{{"answer": "{canned_response}"}}'}
+
+            def parse_response(self, raw_response: dict[str, Any]) -> str:
+                return raw_response["message"]
+
         if test_model:
-            m = Model("test", canned_response=canned_response)
+            m = TestLanguageModelGood()
             return m
         else:
             return Model(skip_api_key_check=True)

edsl/language_models/ModelList.py

@@ -40,8 +40,8 @@ class ModelList(Base, UserList):
     def __hash__(self):
         """Return a hash of the ModelList. This is used for comparison of ModelLists.
 
-        >>> isinstance(hash(Model()), int)
-        True
+        >>> hash(ModelList.example())
+        1423518243781418961
 
         """
         from edsl.utilities.utilities import dict_hash

edsl/language_models/RegisterLanguageModelsMeta.py

@@ -47,6 +47,13 @@ class RegisterLanguageModelsMeta(ABCMeta):
                 must_be_async=True,
             )
             # LanguageModel children have to implement the parse_response method
+            RegisterLanguageModelsMeta.verify_method(
+                candidate_class=cls,
+                method_name="parse_response",
+                expected_return_type=str,
+                required_parameters=[("raw_response", dict[str, Any])],
+                must_be_async=False,
+            )
             RegisterLanguageModelsMeta._registry[model_name] = cls
 
     @classmethod
@@ -91,7 +98,7 @@ class RegisterLanguageModelsMeta(ABCMeta):
 
         required_parameters = required_parameters or []
         method = getattr(candidate_class, method_name)
-        # signature = inspect.signature(method)
+        signature = inspect.signature(method)
 
         RegisterLanguageModelsMeta._check_return_type(method, expected_return_type)
 
@@ -99,11 +106,11 @@ class RegisterLanguageModelsMeta(ABCMeta):
             RegisterLanguageModelsMeta._check_is_coroutine(method)
 
         # Check the parameters
-        # params = signature.parameters
-        # for param_name, param_type in required_parameters:
-        #     RegisterLanguageModelsMeta._verify_parameter(
-        #         params, param_name, param_type, method_name
-        #     )
+        params = signature.parameters
+        for param_name, param_type in required_parameters:
+            RegisterLanguageModelsMeta._verify_parameter(
+                params, param_name, param_type, method_name
+            )
 
     @staticmethod
     def _check_method_defined(cls, method_name):
@@ -160,15 +167,23 @@ class RegisterLanguageModelsMeta(ABCMeta):
         Check if the return type of a method is as expected.
 
         Example:
+        >>> class M:
+        ...     async def f(self) -> str: pass
+        >>> RegisterLanguageModelsMeta._check_return_type(M.f, str)
+        >>> class N:
+        ...     async def f(self) -> int: pass
+        >>> RegisterLanguageModelsMeta._check_return_type(N.f, str)
+        Traceback (most recent call last):
+        ...
+        TypeError: Return type of f must be <class 'str'>. Got <class 'int'>.
         """
-        pass
-        # if inspect.isroutine(method):
-        #     # return_type = inspect.signature(method).return_annotation
-        #     return_type = get_type_hints(method)["return"]
-        #     if return_type != expected_return_type:
-        #         raise TypeError(
-        #             f"Return type of {method.__name__} must be {expected_return_type}. Got {return_type}."
-        #         )
+        if inspect.isroutine(method):
+            # return_type = inspect.signature(method).return_annotation
+            return_type = get_type_hints(method)["return"]
+            if return_type != expected_return_type:
+                raise TypeError(
+                    f"Return type of {method.__name__} must be {expected_return_type}. Got {return_type}."
+                )
 
     @classmethod
     def model_names_to_classes(cls):