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

edsl/language_models/KeyLookup.py

@@ -0,0 +1,30 @@
+import os
+from collections import UserDict
+
+from edsl.enums import service_to_api_keyname
+from edsl.exceptions import MissingAPIKeyError
+
+
+class KeyLookup(UserDict):
+    @classmethod
+    def from_os_environ(cls):
+        """Create an instance of KeyLookupAPI with keys from os.environ"""
+        return cls({key: value for key, value in os.environ.items()})
+
+    def get_api_token(self, service: str, remote: bool = False):
+        key_name = service_to_api_keyname.get(service, "NOT FOUND")
+
+        if service == "bedrock":
+            api_token = [self.get(key_name[0]), self.get(key_name[1])]
+            missing_token = any(token is None for token in api_token)
+        else:
+            api_token = self.get(key_name)
+            missing_token = api_token is None
+
+        if missing_token and service != "test" and not remote:
+            raise MissingAPIKeyError(
+                f"""The key for service: `{service}` is not set.
+                    Need a key with name {key_name} in your .env file."""
+            )
+
+        return api_token

edsl/language_models/LanguageModel.py

@@ -21,6 +21,7 @@ import os
 from typing import (
     Coroutine,
     Any,
+    Callable,
     Type,
     Union,
     List,
@@ -31,6 +32,8 @@ from typing import (
 )
 from abc import ABC, abstractmethod
 
+from json_repair import repair_json
+
 from edsl.data_transfer_models import (
     ModelResponse,
     ModelInputs,
@@ -42,24 +45,61 @@ if TYPE_CHECKING:
     from edsl.data.Cache import Cache
     from edsl.scenarios.FileStore import FileStore
     from edsl.questions.QuestionBase import QuestionBase
-    from edsl.language_models.key_management.KeyLookup import KeyLookup
-
-from edsl.enums import InferenceServiceType
 
-from edsl.utilities.decorators import (
-    sync_wrapper,
-    jupyter_nb_handler,
-)
-from edsl.utilities.remove_edsl_version import remove_edsl_version
+from edsl.config import CONFIG
+from edsl.utilities.decorators import sync_wrapper, jupyter_nb_handler
+from edsl.utilities.decorators import remove_edsl_version
 
-from edsl.Base import PersistenceMixin, RepresentationMixin
+from edsl.Base import PersistenceMixin
 from edsl.language_models.RegisterLanguageModelsMeta import RegisterLanguageModelsMeta
+from edsl.language_models.KeyLookup import KeyLookup
+from edsl.exceptions.language_models import LanguageModelBadResponseError
 
-from edsl.language_models.key_management.KeyLookupCollection import (
-    KeyLookupCollection,
-)
+TIMEOUT = float(CONFIG.get("EDSL_API_TIMEOUT"))
+
+
+# you might be tempated to move this to be a static method of LanguageModel, but this doesn't work
+# for reasons I don't understand. So leave it here.
+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}"
+                )
 
-from edsl.language_models.RawResponseHandler import RawResponseHandler
+            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):
@@ -77,21 +117,8 @@ def handle_key_error(func):
     return wrapper
 
 
-class classproperty:
-    def __init__(self, method):
-        self.method = method
-
-    def __get__(self, instance, cls):
-        return self.method(cls)
-
-
-from edsl.Base import HashingMixin
-
-
 class LanguageModel(
     PersistenceMixin,
-    RepresentationMixin,
-    HashingMixin,
     ABC,
     metaclass=RegisterLanguageModelsMeta,
 ):
@@ -101,22 +128,15 @@ class LanguageModel(
     key_sequence = (
         None  # This should be something like ["choices", 0, "message", "content"]
     )
-
-    DEFAULT_RPM = 100
-    DEFAULT_TPM = 1000
-
-    @classproperty
-    def response_handler(cls):
-        key_sequence = cls.key_sequence
-        usage_sequence = cls.usage_sequence if hasattr(cls, "usage_sequence") else None
-        return RawResponseHandler(key_sequence, usage_sequence)
+    __rate_limits = None
+    _safety_factor = 0.8
 
     def __init__(
         self,
-        tpm: Optional[float] = None,
-        rpm: Optional[float] = None,
+        tpm: float = None,
+        rpm: float = None,
         omit_system_prompt_if_empty_string: bool = True,
-        key_lookup: Optional["KeyLookup"] = None,
+        key_lookup: Optional[KeyLookup] = None,
         **kwargs,
     ):
         """Initialize the LanguageModel."""
@@ -127,9 +147,7 @@ class LanguageModel(
         self.remote = False
         self.omit_system_prompt_if_empty = omit_system_prompt_if_empty_string
 
-        self.key_lookup = self._set_key_lookup(key_lookup)
-        self.model_info = self.key_lookup.get(self._inference_service_)
-
+        # self._rpm / _tpm comes from the class
         if rpm is not None:
             self._rpm = rpm
 
@@ -143,75 +161,49 @@ class LanguageModel(
             if key not in parameters:
                 setattr(self, key, value)
 
-        if kwargs.get("skip_api_key_check", False):
+        if "use_cache" in kwargs:
+            warnings.warn(
+                "The use_cache parameter is deprecated. Use the Cache class instead."
+            )
+
+        if skip_api_key_check := kwargs.get("skip_api_key_check", False):
             # Skip the API key check. Sometimes this is useful for testing.
             self._api_token = None
 
-    def _set_key_lookup(self, key_lookup: "KeyLookup") -> "KeyLookup":
-        """Set the key lookup."""
         if key_lookup is not None:
-            return key_lookup
+            self.key_lookup = key_lookup
         else:
-            klc = KeyLookupCollection()
-            klc.add_key_lookup(fetch_order=("config", "env"))
-            return klc.get(("config", "env"))
-
-    def set_key_lookup(self, key_lookup: "KeyLookup") -> None:
-        """Set the key lookup, later"""
-        if hasattr(self, "_api_token"):
-            del self._api_token
-        self.key_lookup = key_lookup
-
-    def ask_question(self, question: "QuestionBase") -> str:
-        """Ask a question and return the response.
+            self.key_lookup = KeyLookup.from_os_environ()
 
-        :param question: The question to ask.
-        """
+    def ask_question(self, question):
         user_prompt = question.get_instructions().render(question.data).text
         system_prompt = "You are a helpful agent pretending to be a human."
         return self.execute_model_call(user_prompt, system_prompt)
 
-    @property
-    def rpm(self):
-        if not hasattr(self, "_rpm"):
-            if self.model_info is None:
-                self._rpm = self.DEFAULT_RPM
-            else:
-                self._rpm = self.model_info.rpm
-        return self._rpm
-
-    @property
-    def tpm(self):
-        if not hasattr(self, "_tpm"):
-            if self.model_info is None:
-                self._tpm = self.DEFAULT_TPM
-            else:
-                self._tpm = self.model_info.tpm
-        return self._tpm
-
-    # in case we want to override the default values
-    @tpm.setter
-    def tpm(self, value):
-        self._tpm = value
-
-    @rpm.setter
-    def rpm(self, value):
-        self._rpm = value
+    def set_key_lookup(self, key_lookup: KeyLookup) -> None:
+        del self._api_token
+        self.key_lookup = key_lookup
 
     @property
     def api_token(self) -> str:
         if not hasattr(self, "_api_token"):
-            info = self.key_lookup.get(self._inference_service_, None)
-            if info is None:
-                raise ValueError(
-                    f"No key found for service '{self._inference_service_}'"
-                )
-            self._api_token = info.api_token
+            self._api_token = self.key_lookup.get_api_token(
+                self._inference_service_, self.remote
+            )
         return self._api_token
 
     def __getitem__(self, key):
         return getattr(self, key)
 
+    def _repr_html_(self) -> str:
+        d = {"model": self.model}
+        d.update(self.parameters)
+        data = [[k, v] for k, v in d.items()]
+        from tabulate import tabulate
+
+        table = str(tabulate(data, headers=["keys", "values"], tablefmt="html"))
+        return f"<pre>{table}</pre>"
+
     def hello(self, verbose=False):
         """Runs a simple test to check if the model is working."""
         token = self.api_token
@@ -240,12 +232,7 @@ class LanguageModel(
         return key_value is not None
 
     def __hash__(self) -> str:
-        """Allow the model to be used as a key in a dictionary.
-
-        >>> m = LanguageModel.example()
-        >>> hash(m)
-        1811901442659237949
-        """
+        """Allow the model to be used as a key in a dictionary."""
         from edsl.utilities.utilities import dict_hash
 
         return dict_hash(self.to_dict(add_edsl_version=False))
@@ -261,6 +248,46 @@ class LanguageModel(
         """
         return self.model == other.model and self.parameters == other.parameters
 
+    def set_rate_limits(self, rpm=None, tpm=None) -> None:
+        """Set the rate limits for the model.
+
+        >>> m = LanguageModel.example()
+        >>> m.set_rate_limits(rpm=100, tpm=1000)
+        >>> m.RPM
+        100
+        """
+        if rpm is not None:
+            self._rpm = rpm
+        if tpm is not None:
+            self._tpm = tpm
+        return None
+
+    @property
+    def RPM(self):
+        """Model's requests-per-minute limit."""
+        return self._rpm
+
+    @property
+    def TPM(self):
+        """Model's tokens-per-minute limit."""
+        return self._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):
         """Return a dictionary of parameters, with passed parameters taking precedence over defaults.
@@ -283,7 +310,16 @@ class LanguageModel(
 
     @abstractmethod
     async def async_execute_model_call(user_prompt: str, system_prompt: str):
-        """Execute the model call and returns a coroutine."""
+        """Execute the model call and returns a coroutine.
+
+        >>> 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'}], ...}
+
+        >>> m.execute_model_call("Hello, model!", "You are a helpful agent.")
+        {'message': [{'text': 'Hello world'}], ...}
+        """
         pass
 
     async def remote_async_execute_model_call(
@@ -300,7 +336,12 @@ class LanguageModel(
 
     @jupyter_nb_handler
     def execute_model_call(self, *args, **kwargs) -> Coroutine:
-        """Execute the model call and returns the result as a coroutine."""
+        """Execute the model call and returns the result as a coroutine.
+
+        >>> m = LanguageModel.example(test_model = True)
+        >>> m.execute_model_call(user_prompt = "Hello, model!", system_prompt = "You are a helpful agent.")
+
+        """
 
         async def main():
             results = await asyncio.gather(
@@ -312,25 +353,58 @@ class LanguageModel(
 
     @classmethod
     def get_generated_token_string(cls, raw_response: dict[str, Any]) -> str:
-        """Return the generated token string from the raw response.
-
-        >>> m = LanguageModel.example(test_model = True)
-        >>> raw_response = m.execute_model_call("Hello, model!", "You are a helpful agent.")
-        >>> m.get_generated_token_string(raw_response)
-        'Hello world'
-
-        """
-        return cls.response_handler.get_generated_token_string(raw_response)
+        """Return the generated token string from the raw response."""
+        return extract_item_from_raw_response(raw_response, cls.key_sequence)
 
     @classmethod
     def get_usage_dict(cls, raw_response: dict[str, Any]) -> dict[str, Any]:
         """Return the usage dictionary from the raw response."""
-        return cls.response_handler.get_usage_dict(raw_response)
+        if not hasattr(cls, "usage_sequence"):
+            raise NotImplementedError(
+                "This inference service does not have a usage_sequence."
+            )
+        return extract_item_from_raw_response(raw_response, cls.usage_sequence)
+
+    @staticmethod
+    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
 
     @classmethod
     def parse_response(cls, raw_response: dict[str, Any]) -> EDSLOutput:
         """Parses the API response and returns the response text."""
-        return cls.response_handler.parse_response(raw_response)
+        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": cls.convert_answer(generated_token_string),
+                "generated_tokens": generated_token_string,
+                "comment": None,
+            }
+        else:
+            edsl_dict = {
+                "answer": cls.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,
@@ -347,8 +421,6 @@ class LanguageModel(
         :param system_prompt: The system's prompt.
         :param iteration: The iteration number.
         :param cache: The cache to use.
-        :param files_list: The list of files to use.
-        :param invigilator: The invigilator to use.
 
         If the cache isn't being used, it just returns a 'fresh' call to the LLM.
         But if cache is being used, it first checks the database to see if the response is already there.
@@ -391,10 +463,6 @@ class LanguageModel(
                 "system_prompt": system_prompt,
                 "files_list": files_list,
             }
-            from edsl.config import CONFIG
-
-            TIMEOUT = float(CONFIG.get("EDSL_API_TIMEOUT"))
-
             response = await asyncio.wait_for(f(**params), timeout=TIMEOUT)
             new_cache_key = cache.store(
                 **cache_call_params, response=response
@@ -402,6 +470,7 @@ class LanguageModel(
             assert new_cache_key == cache_key  # should be the same
 
         cost = self.cost(response)
+
         return ModelResponse(
             response=response,
             cache_used=cache_used,
@@ -440,9 +509,9 @@ class LanguageModel(
 
         :param user_prompt: The user's prompt.
         :param system_prompt: The system's prompt.
-        :param cache: The cache to use.
         :param iteration: The iteration number.
-        :param files_list: The list of files to use.
+        :param cache: The cache to use.
+        :param encoded_image: The encoded image to use.
 
         """
         params = {
@@ -456,11 +525,8 @@ class LanguageModel(
             params.update({"invigilator": kwargs["invigilator"]})
 
         model_inputs = ModelInputs(user_prompt=user_prompt, system_prompt=system_prompt)
-        model_outputs: ModelResponse = (
-            await self._async_get_intended_model_call_outcome(**params)
-        )
-        edsl_dict: EDSLOutput = self.parse_response(model_outputs.response)
-
+        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,
@@ -471,36 +537,60 @@ class LanguageModel(
     get_response = sync_wrapper(async_get_response)
 
     def cost(self, raw_response: dict[str, Any]) -> Union[float, str]:
-        """Return the dollar cost of a raw response.
-
-        :param raw_response: The raw response from the model.
-        """
+        """Return the dollar cost of a raw response."""
 
         usage = self.get_usage_dict(raw_response)
-        from edsl.language_models.PriceManager import PriceManager
-
-        price_manger = PriceManager()
-        return price_manger.calculate_cost(
-            inference_service=self._inference_service_,
-            model=self.model,
-            usage=usage,
-            input_token_name=self.input_token_name,
-            output_token_name=self.output_token_name,
-        )
+        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
 
     def to_dict(self, add_edsl_version: bool = True) -> dict[str, Any]:
         """Convert instance to a dictionary
 
-        :param add_edsl_version: Whether to add the EDSL version to the dictionary.
-
         >>> 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'}
         """
-        d = {
-            "model": self.model,
-            "parameters": self.parameters,
-        }
+        d = {"model": self.model, "parameters": self.parameters}
         if add_edsl_version:
             from edsl import __version__
 
@@ -512,13 +602,13 @@ class LanguageModel(
     @remove_edsl_version
     def from_dict(cls, data: dict) -> Type[LanguageModel]:
         """Convert dictionary to a LanguageModel child instance."""
-        from edsl.language_models.model import get_model_class
+        from edsl.language_models.registry import get_model_class
 
         model_class = get_model_class(data["model"])
         return model_class(**data)
 
     def __repr__(self) -> str:
-        """Return a representation of the object."""
+        """Return a string representation of the object."""
         param_string = ", ".join(
             f"{key} = {value}" for key, value in self.parameters.items()
         )
@@ -560,7 +650,7 @@ class LanguageModel(
         Exception report saved to ...
         Also see: ...
         """
-        from edsl.language_models.model import Model
+        from edsl import Model
 
         if test_model:
             m = Model(
@@ -570,54 +660,6 @@ class LanguageModel(
         else:
             return Model(skip_api_key_check=True)
 
-    def from_cache(self, cache: "Cache") -> LanguageModel:
-
-        from copy import deepcopy
-        from types import MethodType
-        from edsl import Cache
-
-        new_instance = deepcopy(self)
-        print("Cache entries", len(cache))
-        new_instance.cache = Cache(
-            data={k: v for k, v in cache.items() if v.model == self.model}
-        )
-        print("Cache entries with same model", len(new_instance.cache))
-
-        new_instance.user_prompts = [
-            ce.user_prompt for ce in new_instance.cache.values()
-        ]
-        new_instance.system_prompts = [
-            ce.system_prompt for ce in new_instance.cache.values()
-        ]
-
-        async def async_execute_model_call(self, user_prompt: str, system_prompt: str):
-            cache_call_params = {
-                "model": str(self.model),
-                "parameters": self.parameters,
-                "system_prompt": system_prompt,
-                "user_prompt": user_prompt,
-                "iteration": 1,
-            }
-            cached_response, cache_key = cache.fetch(**cache_call_params)
-            response = json.loads(cached_response)
-            cost = 0
-            return ModelResponse(
-                response=response,
-                cache_used=True,
-                cache_key=cache_key,
-                cached_response=cached_response,
-                cost=cost,
-            )
-
-        # Bind the new method to the copied instance
-        setattr(
-            new_instance,
-            "async_execute_model_call",
-            MethodType(async_execute_model_call, new_instance),
-        )
-
-        return new_instance
-
 
 if __name__ == "__main__":
     """Run the module's test suite."""