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

edsl/inference_services/MistralAIService.py

@@ -0,0 +1,120 @@
+import os
+from typing import Any, List
+from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
+from edsl.language_models.LanguageModel import LanguageModel
+import asyncio
+from mistralai import Mistral
+
+from edsl.exceptions.language_models import LanguageModelBadResponseError
+
+
+class MistralAIService(InferenceServiceABC):
+    """Mistral AI service class."""
+
+    key_sequence = ["choices", 0, "message", "content"]
+    usage_sequence = ["usage"]
+
+    _inference_service_ = "mistral"
+    _env_key_name_ = "MISTRAL_API_KEY"  # Environment variable for Mistral API key
+    input_token_name = "prompt_tokens"
+    output_token_name = "completion_tokens"
+
+    _sync_client_instance = None
+    _async_client_instance = None
+
+    _sync_client = Mistral
+    _async_client = Mistral
+
+    _models_list_cache: List[str] = []
+    model_exclude_list = []
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        # so subclasses have to create their own instances of the clients
+        cls._sync_client_instance = None
+        cls._async_client_instance = None
+
+    @classmethod
+    def sync_client(cls):
+        if cls._sync_client_instance is None:
+            cls._sync_client_instance = cls._sync_client(
+                api_key=os.getenv(cls._env_key_name_)
+            )
+        return cls._sync_client_instance
+
+    @classmethod
+    def async_client(cls):
+        if cls._async_client_instance is None:
+            cls._async_client_instance = cls._async_client(
+                api_key=os.getenv(cls._env_key_name_)
+            )
+        return cls._async_client_instance
+
+    @classmethod
+    def available(cls) -> list[str]:
+        if not cls._models_list_cache:
+            cls._models_list_cache = [
+                m.id for m in cls.sync_client().models.list().data
+            ]
+
+        return cls._models_list_cache
+
+    @classmethod
+    def create_model(
+        cls, model_name: str = "mistral", model_class_name=None
+    ) -> LanguageModel:
+        if model_class_name is None:
+            model_class_name = cls.to_class_name(model_name)
+
+        class LLM(LanguageModel):
+            """
+            Child class of LanguageModel for interacting with Mistral models.
+            """
+
+            key_sequence = cls.key_sequence
+            usage_sequence = cls.usage_sequence
+
+            input_token_name = cls.input_token_name
+            output_token_name = cls.output_token_name
+
+            _inference_service_ = cls._inference_service_
+            _model_ = model_name
+            _parameters_ = {
+                "temperature": 0.5,
+                "max_tokens": 512,
+                "top_p": 0.9,
+            }
+
+            _tpm = cls.get_tpm(cls)
+            _rpm = cls.get_rpm(cls)
+
+            def sync_client(self):
+                return cls.sync_client()
+
+            def async_client(self):
+                return cls.async_client()
+
+            async def async_execute_model_call(
+                self, user_prompt: str, system_prompt: str = ""
+            ) -> dict[str, Any]:
+                """Calls the Mistral API and returns the API response."""
+                s = self.async_client()
+
+                try:
+                    res = await s.chat.complete_async(
+                        model=model_name,
+                        messages=[
+                            {
+                                "content": user_prompt,
+                                "role": "user",
+                            },
+                        ],
+                    )
+                except Exception as e:
+                    raise LanguageModelBadResponseError(f"Error with Mistral API: {e}")
+
+                return res.model_dump()
+
+        LLM.__name__ = model_class_name
+
+        return LLM

edsl/inference_services/OllamaService.py

@@ -0,0 +1,18 @@
+import aiohttp
+import json
+import requests
+from typing import Any, List
+
+# from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
+from edsl.language_models import LanguageModel
+
+from edsl.inference_services.OpenAIService import OpenAIService
+
+
+class OllamaService(OpenAIService):
+    """DeepInfra service class."""
+
+    _inference_service_ = "ollama"
+    _env_key_name_ = "DEEP_INFRA_API_KEY"
+    _base_url_ = "http://localhost:11434/v1"
+    _models_list_cache: List[str] = []

edsl/inference_services/OpenAIService.py

@@ -1,12 +1,15 @@
-from typing import Any, List
-import re
+from __future__ import annotations
+from typing import Any, List, Optional
 import os
-#from openai import AsyncOpenAI
+
 import openai
 
 from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
 from edsl.language_models import LanguageModel
 from edsl.inference_services.rate_limits_cache import rate_limits
+from edsl.utilities.utilities import fix_partial_correct_response
+
+from edsl.config import CONFIG
 
 
 class OpenAIService(InferenceServiceABC):
@@ -18,20 +21,37 @@ class OpenAIService(InferenceServiceABC):
 
     _sync_client_ = openai.OpenAI
     _async_client_ = openai.AsyncOpenAI
-    
+
+    _sync_client_instance = None
+    _async_client_instance = None
+
+    key_sequence = ["choices", 0, "message", "content"]
+    usage_sequence = ["usage"]
+    input_token_name = "prompt_tokens"
+    output_token_name = "completion_tokens"
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        # so subclasses have to create their own instances of the clients
+        cls._sync_client_instance = None
+        cls._async_client_instance = None
+
     @classmethod
     def sync_client(cls):
-        return cls._sync_client_(
-            api_key = os.getenv(cls._env_key_name_), 
-            base_url = cls._base_url_)
-    
+        if cls._sync_client_instance is None:
+            cls._sync_client_instance = cls._sync_client_(
+                api_key=os.getenv(cls._env_key_name_), base_url=cls._base_url_
+            )
+        return cls._sync_client_instance
+
     @classmethod
     def async_client(cls):
-        return cls._async_client_(
-            api_key = os.getenv(cls._env_key_name_), 
-            base_url = cls._base_url_)
+        if cls._async_client_instance is None:
+            cls._async_client_instance = cls._async_client_(
+                api_key=os.getenv(cls._env_key_name_), base_url=cls._base_url_
+            )
+        return cls._async_client_instance
 
-    # TODO: Make this a coop call
     model_exclude_list = [
         "whisper-1",
         "davinci-002",
@@ -46,6 +66,8 @@ class OpenAIService(InferenceServiceABC):
         "text-embedding-3-small",
         "text-embedding-ada-002",
         "ft:davinci-002:mit-horton-lab::8OfuHgoo",
+        "gpt-3.5-turbo-instruct-0914",
+        "gpt-3.5-turbo-instruct",
     ]
     _models_list_cache: List[str] = []
 
@@ -59,27 +81,15 @@ class OpenAIService(InferenceServiceABC):
 
     @classmethod
     def available(cls) -> List[str]:
-        #from openai import OpenAI
-
         if not cls._models_list_cache:
             try:
-                #client = OpenAI(api_key = os.getenv(cls._env_key_name_), base_url = cls._base_url_)
                 cls._models_list_cache = [
                     m.id
-                    for m in cls.get_model_list() 
+                    for m in cls.get_model_list()
                     if m.id not in cls.model_exclude_list
                 ]
             except Exception as e:
                 raise
-                # print(
-                #     f"""Error retrieving models: {e}.
-                #     See instructions about storing your API keys: https://docs.expectedparrot.com/en/latest/api_keys.html"""
-                # )
-                # cls._models_list_cache = [
-                #     "gpt-3.5-turbo",
-                #     "gpt-4-1106-preview",
-                #     "gpt-4",
-                # ]  # Fallback list
         return cls._models_list_cache
 
     @classmethod
@@ -92,6 +102,14 @@ class OpenAIService(InferenceServiceABC):
             Child class of LanguageModel for interacting with OpenAI models
             """
 
+            key_sequence = cls.key_sequence
+            usage_sequence = cls.usage_sequence
+            input_token_name = cls.input_token_name
+            output_token_name = cls.output_token_name
+
+            _rpm = cls.get_rpm(cls)
+            _tpm = cls.get_tpm(cls)
+
             _inference_service_ = cls._inference_service_
             _model_ = model_name
             _parameters_ = {
@@ -106,21 +124,15 @@ class OpenAIService(InferenceServiceABC):
 
             def sync_client(self):
                 return cls.sync_client()
-            
+
             def async_client(self):
                 return cls.async_client()
 
             @classmethod
             def available(cls) -> list[str]:
-                #import openai
-                #client = openai.OpenAI(api_key = os.getenv(cls._env_key_name_), base_url = cls._base_url_)
-                #return client.models.list()
                 return cls.sync_client().models.list()
-            
-            def get_headers(self) -> dict[str, Any]:
-                #from openai import OpenAI
 
-                #client = OpenAI(api_key = os.getenv(cls._env_key_name_), base_url = cls._base_url_)
+            def get_headers(self) -> dict[str, Any]:
                 client = self.sync_client()
                 response = client.chat.completions.with_raw_response.create(
                     messages=[
@@ -157,6 +169,9 @@ class OpenAIService(InferenceServiceABC):
                 user_prompt: str,
                 system_prompt: str = "",
                 encoded_image=None,
+                invigilator: Optional[
+                    "InvigilatorAI"
+                ] = None,  # TBD - can eventually be used for function-calling
             ) -> dict[str, Any]:
                 """Calls the OpenAI API and returns the API response."""
                 if encoded_image:
@@ -171,17 +186,16 @@ class OpenAIService(InferenceServiceABC):
                     )
                 else:
                     content = user_prompt
-                # self.client = AsyncOpenAI(
-                #     api_key = os.getenv(cls._env_key_name_), 
-                #     base_url = cls._base_url_
-                #     )
                 client = self.async_client()
+                messages = [
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": content},
+                ]
+                if system_prompt == "" and self.omit_system_prompt_if_empty:
+                    messages = messages[1:]
                 params = {
                     "model": self.model,
-                    "messages": [
-                        {"role": "system", "content": system_prompt},
-                        {"role": "user", "content": content},
-                    ],
+                    "messages": messages,
                     "temperature": self.temperature,
                     "max_tokens": self.max_tokens,
                     "top_p": self.top_p,
@@ -193,21 +207,6 @@ class OpenAIService(InferenceServiceABC):
                 response = await client.chat.completions.create(**params)
                 return response.model_dump()
 
-            @staticmethod
-            def parse_response(raw_response: dict[str, Any]) -> str:
-                """Parses the API response and returns the response text."""
-                try:
-                    response = raw_response["choices"][0]["message"]["content"]
-                except KeyError:
-                    print("Tried to parse response but failed:")
-                    print(raw_response)
-                pattern = r"^```json(?:\\n|\n)(.+?)(?:\\n|\n)```$"
-                match = re.match(pattern, response, re.DOTALL)
-                if match:
-                    return match.group(1)
-                else:
-                    return response
-
         LLM.__name__ = "LanguageModel"
 
         return LLM

edsl/inference_services/TestService.py

@@ -0,0 +1,80 @@
+from typing import Any, List
+import os
+import asyncio
+from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
+from edsl.language_models import LanguageModel
+from edsl.inference_services.rate_limits_cache import rate_limits
+from edsl.utilities.utilities import fix_partial_correct_response
+
+from edsl.enums import InferenceServiceType
+import random
+
+
+class TestService(InferenceServiceABC):
+    """OpenAI service class."""
+
+    _inference_service_ = "test"
+    _env_key_name_ = None
+    _base_url_ = None
+
+    _sync_client_ = None
+    _async_client_ = None
+
+    _sync_client_instance = None
+    _async_client_instance = None
+
+    key_sequence = None
+    usage_sequence = None
+    model_exclude_list = []
+    input_token_name = "prompt_tokens"
+    output_token_name = "completion_tokens"
+
+    @classmethod
+    def available(cls) -> list[str]:
+        return ["test"]
+
+    @classmethod
+    def create_model(cls, model_name, model_class_name=None) -> LanguageModel:
+        throw_exception = False
+
+        class TestServiceLanguageModel(LanguageModel):
+            _model_ = "test"
+            _parameters_ = {"temperature": 0.5}
+            _inference_service_ = InferenceServiceType.TEST.value
+            usage_sequence = ["usage"]
+            key_sequence = ["message", 0, "text"]
+            input_token_name = cls.input_token_name
+            output_token_name = cls.output_token_name
+            _rpm = 1000
+            _tpm = 100000
+
+            @property
+            def _canned_response(self):
+                if hasattr(self, "canned_response"):
+                    return self.canned_response
+                else:
+                    return "Hello, world"
+
+            async def async_execute_model_call(
+                self,
+                user_prompt: str,
+                system_prompt: str,
+                encoded_image=None,
+            ) -> dict[str, Any]:
+                await asyncio.sleep(0.1)
+                # return {"message": """{"answer": "Hello, world"}"""}
+
+                if hasattr(self, "throw_exception") and self.throw_exception:
+                    if hasattr(self, "exception_probability"):
+                        p = self.exception_probability
+                    else:
+                        p = 1
+
+                    if random.random() < p:
+                        raise Exception("This is a test error")
+                return {
+                    "message": [{"text": f"{self._canned_response}"}],
+                    "usage": {"prompt_tokens": 1, "completion_tokens": 1},
+                }
+
+        return TestServiceLanguageModel

edsl/inference_services/TogetherAIService.py

@@ -0,0 +1,170 @@
+import aiohttp
+import json
+import requests
+from typing import Any, List
+
+# from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
+from edsl.language_models import LanguageModel
+
+from edsl.inference_services.OpenAIService import OpenAIService
+import openai
+
+
+class TogetherAIService(OpenAIService):
+    """DeepInfra service class."""
+
+    _inference_service_ = "together"
+    _env_key_name_ = "TOGETHER_API_KEY"
+    _base_url_ = "https://api.together.xyz/v1"
+    _models_list_cache: List[str] = []
+
+    # These are non-serverless models. There was no api param to filter them
+    model_exclude_list = [
+        "EleutherAI/llemma_7b",
+        "HuggingFaceH4/zephyr-7b-beta",
+        "Nexusflow/NexusRaven-V2-13B",
+        "NousResearch/Hermes-2-Theta-Llama-3-70B",
+        "NousResearch/Nous-Capybara-7B-V1p9",
+        "NousResearch/Nous-Hermes-13b",
+        "NousResearch/Nous-Hermes-2-Mistral-7B-DPO",
+        "NousResearch/Nous-Hermes-2-Mixtral-8x7B-SFT",
+        "NousResearch/Nous-Hermes-Llama2-13b",
+        "NousResearch/Nous-Hermes-Llama2-70b",
+        "NousResearch/Nous-Hermes-llama-2-7b",
+        "NumbersStation/nsql-llama-2-7B",
+        "Open-Orca/Mistral-7B-OpenOrca",
+        "Phind/Phind-CodeLlama-34B-Python-v1",
+        "Phind/Phind-CodeLlama-34B-v2",
+        "Qwen/Qwen1.5-0.5B",
+        "Qwen/Qwen1.5-0.5B-Chat",
+        "Qwen/Qwen1.5-1.8B",
+        "Qwen/Qwen1.5-1.8B-Chat",
+        "Qwen/Qwen1.5-14B",
+        "Qwen/Qwen1.5-14B-Chat",
+        "Qwen/Qwen1.5-32B",
+        "Qwen/Qwen1.5-32B-Chat",
+        "Qwen/Qwen1.5-4B",
+        "Qwen/Qwen1.5-4B-Chat",
+        "Qwen/Qwen1.5-72B",
+        "Qwen/Qwen1.5-7B",
+        "Qwen/Qwen1.5-7B-Chat",
+        "Qwen/Qwen2-1.5B",
+        "Qwen/Qwen2-1.5B-Instruct",
+        "Qwen/Qwen2-72B",
+        "Qwen/Qwen2-7B",
+        "Qwen/Qwen2-7B-Instruct",
+        "SG161222/Realistic_Vision_V3.0_VAE",
+        "Snowflake/snowflake-arctic-instruct",
+        "Undi95/ReMM-SLERP-L2-13B",
+        "Undi95/Toppy-M-7B",
+        "WizardLM/WizardCoder-Python-34B-V1.0",
+        "WizardLM/WizardLM-13B-V1.2",
+        "WizardLM/WizardLM-70B-V1.0",
+        "allenai/OLMo-7B",
+        "allenai/OLMo-7B-Instruct",
+        "bert-base-uncased",
+        "codellama/CodeLlama-13b-Instruct-hf",
+        "codellama/CodeLlama-13b-Python-hf",
+        "codellama/CodeLlama-13b-hf",
+        "codellama/CodeLlama-34b-Python-hf",
+        "codellama/CodeLlama-34b-hf",
+        "codellama/CodeLlama-70b-Instruct-hf",
+        "codellama/CodeLlama-70b-Python-hf",
+        "codellama/CodeLlama-70b-hf",
+        "codellama/CodeLlama-7b-Instruct-hf",
+        "codellama/CodeLlama-7b-Python-hf",
+        "codellama/CodeLlama-7b-hf",
+        "cognitivecomputations/dolphin-2.5-mixtral-8x7b",
+        "deepseek-ai/deepseek-coder-33b-instruct",
+        "garage-bAInd/Platypus2-70B-instruct",
+        "google/gemma-2b",
+        "google/gemma-7b",
+        "google/gemma-7b-it",
+        "gradientai/Llama-3-70B-Instruct-Gradient-1048k",
+        "hazyresearch/M2-BERT-2k-Retrieval-Encoder-V1",
+        "huggyllama/llama-13b",
+        "huggyllama/llama-30b",
+        "huggyllama/llama-65b",
+        "huggyllama/llama-7b",
+        "lmsys/vicuna-13b-v1.3",
+        "lmsys/vicuna-13b-v1.5",
+        "lmsys/vicuna-13b-v1.5-16k",
+        "lmsys/vicuna-7b-v1.3",
+        "lmsys/vicuna-7b-v1.5",
+        "meta-llama/Llama-2-13b-hf",
+        "meta-llama/Llama-2-70b-chat-hf",
+        "meta-llama/Llama-2-7b-hf",
+        "meta-llama/Llama-3-70b-hf",
+        "meta-llama/Llama-3-8b-hf",
+        "meta-llama/Meta-Llama-3-70B",
+        "meta-llama/Meta-Llama-3-70B-Instruct",
+        "meta-llama/Meta-Llama-3-8B-Instruct",
+        "meta-llama/Meta-Llama-3.1-70B-Instruct-Reference",
+        "meta-llama/Meta-Llama-3.1-70B-Reference",
+        "meta-llama/Meta-Llama-3.1-8B-Reference",
+        "microsoft/phi-2",
+        "mistralai/Mixtral-8x22B",
+        "openchat/openchat-3.5-1210",
+        "prompthero/openjourney",
+        "runwayml/stable-diffusion-v1-5",
+        "sentence-transformers/msmarco-bert-base-dot-v5",
+        "snorkelai/Snorkel-Mistral-PairRM-DPO",
+        "stabilityai/stable-diffusion-2-1",
+        "teknium/OpenHermes-2-Mistral-7B",
+        "teknium/OpenHermes-2p5-Mistral-7B",
+        "togethercomputer/CodeLlama-13b-Instruct",
+        "togethercomputer/CodeLlama-13b-Python",
+        "togethercomputer/CodeLlama-34b",
+        "togethercomputer/CodeLlama-34b-Python",
+        "togethercomputer/CodeLlama-7b-Instruct",
+        "togethercomputer/CodeLlama-7b-Python",
+        "togethercomputer/Koala-13B",
+        "togethercomputer/Koala-7B",
+        "togethercomputer/LLaMA-2-7B-32K",
+        "togethercomputer/SOLAR-10.7B-Instruct-v1.0-int4",
+        "togethercomputer/StripedHyena-Hessian-7B",
+        "togethercomputer/alpaca-7b",
+        "togethercomputer/evo-1-131k-base",
+        "togethercomputer/evo-1-8k-base",
+        "togethercomputer/guanaco-13b",
+        "togethercomputer/guanaco-33b",
+        "togethercomputer/guanaco-65b",
+        "togethercomputer/guanaco-7b",
+        "togethercomputer/llama-2-13b",
+        "togethercomputer/llama-2-70b-chat",
+        "togethercomputer/llama-2-7b",
+        "wavymulder/Analog-Diffusion",
+        "zero-one-ai/Yi-34B",
+        "zero-one-ai/Yi-34B-Chat",
+        "zero-one-ai/Yi-6B",
+    ]
+
+    _sync_client_ = openai.OpenAI
+    _async_client_ = openai.AsyncOpenAI
+
+    @classmethod
+    def get_model_list(cls):
+        # Togheter.ai has a different response in model list then openai
+        # and the OpenAI class returns an error when calling .models.list()
+        import requests
+        import os
+
+        url = "https://api.together.xyz/v1/models?filter=serverless"
+        token = os.getenv(cls._env_key_name_)
+        headers = {"accept": "application/json", "authorization": f"Bearer {token}"}
+
+        response = requests.get(url, headers=headers)
+        return response.json()
+
+    @classmethod
+    def available(cls) -> List[str]:
+        if not cls._models_list_cache:
+            try:
+                cls._models_list_cache = [
+                    m["id"]
+                    for m in cls.get_model_list()
+                    if m["id"] not in cls.model_exclude_list
+                ]
+            except Exception as e:
+                raise
+        return cls._models_list_cache

edsl/inference_services/models_available_cache.py

@@ -66,4 +66,29 @@ models_available = {
         "openchat/openchat_3.5",
     ],
     "google": ["gemini-pro"],
+    "bedrock": [
+        "amazon.titan-tg1-large",
+        "amazon.titan-text-lite-v1",
+        "amazon.titan-text-express-v1",
+        "anthropic.claude-instant-v1",
+        "anthropic.claude-v2:1",
+        "anthropic.claude-v2",
+        "anthropic.claude-3-sonnet-20240229-v1:0",
+        "anthropic.claude-3-haiku-20240307-v1:0",
+        "anthropic.claude-3-opus-20240229-v1:0",
+        "anthropic.claude-3-5-sonnet-20240620-v1:0",
+        "cohere.command-text-v14",
+        "cohere.command-r-v1:0",
+        "cohere.command-r-plus-v1:0",
+        "cohere.command-light-text-v14",
+        "meta.llama3-8b-instruct-v1:0",
+        "meta.llama3-70b-instruct-v1:0",
+        "meta.llama3-1-8b-instruct-v1:0",
+        "meta.llama3-1-70b-instruct-v1:0",
+        "meta.llama3-1-405b-instruct-v1:0",
+        "mistral.mistral-7b-instruct-v0:2",
+        "mistral.mixtral-8x7b-instruct-v0:1",
+        "mistral.mistral-large-2402-v1:0",
+        "mistral.mistral-large-2407-v1:0",
+    ],
 }

edsl/inference_services/registry.py

@@ -7,7 +7,25 @@ from edsl.inference_services.AnthropicService import AnthropicService
 from edsl.inference_services.DeepInfraService import DeepInfraService
 from edsl.inference_services.GoogleService import GoogleService
 from edsl.inference_services.GroqService import GroqService
+from edsl.inference_services.AwsBedrock import AwsBedrockService
+from edsl.inference_services.AzureAI import AzureAIService
+from edsl.inference_services.OllamaService import OllamaService
+from edsl.inference_services.TestService import TestService
+from edsl.inference_services.MistralAIService import MistralAIService
+from edsl.inference_services.TogetherAIService import TogetherAIService
 
 default = InferenceServicesCollection(
-    [OpenAIService, AnthropicService, DeepInfraService, GoogleService, GroqService]
+    [
+        OpenAIService,
+        AnthropicService,
+        DeepInfraService,
+        GoogleService,
+        GroqService,
+        AwsBedrockService,
+        AzureAIService,
+        OllamaService,
+        TestService,
+        MistralAIService,
+        TogetherAIService,
+    ]
 )

edsl/jobs/Answers.py

@@ -2,24 +2,22 @@
 
 from collections import UserDict
 from rich.table import Table
+from edsl.data_transfer_models import EDSLResultObjectInput
 
 
 class Answers(UserDict):
     """Helper class to hold the answers to a survey."""
 
-    def add_answer(self, response, question) -> None:
-        """Add a response to the answers dictionary.
-
-        >>> from edsl import QuestionFreeText
-        >>> q = QuestionFreeText.example()
-        >>> answers = Answers()
-        >>> answers.add_answer({"answer": "yes"}, q)
-        >>> answers[q.question_name]
-        'yes'
-        """
-        answer = response.get("answer")
-        comment = response.pop("comment", None)
+    def add_answer(
+        self, response: EDSLResultObjectInput, question: "QuestionBase"
+    ) -> None:
+        """Add a response to the answers dictionary."""
+        answer = response.answer
+        comment = response.comment
+        generated_tokens = response.generated_tokens
         # record the answer
+        if generated_tokens:
+            self[question.question_name + "_generated_tokens"] = generated_tokens
         self[question.question_name] = answer
         if comment:
             self[question.question_name + "_comment"] = comment