edsl 0.1.31.dev4__py3-none-any.whl → 0.1.33__py3-none-any.whl

edsl/language_models/LanguageModel.py

@@ -1,4 +1,16 @@
-"""This module contains the LanguageModel class, which is an abstract base class for all language models."""
+"""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': ...}
+
+"""
 
 from __future__ import annotations
 import warnings
@@ -8,45 +20,103 @@ import json
 import time
 import os
 import hashlib
-from typing import Coroutine, Any, Callable, Type, List, get_type_hints
+from typing import (
+    Coroutine,
+    Any,
+    Callable,
+    Type,
+    Union,
+    List,
+    get_type_hints,
+    TypedDict,
+    Optional,
+)
 from abc import ABC, abstractmethod
 
-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
+from json_repair import repair_json
 
-    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
+from edsl.data_transfer_models import (
+    ModelResponse,
+    ModelInputs,
+    EDSLOutput,
+    AgentResponseDict,
+)
 
-    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):
@@ -90,21 +160,29 @@ 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, **kwargs):
+    def __init__(
+        self, tpm=None, rpm=None, omit_system_prompt_if_empty_string=True, **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)
@@ -131,17 +209,20 @@ 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")
-            self._api_token = os.getenv(key_name)
-            if (
-                self._api_token is None
-                and self._inference_service_ != "test"
-                and not self.remote
-            ):
+            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
+                missing_token = any(token is None for token in self._api_token)
+            else:
+                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")
                 raise MissingAPIKeyError(
                     f"""The key for service: `{self._inference_service_}` is not set.
-                    Need a key with name {key_name} in your .env file.
-                    """
+                        Need a key with name {key_name} in your .env file."""
                 )
+
         return self._api_token
 
     def __getitem__(self, key):
@@ -159,8 +240,7 @@ class LanguageModel(
         if verbose:
             print(f"Current key is {masked}")
         return self.execute_model_call(
-            user_prompt="Hello, model!", 
-            system_prompt="You are a helpful agent."
+            user_prompt="Hello, model!", system_prompt="You are a helpful agent."
         )
 
     def has_valid_api_key(self) -> bool:
@@ -204,42 +284,58 @@ class LanguageModel(
         >>> m = LanguageModel.example()
         >>> m.set_rate_limits(rpm=100, tpm=1000)
         >>> m.RPM
-        80.0
+        100
         """
-        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
+        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
 
     @property
     def RPM(self):
         """Model's requests-per-minute limit."""
-        self._set_rate_limits()
-        return self._safety_factor * self.__rate_limits["rpm"]
+        # self._set_rate_limits()
+        # return self._safety_factor * self.__rate_limits["rpm"]
+        return self._rpm
 
     @property
     def TPM(self):
-        """Model's tokens-per-minute limit.
+        """Model's tokens-per-minute limit."""
+        # self._set_rate_limits()
+        # return self._safety_factor * self.__rate_limits["tpm"]
+        return self._tpm
 
-        >>> m = LanguageModel.example()
-        >>> m.TPM > 0
-        True
-        """
-        self._set_rate_limits()
-        return self._safety_factor * self.__rate_limits["tpm"]
+    @property
+    def rpm(self):
+        return self._rpm
+
+    @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
 
     @staticmethod
     def _overide_default_parameters(passed_parameter_dict, default_parameter_dict):
@@ -250,14 +346,16 @@ class LanguageModel(
         >>> LanguageModel._overide_default_parameters(passed_parameter_dict={"temperature": 0.5}, default_parameter_dict={"temperature":0.9, "max_tokens": 1000})
         {'temperature': 0.5, 'max_tokens': 1000}
         """
-        #parameters = dict({})
+        # parameters = dict({})
+
+        return {
+            parameter_name: passed_parameter_dict.get(parameter_name, default_value)
+            for parameter_name, default_value in default_parameter_dict.items()
+        }
 
-        return {parameter_name: passed_parameter_dict.get(parameter_name, default_value) 
-                for parameter_name, default_value in default_parameter_dict.items()}
-    
-    def __call__(self, user_prompt:str, system_prompt:str):
+    def __call__(self, user_prompt: str, system_prompt: str):
         return self.execute_model_call(user_prompt, system_prompt)
-    
+
     @abstractmethod
     async def async_execute_model_call(user_prompt: str, system_prompt: str):
         """Execute the model call and returns a coroutine.
@@ -265,11 +363,10 @@ 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': '{"answer": "Hello world"}'}
+        {'message': [{'text': 'Hello world'}], ...}
 
         >>> m.execute_model_call("Hello, model!", "You are a helpful agent.")
-        {'message': '{"answer": "Hello world"}'}
-
+        {'message': [{'text': 'Hello world'}], ...}
         """
         pass
 
@@ -302,66 +399,40 @@ class LanguageModel(
 
         return main()
 
-    @abstractmethod
-    def parse_response(raw_response: dict[str, Any]) -> str:
-        """Parse the response and returns the response text.
-
-        >>> 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,
-        }
+    @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)
 
-        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
+    @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."
             )
-            if not success:
-                raise Exception(
-                    f"""Even the repair failed. The error was: {e}. The response was: {answer_portion}."""
-                )
-    
-        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
-        )
+        return extract_item_from_raw_response(raw_response, cls.usage_sequence)
 
+    @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)
 
     async def _async_get_intended_model_call_outcome(
         self,
@@ -370,7 +441,7 @@ class LanguageModel(
         cache: "Cache",
         iteration: int = 0,
         encoded_image=None,
-    ) -> IntendedModelCallOutcome:
+    ) -> ModelResponse:
         """Handle caching of responses.
 
         :param user_prompt: The user's prompt.
@@ -389,23 +460,23 @@ 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())
-        IntendedModelCallOutcome(response = {'message': '{"answer": "Hello world"}'}, cache_used = False, cache_key = '24ff6ac2bc2f1729f817f261e0792577')
-        """
+        ModelResponse(...)"""
 
         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 + "" if not encoded_image else f" {image_hash}",
+            "user_prompt": user_prompt,
             "iteration": iteration,
         }
         cached_response, cache_key = cache.fetch(**cache_call_params)
-    
-        if (cache_used := cached_response is not None):
+
+        if cache_used := cached_response is not None:
             response = json.loads(cached_response)
         else:
             f = (
@@ -413,18 +484,33 @@ class LanguageModel(
                 if hasattr(self, "remote") and self.remote
                 else self.async_execute_model_call
             )
-            params = {"user_prompt": user_prompt, "system_prompt": system_prompt, 
-            **({"encoded_image": encoded_image} if encoded_image else {})
+            params = {
+                "user_prompt": user_prompt,
+                "system_prompt": system_prompt,
+                **({"encoded_image": encoded_image} if encoded_image else {}),
             }
-            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
-     
-        return IntendedModelCallOutcome(response = response, cache_used = cache_used, cache_key = cache_key)
+            # response = await f(**params)
+            response = await asyncio.wait_for(f(**params), timeout=TIMEOUT)
+            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,
+        )
 
-    _get_intended_model_call_outcome = sync_wrapper(_async_get_intended_model_call_outcome)
+    _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,
@@ -443,7 +529,7 @@ class LanguageModel(
         self,
         user_prompt: str,
         system_prompt: str,
-        cache: 'Cache',
+        cache: "Cache",
         iteration: int = 1,
         encoded_image=None,
     ) -> dict:
@@ -461,16 +547,68 @@ class LanguageModel(
             "system_prompt": system_prompt,
             "iteration": iteration,
             "cache": cache,
-            **({"encoded_image": encoded_image} if encoded_image else {})
-        }        
-        model_call_outcome = await self._async_get_intended_model_call_outcome(**params)
-        return await self._async_prepare_response(model_call_outcome, 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)
 
     get_response = sync_wrapper(async_get_response)
 
-    def cost(self, raw_response: dict[str, Any]) -> float:
+    def cost(self, raw_response: dict[str, Any]) -> Union[float, str]:
         """Return the dollar cost of a raw response."""
-        raise NotImplementedError
+
+        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
 
     #######################
     # SERIALIZATION METHODS
@@ -484,7 +622,7 @@ class LanguageModel(
 
         >>> m = LanguageModel.example()
         >>> m.to_dict()
-        {'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'}
+        {'model': '...', 'parameters': {'temperature': ..., 'max_tokens': ..., 'top_p': ..., 'frequency_penalty': ..., 'presence_penalty': ..., 'logprobs': False, 'top_logprobs': ...}, 'edsl_version': '...', 'edsl_class_name': 'LanguageModel'}
         """
         return self._to_dict()
 
@@ -560,26 +698,8 @@ 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 = TestLanguageModelGood()
+            m = Model("test", canned_response=canned_response)
             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.
 
-        >>> hash(ModelList.example())
-        1423518243781418961
+        >>> isinstance(hash(Model()), int)
+        True
 
         """
         from edsl.utilities.utilities import dict_hash