edsl 0.1.15__py3-none-any.whl → 0.1.40__py3-none-any.whl

edsl/language_models/LanguageModel.py

@@ -1,310 +1,307 @@
-from __future__ import annotations
-from functools import wraps
-import io
-import asyncio
-import json
-import time
-import inspect
-from typing import Coroutine
-from abc import ABC, abstractmethod, ABCMeta
-from rich.console import Console
-from rich.table import Table
+"""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. 
 
-from edsl.trackers.TrackerAPI import TrackerAPI
-from queue import Queue
-from typing import Any, Callable, Type, List
-from edsl.data import CRUDOperations, CRUD
-from edsl.exceptions import LanguageModelResponseNotJSONError
-from edsl.language_models.schemas import model_prices
-from edsl.utilities.decorators import sync_wrapper, jupyter_nb_handler
+edsl_augmented_response: The JSON response from model, but augmented with EDSL-specific information, 
+such as the cache key, token usage, etc. 
 
-from edsl.language_models.repair import repair
-from typing import get_type_hints
+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 edsl.exceptions.language_models import LanguageModelAttributeTypeError
-from edsl.enums import LanguageModelType, InferenceServiceType
+"""
 
-from edsl.Base import RichPrintingMixin, PersistenceMixin
+from __future__ import annotations
+import warnings
+from functools import wraps
+import asyncio
+import json
+import os
+from typing import (
+    Coroutine,
+    Any,
+    Type,
+    Union,
+    List,
+    get_type_hints,
+    TypedDict,
+    Optional,
+    TYPE_CHECKING,
+)
+from abc import ABC, abstractmethod
+
+from edsl.data_transfer_models import (
+    ModelResponse,
+    ModelInputs,
+    EDSLOutput,
+    AgentResponseDict,
+)
+
+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.Base import PersistenceMixin, RepresentationMixin
+from edsl.language_models.RegisterLanguageModelsMeta import RegisterLanguageModelsMeta
+
+from edsl.language_models.key_management.KeyLookupCollection import (
+    KeyLookupCollection,
+)
+
+from edsl.language_models.RawResponseHandler import RawResponseHandler
 
 
 def handle_key_error(func):
+    """Handle KeyError exceptions."""
+
     @wraps(func)
     def wrapper(*args, **kwargs):
         try:
             return func(*args, **kwargs)
             assert True == False
         except KeyError as e:
-            # Handle the KeyError exception
             return f"""KeyError occurred: {e}. This is most likely because the model you are using 
             returned a JSON object we were not expecting."""
 
     return wrapper
 
 
-class RegisterLanguageModelsMeta(ABCMeta):
-    "Metaclass to register output elements in a registry i.e., those that have a parent"
-    _registry = {}  # Initialize the registry as a dictionary
-    REQUIRED_CLASS_ATTRIBUTES = ["_model_", "_parameters_", "_inference_service_"]
+class classproperty:
+    def __init__(self, method):
+        self.method = method
 
-    def __init__(cls, name, bases, dct):
-        super(RegisterLanguageModelsMeta, cls).__init__(name, bases, dct)
-        # if name != "LanguageModel":
-        if (model_name := getattr(cls, "_model_", None)) is not None:
-            RegisterLanguageModelsMeta.check_required_class_variables(
-                cls, RegisterLanguageModelsMeta.REQUIRED_CLASS_ATTRIBUTES
-            )
+    def __get__(self, instance, cls):
+        return self.method(cls)
 
-            ## Check that model name is valid
-            if not LanguageModelType.is_value_valid(model_name):
-                acceptable_values = [item.value for item in LanguageModelType]
-                raise LanguageModelAttributeTypeError(
-                    f"""A LanguageModel's model must be one of {LanguageModelType} values, which are
-                    {acceptable_values}. You passed {model_name}."""
-                )
 
-            if not InferenceServiceType.is_value_valid(
-                inference_service := getattr(cls, "_inference_service_", None)
-            ):
-                acceptable_values = [item.value for item in InferenceServiceType]
-                raise LanguageModelAttributeTypeError(
-                    f"""A LanguageModel's model must have an _inference_service_ value from 
-                    {acceptable_values}. You passed {inference_service}."""
-                )
+from edsl.Base import HashingMixin
 
-            # LanguageModel children have to implement the async_execute_model_call method
-            RegisterLanguageModelsMeta.verify_method(
-                candidate_class=cls,
-                method_name="async_execute_model_call",
-                expected_return_type=dict[str, Any],
-                required_parameters=[("user_prompt", str), ("system_prompt", str)],
-                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
-    def get_registered_classes(cls):
-        return cls._registry
+class LanguageModel(
+    PersistenceMixin,
+    RepresentationMixin,
+    HashingMixin,
+    ABC,
+    metaclass=RegisterLanguageModelsMeta,
+):
+    """ABC for Language Models."""
 
-    @staticmethod
-    def check_required_class_variables(
-        candidate_class: LanguageModel, required_attributes: List[str] = None
+    _model_ = None
+    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)
+
+    def __init__(
+        self,
+        tpm: Optional[float] = None,
+        rpm: Optional[float] = None,
+        omit_system_prompt_if_empty_string: bool = True,
+        key_lookup: Optional["KeyLookup"] = None,
+        **kwargs,
     ):
-        """Checks if a class has the required attributes
-        >>> class M:
-        ...     _model_ = "m"
-        ...     _parameters_ = {}
-        >>> RegisterLanguageModelsMeta.check_required_class_variables(M, ["_model_", "_parameters_"])
-        >>> class M2:
-        ...     _model_ = "m"
-        >>> RegisterLanguageModelsMeta.check_required_class_variables(M2, ["_model_", "_parameters_"])
-        Traceback (most recent call last):
-        ...
-        Exception: Class M2 does not have required attribute _parameters_
-        """
-        required_attributes = required_attributes or []
-        for attribute in required_attributes:
-            if not hasattr(candidate_class, attribute):
-                raise Exception(
-                    f"Class {candidate_class.__name__} does not have required attribute {attribute}"
-                )
+        """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
 
-    @staticmethod
-    def verify_method(
-        candidate_class: LanguageModel,
-        method_name: str,
-        expected_return_type: Any,
-        required_parameters: List[tuple[str, Any]] = None,
-        must_be_async: bool = False,
-    ):
-        RegisterLanguageModelsMeta._check_method_defined(candidate_class, method_name)
+        self.key_lookup = self._set_key_lookup(key_lookup)
+        self.model_info = self.key_lookup.get(self._inference_service_)
 
-        required_parameters = required_parameters or []
-        method = getattr(candidate_class, method_name)
-        signature = inspect.signature(method)
+        if rpm is not None:
+            self._rpm = rpm
 
-        RegisterLanguageModelsMeta._check_return_type(method, expected_return_type)
+        if tpm is not None:
+            self._tpm = tpm
 
-        if must_be_async:
-            RegisterLanguageModelsMeta._check_is_coroutine(method)
+        for key, value in parameters.items():
+            setattr(self, key, value)
 
-        # Check the parameters
-        params = signature.parameters
-        for param_name, param_type in required_parameters:
-            RegisterLanguageModelsMeta._verify_parameter(
-                params, param_name, param_type, method_name
-            )
+        for key, value in kwargs.items():
+            if key not in parameters:
+                setattr(self, key, value)
 
-    @staticmethod
-    def _check_method_defined(cls, method_name):
-        """Checks if a method is defined in a class
-        >>> class M:
-        ...     def f(self): pass
-        >>> RegisterLanguageModelsMeta._check_method_defined(M, "f")
-        >>> RegisterLanguageModelsMeta._check_method_defined(M, "g")
-        Traceback (most recent call last):
-        ...
-        NotImplementedError: g method must be implemented
-        """
-        if not hasattr(cls, method_name):
-            raise NotImplementedError(f"{method_name} method must be implemented")
+        if kwargs.get("skip_api_key_check", False):
+            # Skip the API key check. Sometimes this is useful for testing.
+            self._api_token = None
 
-    @staticmethod
-    def _check_is_coroutine(func: Callable):
-        """
-        Checks to make sure it's a coroutine function
-        >>> def f(): pass
-        >>> RegisterLanguageModelsMeta._check_is_coroutine(f)
-        Traceback (most recent call last):
-        ...
-        TypeError: A LangugeModel class with method f must be an asynchronous method
-        """
-        if not inspect.iscoroutinefunction(func):
-            raise TypeError(
-                f"A LangugeModel class with method {func.__name__} must be an asynchronous method"
-            )
+    def _set_key_lookup(self, key_lookup: "KeyLookup") -> "KeyLookup":
+        """Set the key lookup."""
+        if key_lookup is not None:
+            return key_lookup
+        else:
+            klc = KeyLookupCollection()
+            klc.add_key_lookup(fetch_order=("config", "env"))
+            return klc.get(("config", "env"))
 
-    @staticmethod
-    def _verify_parameter(params, param_name, param_type, method_name):
-        if param_name not in params:
-            raise TypeError(
-                f"""Parameter "{param_name}" of method "{method_name}" must be defined.
-                """
-            )
-        if params[param_name].annotation != param_type:
-            raise TypeError(
-                f"""Parameter "{param_name}" of method "{method_name}" must be of type {param_type.__name__}.
-                Got {params[param_name].annotation} instead.
-                """
-            )
+    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
 
-    @staticmethod
-    def _check_return_type(method, expected_return_type):
-        """
-        Checks if the return type of a method is as expected
-        >>> 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'>
+    def ask_question(self, question: "QuestionBase") -> str:
+        """Ask a question and return the response.
+
+        :param question: The question to ask.
         """
-        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}."
-                )
+        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)
 
-    @classmethod
-    def model_names_to_classes(cls):
-        d = {}
-        for classname, cls in cls._registry.items():
-            if hasattr(cls, "_model_"):
-                d[cls._model_] = cls
+    @property
+    def rpm(self):
+        if not hasattr(self, "_rpm"):
+            if self.model_info is None:
+                self._rpm = self.DEFAULT_RPM
             else:
-                raise Exception(
-                    f"Class {classname} does not have a _model_ class attribute"
-                )
-        return d
+                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
 
-class LanguageModel(
-    RichPrintingMixin, PersistenceMixin, ABC, metaclass=RegisterLanguageModelsMeta
-):
-    """ABC for LLM subclasses."""
+    # in case we want to override the default values
+    @tpm.setter
+    def tpm(self, value):
+        self._tpm = value
 
-    _model_ = None
+    @rpm.setter
+    def rpm(self, value):
+        self._rpm = value
 
-    __rate_limits = None
-    # TODO: Use the OpenAI Teir 1 rate limits
-    __default_rate_limits = {"rpm": 10_000, "tpm": 2_000_000}
-    _safety_factor = 0.8
+    @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
+        return self._api_token
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def hello(self, verbose=False):
+        """Runs a simple test to check if the model is working."""
+        token = self.api_token
+        masked = token[: min(8, len(token))] + "..."
+        if verbose:
+            print(f"Current key is {masked}")
+        return self.execute_model_call(
+            user_prompt="Hello, model!", system_prompt="You are a helpful agent."
+        )
 
-    def __init__(self, crud: CRUDOperations = CRUD, **kwargs):
-        """
-        Attributes:
-        - all attributes inherited from subclasses
-        - lock: lock for this model to ensure TODO
-        - api_queue: queue that records messages about API calls the model makes. Used by `InterviewManager` to update details about state of model.
+    def has_valid_api_key(self) -> bool:
+        """Check if the model has a valid API key.
+
+        >>> LanguageModel.example().has_valid_api_key() : # doctest: +SKIP
+        True
+
+        This method is used to check if the model has a valid API key.
         """
-        self.model = getattr(self, "_model_", None)
-        default_parameters = getattr(self, "_parameters_", None)
-        parameters = self._overide_default_parameters(kwargs, default_parameters)
-        self.parameters = parameters
+        from edsl.enums import service_to_api_keyname
 
-        for key, value in parameters.items():
-            setattr(self, key, value)
+        if self._model_ == "test":
+            return True
 
-        for key, value in kwargs.items():
-            if key not in parameters:
-                setattr(self, key, value)
+        key_name = service_to_api_keyname.get(self._inference_service_, "NOT FOUND")
+        key_value = os.getenv(key_name)
+        return key_value is not None
 
-        # TODO: This can very likely be removed
-        self.api_queue = Queue()
-        self.crud = crud
+    def __hash__(self) -> str:
+        """Allow the model to be used as a key in a dictionary.
 
-    def __hash__(self):
-        "Allows the model to be used as a key in a dictionary"
-        return hash(self.model + str(self.parameters))
+        >>> m = LanguageModel.example()
+        >>> hash(m)
+        1811901442659237949
+        """
+        from edsl.utilities.utilities import dict_hash
 
-    def __eq__(self, other):
-        return self.model == other.model and self.parameters == other.parameters
+        return dict_hash(self.to_dict(add_edsl_version=False))
 
-    def _set_rate_limits(self) -> None:
-        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
+    def __eq__(self, other) -> bool:
+        """Check is two models are the same.
 
-    @property
-    def RPM(self):
-        "Model's requests-per-minute limit"
-        self._set_rate_limits()
-        return self._safety_factor * self.__rate_limits["rpm"]
+        >>> m1 = LanguageModel.example()
+        >>> m2 = LanguageModel.example()
+        >>> m1 == m2
+        True
 
-    @property
-    def TPM(self):
-        "Model's tokens-per-minute limit"
-        self._set_rate_limits()
-        return self._safety_factor * self.__rate_limits["tpm"]
+        """
+        return self.model == other.model and self.parameters == other.parameters
 
     @staticmethod
     def _overide_default_parameters(passed_parameter_dict, default_parameter_dict):
-        """Returns a dictionary of parameters, with passed parameters taking precedence over defaults.
+        """Return a dictionary of parameters, with passed parameters taking precedence over defaults.
 
         >>> LanguageModel._overide_default_parameters(passed_parameter_dict={"temperature": 0.5}, default_parameter_dict={"temperature":0.9})
         {'temperature': 0.5}
         >>> 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({})
-        for parameter, default_value in default_parameter_dict.items():
-            if parameter in passed_parameter_dict:
-                parameters[parameter] = passed_parameter_dict[parameter]
-            else:
-                parameters[parameter] = default_value
-        return parameters
+        # this is the case when data is loaded from a dict after serialization
+        if "parameters" in passed_parameter_dict:
+            passed_parameter_dict = passed_parameter_dict["parameters"]
+        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):
+        return self.execute_model_call(user_prompt, system_prompt)
 
     @abstractmethod
-    async def async_execute_model_call():
+    async def async_execute_model_call(user_prompt: str, system_prompt: str):
+        """Execute the model call and returns a coroutine."""
         pass
 
+    async def remote_async_execute_model_call(
+        self, user_prompt: str, system_prompt: str
+    ):
+        """Execute the model call and returns the result as a coroutine, using Coop."""
+        from edsl.coop import Coop
+
+        client = Coop()
+        response_data = await client.remote_async_execute_model_call(
+            self.to_dict(), user_prompt, system_prompt
+        )
+        return response_data
+
     @jupyter_nb_handler
     def execute_model_call(self, *args, **kwargs) -> Coroutine:
+        """Execute the model call and returns the result as a coroutine."""
+
         async def main():
             results = await asyncio.gather(
                 self.async_execute_model_call(*args, **kwargs)
@@ -313,193 +310,317 @@ class LanguageModel(
 
         return main()
 
-    @abstractmethod
-    def parse_response(raw_response: dict[str, Any]) -> str:
-        """Parses the API response and returns the response text.
-        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:
-
-        {
-            "id": "chatcmpl-8eORaeuVb4po9WQRjKEFY6w7v6cTm",
-            "choices": [
-                {
-                    "finish_reason": "stop",
-                    "index": 0,
-                    "logprobs": None,
-                    "message": {
-                        "content": "Hello! How can I assist you today? If you have any questions or need information on a particular topic, feel free to ask.",
-                        "role": "assistant",
-                        "function_call": None,
-                        "tool_calls": None,
-                    },
-                }
-            ],
-            "created": 1704637774,
-            "model": "gpt-4-1106-preview",
-            "object": "chat.completion",
-            "system_fingerprint": "fp_168383a679",
-            "usage": {"completion_tokens": 27, "prompt_tokens": 13, "total_tokens": 40},
-        }
+    @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'
 
-        To actually tract the response, we need to grab
-            data["choices[0]"]["message"]["content"].
         """
-        raise NotImplementedError
+        return cls.response_handler.get_generated_token_string(raw_response)
 
-    def _update_response_with_tracking(
-        self, response, start_time, cached_response=False
-    ):
-        end_time = time.time()
-        response["elapsed_time"] = end_time - start_time
-        response["timestamp"] = end_time
-        self._post_tracker_event(response)
-        response["cached_response"] = cached_response
-        return response
-
-    async def async_get_raw_response(
-        self, user_prompt: str, system_prompt: str = ""
-    ) -> dict[str, Any]:
-        """This is some middle-ware that handles the caching of responses.
-        If the cache isn't being used, it just returns a 'fresh' call to the LLM,
-        but appends some tracking information to the response (using the _update_response_with_tracking method).
+    @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)
+
+    @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)
+
+    async def _async_get_intended_model_call_outcome(
+        self,
+        user_prompt: str,
+        system_prompt: str,
+        cache: Cache,
+        iteration: int = 0,
+        files_list: Optional[List[FileStore]] = None,
+        invigilator=None,
+    ) -> ModelResponse:
+        """Handle caching of responses.
+
+        :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 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.
         If it is, it returns the cached response, but again appends some tracking information.
         If it isn't, it calls the LLM, saves the response to the database, and returns the response with tracking information.
 
         If self.use_cache is True, then attempts to retrieve the response from the database;
-        if not in the DB, calls the LLM and writes the response to the DB."""
-        start_time = time.time()
-
-        if not self.use_cache:
-            response = await self.async_execute_model_call(user_prompt, system_prompt)
-            return self._update_response_with_tracking(response, start_time, False)
-
-        cached_response = self.crud.get_LLMOutputData(
-            model=str(self.model),
-            parameters=str(self.parameters),
-            system_prompt=system_prompt,
-            prompt=user_prompt,
-        )
+        if not in the DB, calls the LLM and writes the response to the DB.
+
+        >>> 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(...)"""
 
-        if cached_response:
+        if files_list:
+            files_hash = "+".join([str(hash(file)) for file in files_list])
+            user_prompt_with_hashes = user_prompt + f" {files_hash}"
+        else:
+            user_prompt_with_hashes = user_prompt
+
+        cache_call_params = {
+            "model": str(self.model),
+            "parameters": self.parameters,
+            "system_prompt": system_prompt,
+            "user_prompt": user_prompt_with_hashes,
+            "iteration": iteration,
+        }
+        cached_response, cache_key = cache.fetch(**cache_call_params)
+
+        if cache_used := cached_response is not None:
             response = json.loads(cached_response)
-            cache_used = True
         else:
-            response = await self.async_execute_model_call(user_prompt, system_prompt)
-            self._save_response_to_db(user_prompt, system_prompt, response)
-            cache_used = False
+            f = (
+                self.remote_async_execute_model_call
+                if hasattr(self, "remote") and self.remote
+                else self.async_execute_model_call
+            )
+            params = {
+                "user_prompt": user_prompt,
+                "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
+            )  # 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 self._update_response_with_tracking(response, start_time, cache_used)
+    _get_intended_model_call_outcome = sync_wrapper(
+        _async_get_intended_model_call_outcome
+    )
 
-    get_raw_response = sync_wrapper(async_get_raw_response)
+    def simple_ask(
+        self,
+        question: QuestionBase,
+        system_prompt="You are a helpful agent pretending to be a human.",
+        top_logprobs=2,
+    ):
+        """Ask a question and return the response."""
+        self.logprobs = True
+        self.top_logprobs = top_logprobs
+        return self.execute_model_call(
+            user_prompt=question.human_readable(), system_prompt=system_prompt
+        )
 
-    def _save_response_to_db(self, prompt, system_prompt, response):
-        try:
-            output = json.dumps(response)
-        except json.JSONDecodeError:
-            raise LanguageModelResponseNotJSONError
-        self.crud.write_LLMOutputData(
-            model=str(self.model),
-            parameters=str(self.parameters),
-            system_prompt=system_prompt,
-            prompt=prompt,
-            output=output,
+    async def async_get_response(
+        self,
+        user_prompt: str,
+        system_prompt: str,
+        cache: Cache,
+        iteration: int = 1,
+        files_list: Optional[List[FileStore]] = None,
+        **kwargs,
+    ) -> dict:
+        """Get response, parse, and return as string.
+
+        :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.
+
+        """
+        params = {
+            "user_prompt": user_prompt,
+            "system_prompt": system_prompt,
+            "iteration": iteration,
+            "cache": cache,
+            "files_list": files_list,
+        }
+        if "invigilator" in kwargs:
+            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)
 
-    async def async_get_response(self, user_prompt: str, system_prompt: str = ""):
-        """Get response, parse, and return as string."""
-        raw_response = await self.async_get_raw_response(user_prompt, system_prompt)
-        response = self.parse_response(raw_response)
-        try:
-            dict_response = json.loads(response)
-        except json.JSONDecodeError as e:
-            # TODO: Turn into logs to generate issues
-            dict_response, success = await repair(response, str(e))
-            if not success:
-                raise Exception("Even the repair failed.")
-
-        dict_response["cached_response"] = raw_response["cached_response"]
-        dict_response["usage"] = raw_response.get("usage", {})
-        dict_response["raw_model_response"] = raw_response
-        return dict_response
+        agent_response_dict = AgentResponseDict(
+            model_inputs=model_inputs,
+            model_outputs=model_outputs,
+            edsl_dict=edsl_dict,
+        )
+        return agent_response_dict
 
     get_response = sync_wrapper(async_get_response)
 
-    #######################
-    # USEFUL METHODS
-    #######################
-    def _post_tracker_event(self, raw_response: dict[str, Any]) -> None:
-        """Parses the API response and sends usage details to the API Queue."""
-        usage = raw_response.get("usage", {})
-        usage.update(
-            {
-                "cached_response": raw_response.get("cached_response", None),
-                "elapsed_time": raw_response.get("elapsed_time", None),
-                "timestamp": raw_response.get("timestamp", None),
-            }
+    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.
+        """
+
+        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,
         )
-        event = TrackerAPI.APICallDetails(details=usage)
-        self.api_queue.put(event)
-
-    def cost(self, raw_response: dict[str, Any]) -> float:
-        """Returns the dollar cost of a raw response."""
-        keys = raw_response["usage"].keys()
-        prices = model_prices.get(self.model)
-        return sum([prices.get(key, 0.0) * raw_response["usage"][key] for key in keys])
-
-    #######################
-    # SERIALIZATION METHODS
-    #######################
-    def to_dict(self) -> dict[str, Any]:
-        """Converts instance to a dictionary."""
-        return {"model": self.model, "parameters": self.parameters}
+
+    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,
+        }
+        if add_edsl_version:
+            from edsl import __version__
+
+            d["edsl_version"] = __version__
+            d["edsl_class_name"] = self.__class__.__name__
+        return d
 
     @classmethod
+    @remove_edsl_version
     def from_dict(cls, data: dict) -> Type[LanguageModel]:
-        """Converts dictionary to a LanguageModel child instance."""
-        from edsl.language_models.registry import get_model_class
+        """Convert dictionary to a LanguageModel child instance."""
+        from edsl.language_models.model import get_model_class
 
         model_class = get_model_class(data["model"])
-        data["use_cache"] = True
         return model_class(**data)
 
-    #######################
-    # DUNDER METHODS
-    #######################
     def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(model = '{self.model}', parameters={self.parameters})"
+        """Return a representation of the object."""
+        param_string = ", ".join(
+            f"{key} = {value}" for key, value in self.parameters.items()
+        )
+        return (
+            f"Model(model_name = '{self.model}'"
+            + (f", {param_string}" if param_string else "")
+            + ")"
+        )
 
     def __add__(self, other_model: Type[LanguageModel]) -> Type[LanguageModel]:
-        """Combine two models into a single model (other_model takes precedence over self)"""
-        print(
+        """Combine two models into a single model (other_model takes precedence over self)."""
+        import warnings
+
+        warnings.warn(
             f"""Warning: one model is replacing another. If you want to run both models, use a single `by` e.g., 
               by(m1, m2, m3) not by(m1).by(m2).by(m3)."""
         )
         return other_model or self
 
-    def rich_print(self):
-        """Displays an object as a table."""
-        table = Table(title="Language Model")
-        table.add_column("Attribute", style="bold")
-        table.add_column("Value")
+    @classmethod
+    def example(
+        cls,
+        test_model: bool = False,
+        canned_response: str = "Hello world",
+        throw_exception: bool = False,
+    ) -> LanguageModel:
+        """Return a default instance of the class.
+
+        >>> from edsl.language_models import LanguageModel
+        >>> m = LanguageModel.example(test_model = True, canned_response = "WOWZA!")
+        >>> isinstance(m, LanguageModel)
+        True
+        >>> from edsl import QuestionFreeText
+        >>> q = QuestionFreeText(question_text = "What is your name?", question_name = 'example')
+        >>> q.by(m).run(cache = False, disable_remote_cache = True, disable_remote_inference = True).select('example').first()
+        'WOWZA!'
+        >>> m = LanguageModel.example(test_model = True, canned_response = "WOWZA!", throw_exception = True)
+        >>> r = q.by(m).run(cache = False, disable_remote_cache = True, disable_remote_inference = True, print_exceptions = True)
+        Exception report saved to ...
+        Also see: ...
+        """
+        from edsl.language_models.model import Model
 
-        to_display = self.__dict__.copy()
-        for attr_name, attr_value in to_display.items():
-            table.add_row(attr_name, repr(attr_value))
+        if test_model:
+            m = Model(
+                "test", canned_response=canned_response, throw_exception=throw_exception
+            )
+            return m
+        else:
+            return Model(skip_api_key_check=True)
 
-        return table
+    def from_cache(self, cache: "Cache") -> LanguageModel:
 
-    @classmethod
-    def example(cls):
-        "Returns a default instance of the class"
-        from edsl import Model
+        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 Model(Model.available()[0])
+        return new_instance
 
 
 if __name__ == "__main__":
-    # import doctest
-    # doctest.testmod()
-    from edsl.language_models import LanguageModel
+    """Run the module's test suite."""
+    import doctest
 
-    print(LanguageModel.example())
+    doctest.testmod(optionflags=doctest.ELLIPSIS)