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

edsl/language_models/LanguageModel.py

@@ -21,7 +21,6 @@ import os
 from typing import (
     Coroutine,
     Any,
-    Callable,
     Type,
     Union,
     List,
@@ -32,8 +31,6 @@ from typing import (
 )
 from abc import ABC, abstractmethod
 
-from json_repair import repair_json
-
 from edsl.data_transfer_models import (
     ModelResponse,
     ModelInputs,
@@ -45,61 +42,22 @@ 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.config import CONFIG
-from edsl.utilities.decorators import sync_wrapper, jupyter_nb_handler
-from edsl.utilities.decorators import remove_edsl_version
+from edsl.utilities.decorators import (
+    sync_wrapper,
+    jupyter_nb_handler,
+)
+from edsl.utilities.remove_edsl_version import remove_edsl_version
 
-from edsl.Base import PersistenceMixin
+from edsl.Base import PersistenceMixin, RepresentationMixin
 from edsl.language_models.RegisterLanguageModelsMeta import RegisterLanguageModelsMeta
-from edsl.language_models.KeyLookup import KeyLookup
-from edsl.exceptions.language_models import LanguageModelBadResponseError
-
-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.key_management.KeyLookupCollection import (
+    KeyLookupCollection,
+)
 
-            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
+from edsl.language_models.RawResponseHandler import RawResponseHandler
 
 
 def handle_key_error(func):
@@ -117,8 +75,21 @@ 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,
 ):
@@ -128,15 +99,22 @@ class LanguageModel(
     key_sequence = (
         None  # This should be something like ["choices", 0, "message", "content"]
     )
-    __rate_limits = None
-    _safety_factor = 0.8
+
+    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)
 
     def __init__(
         self,
-        tpm: float = None,
-        rpm: float = None,
+        tpm: Optional[float] = None,
+        rpm: Optional[float] = None,
         omit_system_prompt_if_empty_string: bool = True,
-        key_lookup: Optional[KeyLookup] = None,
+        key_lookup: Optional["KeyLookup"] = None,
         **kwargs,
     ):
         """Initialize the LanguageModel."""
@@ -147,7 +125,9 @@ class LanguageModel(
         self.remote = False
         self.omit_system_prompt_if_empty = omit_system_prompt_if_empty_string
 
-        # self._rpm / _tpm comes from the class
+        self.key_lookup = self._set_key_lookup(key_lookup)
+        self.model_info = self.key_lookup.get(self._inference_service_)
+
         if rpm is not None:
             self._rpm = rpm
 
@@ -161,49 +141,73 @@ class LanguageModel(
             if key not in parameters:
                 setattr(self, key, value)
 
-        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):
+        if 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:
-            self.key_lookup = key_lookup
+            return key_lookup
         else:
-            self.key_lookup = KeyLookup.from_os_environ()
+            klc = KeyLookupCollection()
+            klc.add_key_lookup(fetch_order=("config", "env"))
+            return klc.get(("config", "env"))
+
+    def set_key_lookup(self, key_lookup: "KeyLookup") -> None:
+        del self._api_token
+        self.key_lookup = key_lookup
+
+    def ask_question(self, question: "QuestionBase") -> str:
+        """Ask a question and return the response.
 
-    def ask_question(self, question):
+        :param question: The question to ask.
+        """
         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)
 
-    def set_key_lookup(self, key_lookup: KeyLookup) -> None:
-        del self._api_token
-        self.key_lookup = key_lookup
+    @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
 
     @property
     def api_token(self) -> str:
         if not hasattr(self, "_api_token"):
-            self._api_token = self.key_lookup.get_api_token(
-                self._inference_service_, self.remote
-            )
+            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
         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
@@ -232,7 +236,12 @@ class LanguageModel(
         return key_value is not None
 
     def __hash__(self) -> str:
-        """Allow the model to be used as a key in a dictionary."""
+        """Allow the model to be used as a key in a dictionary.
+
+        >>> m = LanguageModel.example()
+        >>> hash(m)
+        1811901442659237949
+        """
         from edsl.utilities.utilities import dict_hash
 
         return dict_hash(self.to_dict(add_edsl_version=False))
@@ -248,46 +257,6 @@ 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.
@@ -310,16 +279,7 @@ class LanguageModel(
 
     @abstractmethod
     async def async_execute_model_call(user_prompt: str, system_prompt: str):
-        """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'}], ...}
-        """
+        """Execute the model call and returns a coroutine."""
         pass
 
     async def remote_async_execute_model_call(
@@ -336,12 +296,7 @@ class LanguageModel(
 
     @jupyter_nb_handler
     def execute_model_call(self, *args, **kwargs) -> 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.")
-
-        """
+        """Execute the model call and returns the result as a coroutine."""
 
         async def main():
             results = await asyncio.gather(
@@ -353,58 +308,25 @@ class LanguageModel(
 
     @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)
+        """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)
 
     @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."
-            )
-        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
+        return cls.response_handler.get_usage_dict(raw_response)
 
     @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": 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)
+        return cls.response_handler.parse_response(raw_response)
 
     async def _async_get_intended_model_call_outcome(
         self,
@@ -421,6 +343,8 @@ 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.
@@ -463,6 +387,10 @@ 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
@@ -470,7 +398,6 @@ class LanguageModel(
             assert new_cache_key == cache_key  # should be the same
 
         cost = self.cost(response)
-
         return ModelResponse(
             response=response,
             cache_used=cache_used,
@@ -509,9 +436,9 @@ class LanguageModel(
 
         :param user_prompt: The user's prompt.
         :param system_prompt: The system's prompt.
-        :param iteration: The iteration number.
         :param cache: The cache to use.
-        :param encoded_image: The encoded image to use.
+        :param iteration: The iteration number.
+        :param files_list: The list of files to use.
 
         """
         params = {
@@ -525,8 +452,11 @@ class LanguageModel(
             params.update({"invigilator": kwargs["invigilator"]})
 
         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)
+        model_outputs: ModelResponse = (
+            await self._async_get_intended_model_call_outcome(**params)
+        )
+        edsl_dict: EDSLOutput = self.parse_response(model_outputs.response)
+
         agent_response_dict = AgentResponseDict(
             model_inputs=model_inputs,
             model_outputs=model_outputs,
@@ -537,55 +467,28 @@ 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."""
-
-        usage = self.get_usage_dict(raw_response)
-        from edsl.coop import Coop
+        """Return the dollar cost of a raw response.
 
-        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}"
+        :param raw_response: The raw response from the model.
+        """
 
-        return input_cost + output_cost
+        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,
+        )
 
     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'}
@@ -608,7 +511,7 @@ class LanguageModel(
         return model_class(**data)
 
     def __repr__(self) -> str:
-        """Return a string representation of the object."""
+        """Return a representation of the object."""
         param_string = ", ".join(
             f"{key} = {value}" for key, value in self.parameters.items()
         )
@@ -650,7 +553,7 @@ class LanguageModel(
         Exception report saved to ...
         Also see: ...
         """
-        from edsl import Model
+        from edsl.language_models.registry import Model
 
         if test_model:
             m = Model(

edsl/language_models/ModelList.py

@@ -1,12 +1,12 @@
 from typing import Optional, List
 from collections import UserList
-from edsl import Model
 
-from edsl.language_models import LanguageModel
 from edsl.Base import Base
-from edsl.utilities.decorators import add_edsl_version, remove_edsl_version
-from edsl.utilities.utilities import is_valid_variable_name
-from edsl.utilities.utilities import dict_hash
+from edsl.language_models.registry import Model
+
+# from edsl.language_models import LanguageModel
+from edsl.utilities.remove_edsl_version import remove_edsl_version
+from edsl.utilities.is_valid_variable_name import is_valid_variable_name
 
 
 class ModelList(Base, UserList):
@@ -40,7 +40,7 @@ class ModelList(Base, UserList):
         return f"ModelList({super().__repr__()})"
 
     def _summary(self):
-        return {"EDSL Class": "ModelList", "Number of Models": len(self)}
+        return {"models": len(self)}
 
     def __hash__(self):
         """Return a hash of the ModelList. This is used for comparison of ModelLists.
@@ -54,7 +54,8 @@ class ModelList(Base, UserList):
         return dict_hash(self.to_dict(sort=True, add_edsl_version=False))
 
     def to_scenario_list(self):
-        from edsl import ScenarioList, Scenario
+        from edsl.scenarios.ScenarioList import ScenarioList
+        from edsl.scenarios.Scenario import Scenario
 
         sl = ScenarioList()
         for model in self:
@@ -73,7 +74,7 @@ class ModelList(Base, UserList):
         pretty_labels: Optional[dict] = None,
     ):
         """
-        >>> ModelList.example().table("model")
+        >>> ModelList.example().table('model')
         model
         -------
         gpt-4o
@@ -112,11 +113,6 @@ class ModelList(Base, UserList):
 
         return d
 
-    def _repr_html_(self):
-        """Return an HTML representation of the ModelList."""
-        footer = f"<a href={self.__documentation__}>(docs)</a>"
-        return str(self.summary(format="html")) + footer
-
     @classmethod
     def from_names(self, *args, **kwargs):
         """A a model list from a list of names"""
@@ -133,6 +129,8 @@ class ModelList(Base, UserList):
         >>> newm = ModelList.from_dict(ModelList.example().to_dict())
         >>> assert ModelList.example() == newm
         """
+        from edsl.language_models.LanguageModel import LanguageModel
+
         return cls(data=[LanguageModel.from_dict(model) for model in data["models"]])
 
     def code(self):