palimpzest 0.7.21__py3-none-any.whl → 0.8.1__py3-none-any.whl

palimpzest/query/generators/generators.py

@@ -4,36 +4,28 @@ This file contains the Generator classes and generator factory.
 
 from __future__ import annotations
 
+import json
 import logging
 import os
-import re
 import time
 import warnings
-from abc import ABC, abstractmethod
-from collections import Counter
 from copy import deepcopy
 from typing import Any, Generic, TypeVar
 
+import litellm
+import regex as re  # Use regex instead of re to used variable length lookbehind
 from colorama import Fore, Style
-from openai import OpenAI
-from openai.types.chat.chat_completion import ChatCompletion
-from together import Together
-from together.types.chat_completions import ChatCompletionResponse
+from pydantic.fields import FieldInfo
 
 from palimpzest.constants import (
     MODEL_CARDS,
-    APIClient,
     Cardinality,
     Model,
     PromptStrategy,
 )
-from palimpzest.core.data.dataclasses import GenerationStats
 from palimpzest.core.elements.records import DataRecord
-from palimpzest.core.lib.fields import Field, ListField
+from palimpzest.core.models import GenerationStats
 from palimpzest.prompts import PromptFactory
-from palimpzest.query.generators.api_client_factory import APIClientFactory
-from palimpzest.utils.generation_helpers import get_json_from_answer
-from palimpzest.utils.sandbox import API
 
 # DEFINITIONS
 GenerationOutput = tuple[dict, str | None, GenerationStats, list[dict]]
@@ -43,31 +35,71 @@ InputType = TypeVar("InputType")
 
 logger = logging.getLogger(__name__)
 
-def generator_factory(
-    model: Model, prompt_strategy: PromptStrategy, cardinality: Cardinality, verbose: bool = False
-) -> BaseGenerator:
+def get_json_from_answer(answer: str, model: Model, cardinality: Cardinality) -> dict[str, Any]:
     """
-    Factory function to return the correct generator based on the model, strategy, and cardinality.
+    This function parses an LLM response which is supposed to output a JSON object
+    and optimistically searches for the substring containing the JSON object.
     """
-    if model.is_openai_model():
-        return OpenAIGenerator(model, prompt_strategy, cardinality, verbose)
-
-    elif model.is_together_model():
-        return TogetherGenerator(model, prompt_strategy, cardinality, verbose)
-
-    raise Exception(f"Unsupported model: {model}")
-
-
-def get_api_key(key: str) -> str:
-    # get API key from environment or throw an exception if it's not set
-    if key not in os.environ:
-        raise ValueError("key not found in environment variables")
-
-    return os.environ[key]
-
-
+    # model-specific trimming for LLAMA3 responses
+    if model.is_llama_model():
+        answer = answer.split("---")[0]
+        answer = answer.replace("True", "true")
+        answer = answer.replace("False", "false")
+
+    # split off context / excess, which models sometimes output after answer
+    answer = answer.split("Context:")[0]
+    answer = answer.split("# this is the answer")[0]
+
+    # trim the answer to only include the JSON dictionary
+    if cardinality == Cardinality.ONE_TO_ONE:
+        if not answer.strip().startswith("{"):
+            # Find the start index of the actual JSON string assuming the prefix is followed by the JSON dictionary
+            start_index = answer.find("{")
+            if start_index != -1:
+                # Remove the prefix and any leading characters before the JSON starts
+                answer = answer[start_index:]
+
+        if not answer.strip().endswith("}"):
+            # Find the end index of the actual JSON string assuming the suffix is preceded by the JSON dictionary
+            end_index = answer.rfind("}")
+            if end_index != -1:
+                # Remove the suffix and any trailing characters after the JSON ends
+                answer = answer[: end_index + 1]
+
+    # otherwise, trim the answer to only include the JSON array
+    else:
+        if not answer.strip().startswith("["):
+            # Find the start index of the actual JSON string assuming the prefix is followed by the JSON array
+            start_index = answer.find("[")
+            if start_index != -1:
+                # Remove the prefix and any leading characters before the JSON starts
+                answer = answer[start_index:]
+
+        if not answer.strip().endswith("]"):
+            # Find the end index of the actual JSON string
+            # assuming the suffix is preceded by the JSON object/array
+            end_index = answer.rfind("]")
+            if end_index != -1:
+                # Remove the suffix and any trailing characters after the JSON ends
+                answer = answer[: end_index + 1]
+
+    # Handle weird escaped values. I am not sure why the model
+    # is returning these, but the JSON parser can't take them
+    answer = answer.replace(r"\_", "_")
+    answer = answer.replace("\\n", "\n")
+    # Remove https and http prefixes to not conflict with comment detection
+    # Handle comments in the JSON response. Use regex from // until end of line
+    answer = re.sub(r"(?<!https?:)\/\/.*?$", "", answer, flags=re.MULTILINE)
+    answer = re.sub(r",\n.*\.\.\.$", "", answer, flags=re.MULTILINE)
+    # Sanitize newlines in the JSON response
+    answer = answer.replace("\n", " ")
+
+    # finally, parse and return the JSON object; errors are handled by the caller
+    return json.loads(answer)
+
+# TODO: push parallelism of generations into LiteLLM rather than threadpool in executor
 # TODO: make sure answer parsing works with custom prompts / parsers (can defer this)
-class BaseGenerator(Generic[ContextType, InputType], ABC):
+class Generator(Generic[ContextType, InputType]):
     """
     Abstract base class for Generators.
     """
@@ -76,94 +108,21 @@ class BaseGenerator(Generic[ContextType, InputType], ABC):
         self,
         model: Model,
         prompt_strategy: PromptStrategy,
+        reasoning_effort: str | None = None,
+        api_base: str | None = None,
         cardinality: Cardinality = Cardinality.ONE_TO_ONE,
+        desc: str | None = None,
         verbose: bool = False,
-        system_role: str = "system",
     ):
         self.model = model
         self.model_name = model.value
         self.cardinality = cardinality
         self.prompt_strategy = prompt_strategy
+        self.reasoning_effort = reasoning_effort
+        self.api_base = api_base
+        self.desc = desc
         self.verbose = verbose
-        self.system_role = system_role
-        self.prompt_factory = PromptFactory(prompt_strategy, model, cardinality)
-
-    @abstractmethod
-    def _get_client_or_model(self, **kwargs) -> Any:
-        """Returns a client (or local model) which can be invoked to perform the generation."""
-        pass
-
-    @abstractmethod
-    def _generate_completion(self, client_or_model: Any, payload: dict, **kwargs) -> Any:
-        """Generates a completion object using the client (or local model)."""
-        pass
-
-    @abstractmethod
-    def _get_completion_text(self, completion: Any, **kwargs) -> Any:
-        """Extract the completion text from the completion object."""
-        pass
-
-    @abstractmethod
-    def _get_usage(self, completion: Any, **kwargs) -> Any:
-        """Extract the usage statistics from the completion object."""
-        pass
-
-    @abstractmethod
-    def _get_finish_reason(self, completion: Any, **kwargs) -> Any:
-        """Extract the finish reason from the completion object."""
-        pass
-
-    @abstractmethod
-    def _get_answer_log_probs(self, completion: Any, **kwargs) -> Any:
-        """Extract the log probabilities from the completion object."""
-        pass
-
-    def _generate_payload(self, messages: list[dict], **kwargs) -> dict:
-        """
-        Generates the payload which will be fed into the client (or local model).
-
-        Each message will be a dictionary with the following format:
-        {
-            "role": "user" | "system",
-            "type": "text" | "image",
-            "content": str
-        }
-        """
-        # get basic parameters
-        model = self.model_name
-        temperature = kwargs.get("temperature", 0.0)
-
-        # construct messages and add system prompt if present
-        chat_messages, user_content = [], []
-        for message in messages:
-            # flush user content into a message and add system message
-            if message["role"] == "system":
-                if len(user_content) > 0:
-                    chat_messages.append({"role": "user", "content": user_content})
-                    user_content = []
-
-                chat_messages.append({"role": self.system_role, "content": message["content"]})
-
-            # add user content for text messages
-            elif message["role"] == "user" and message["type"] == "text":
-                user_content.append({"type": "text", "text": message["content"]})
-
-            # add user content for image messages
-            elif message["role"] == "user" and message["type"] == "image":
-                user_content.append({"type": "image_url", "image_url": {"url": message["content"]}})
-
-        # flush any remaining user content into a final message
-        if len(user_content) > 0:
-            chat_messages.append({"role": "user", "content": user_content})
-
-        # construct and return payload
-        payload = {
-            "model": model,
-            "temperature": temperature,
-            "messages": chat_messages,
-        }
-
-        return payload
+        self.prompt_factory = PromptFactory(prompt_strategy, model, cardinality, desc)
 
     def _parse_reasoning(self, completion_text: str, **kwargs) -> str:
         """Extract the reasoning for the generated output from the completion object."""
@@ -183,7 +142,7 @@ class BaseGenerator(Generic[ContextType, InputType], ABC):
         # otherwise, return the full completion text
         return completion_text
 
-    def _prepare_field_answers(self, field_answers: dict | list[dict], fields: dict[str, Field]) -> dict[str, list]:
+    def _prepare_field_answers(self, field_answers: dict | list[dict], fields: dict[str, FieldInfo]) -> dict[str, list]:
         """
         field_answers is a dictionary mapping fields to their values. For one-to-one converts, wrap each
         answer in a list. For one-to-many converts, invert the list of dictionaries into a dictionary with
@@ -205,7 +164,7 @@ class BaseGenerator(Generic[ContextType, InputType], ABC):
 
         return field_answers
 
-    def _check_convert_answer_text(self, answer_text: str, fields: dict[str, Field], throw_exception: bool=False) -> dict | list[dict] | None:
+    def _check_convert_answer_text(self, answer_text: str, fields: dict[str, FieldInfo], throw_exception: bool=False) -> dict | list[dict] | None:
         """
         Try parsing the answer text into a JSON object. If the parsing fails, return None.
         """
@@ -213,18 +172,6 @@ class BaseGenerator(Generic[ContextType, InputType], ABC):
             # extract json from the answer text
             field_answers = get_json_from_answer(answer_text, self.model, self.cardinality)
 
-            # TODO: wrap non-list outputs in a list if expected output is a list
-
-            # common error for one-to-one: if the output is a singleton list which contains a list, but the expected field type
-            # is a list of strings, or a list of floats, i.e. not a list of lists; then extract the inner list
-            if self.cardinality == Cardinality.ONE_TO_ONE:
-                for field, field_type in fields.items():
-                    answer = field_answers[field]
-                    field_type_is_not_list_of_lists = isinstance(field_type, ListField) and not issubclass(field_type.element_type, ListField)
-                    answer_is_list_of_lists = isinstance(answer, list) and len(answer) == 1 and isinstance(answer[0], list)
-                    if field_type_is_not_list_of_lists and answer_is_list_of_lists:
-                        field_answers[field] = answer[0]
-
             # prepare the field answers to match the expected output and return
             return self._prepare_field_answers(field_answers, fields)
 
@@ -234,7 +181,7 @@ class BaseGenerator(Generic[ContextType, InputType], ABC):
 
         return None
 
-    def _check_filter_answer_text(self, answer_text: str) -> dict | None:
+    def _check_bool_answer_text(self, answer_text: str) -> dict | None:
         """
         Return {"passed_operator": True} if and only if "true" is in the answer text.
         Return {"passed_operator": False} if and only if "false" is in the answer text.
@@ -249,7 +196,7 @@ class BaseGenerator(Generic[ContextType, InputType], ABC):
 
         return None
 
-    def _parse_convert_answer(self, completion_text: str, fields: dict[str, Field], json_output: bool) -> dict[str, list]:
+    def _parse_convert_answer(self, completion_text: str, fields: dict[str, FieldInfo], json_output: bool) -> dict[str, list]:
         """Extract the answer from the completion object for convert operations."""
         # if the model followed the default instructions, the completion text will place
         # its answer between "ANSWER:" and "---"
@@ -288,15 +235,15 @@ class BaseGenerator(Generic[ContextType, InputType], ABC):
 
         return self._check_convert_answer_text(completion_text, fields, throw_exception=True)
 
-    def _parse_filter_answer(self, completion_text: str) -> dict[str, list]:
-        """Extract the answer from the completion object for filter operations."""
+    def _parse_bool_answer(self, completion_text: str) -> dict[str, list]:
+        """Extract the answer from the completion object for filter and join operations."""
         # if the model followed the default instructions, the completion text will place
         # its answer between "ANSWER:" and "---"
         regex = re.compile("answer:(.*?)---", re.IGNORECASE | re.DOTALL)
         matches = regex.findall(completion_text)
         if len(matches) > 0:
             answer_text = matches[0].strip()
-            field_answers = self._check_filter_answer_text(answer_text)
+            field_answers = self._check_bool_answer_text(answer_text)
             if field_answers is not None:
                 return field_answers
 
@@ -305,18 +252,18 @@ class BaseGenerator(Generic[ContextType, InputType], ABC):
         matches = regex.findall(completion_text)
         if len(matches) > 0:
             answer_text = matches[0].strip()
-            field_answers = self._check_filter_answer_text(answer_text)
+            field_answers = self._check_bool_answer_text(answer_text)
             if field_answers is not None:
                 return field_answers
 
         # finally, try taking all of the text; throw an exception if this doesn't work
-        field_answers = self._check_filter_answer_text(completion_text)
+        field_answers = self._check_bool_answer_text(completion_text)
         if field_answers is None:
             raise Exception(f"Could not parse answer from completion text: {completion_text}")
 
         return field_answers
 
-    def _parse_answer(self, completion_text: str, fields: dict[str, Field] | None, json_output: bool, **kwargs) -> dict[str, list]:
+    def _parse_answer(self, completion_text: str, fields: dict[str, FieldInfo] | None, json_output: bool, **kwargs) -> dict[str, list]:
         """Extract the answer from the completion object."""
         # use a custom answer parser if provided
         if kwargs.get("parse_answer"):
@@ -328,16 +275,15 @@ class BaseGenerator(Generic[ContextType, InputType], ABC):
 
         # extract the per-field answers from the completion text
         field_answers = (
-            self._parse_filter_answer(completion_text)
-            if self.prompt_strategy.is_bool_prompt()
+            self._parse_bool_answer(completion_text)
+            if self.prompt_strategy.is_bool_prompt() or self.prompt_strategy.is_join_prompt()
             else self._parse_convert_answer(completion_text, fields, json_output)
         )
 
         return field_answers
 
-    def __call__(self, candidate: DataRecord, fields: dict[str, Field] | None, json_output: bool=True, **kwargs) -> GenerationOutput:
+    def __call__(self, candidate: DataRecord, fields: dict[str, FieldInfo] | None, right_candidate: DataRecord | None = None, json_output: bool=True, **kwargs) -> GenerationOutput:
         """Take the input record (`candidate`), generate the output `fields`, and return the generated output."""
-        client = self._get_client_or_model()
         logger.debug(f"Generating for candidate {candidate} with fields {fields}")
 
         # fields can only be None if the user provides an answer parser
@@ -352,23 +298,45 @@ class BaseGenerator(Generic[ContextType, InputType], ABC):
             warnings.warn("Provided `system_prompt` without providing `prompt`; setting `prompt` = `system_prompt`.")  # noqa: B028
 
         # generate a list of messages which can be used to construct a payload
-        messages = self.prompt_factory.create_messages(candidate, fields, **kwargs)
-
-        # create the chat payload
-        chat_payload = self._generate_payload(messages, **kwargs)
+        messages = self.prompt_factory.create_messages(candidate, fields, right_candidate, **kwargs)
 
         # generate the text completion
         start_time = time.time()
         completion = None
         try:
-            completion = self._generate_completion(client, chat_payload, **kwargs)
+            completion_kwargs = {}
+            if not self.model.is_o_model() and not self.model.is_gpt_5_model():
+                completion_kwargs = {"temperature": kwargs.get("temperature", 0.0), **completion_kwargs}
+            if self.prompt_strategy.is_audio_prompt():
+                completion_kwargs = {"modalities": ["text"], **completion_kwargs}
+            if self.model.is_reasoning_model():
+                if self.model.is_vertex_model():
+                    reasoning_effort = self.reasoning_effort
+                    if self.reasoning_effort is None and self.model == Model.GEMINI_2_5_PRO:
+                        reasoning_effort = "low"
+                    elif self.reasoning_effort is None:
+                        reasoning_effort = "disable"
+                    completion_kwargs = {"reasoning_effort": reasoning_effort, **completion_kwargs}
+                elif self.model.is_anthropic_model() and self.reasoning_effort is not None:
+                    completion_kwargs = {"reasoning_effort": self.reasoning_effort, **completion_kwargs}
+                elif self.model.is_openai_model():
+                    reasoning_effort = "minimal" if self.reasoning_effort is None else self.reasoning_effort
+                    completion_kwargs = {"reasoning_effort": reasoning_effort, **completion_kwargs}
+            if self.model.is_vllm_model():
+                completion_kwargs = {"api_base": self.api_base, **completion_kwargs}
+            completion = litellm.completion(model=self.model_name, messages=messages, **completion_kwargs)
             end_time = time.time()
             logger.debug(f"Generated completion in {end_time - start_time:.2f} seconds")
         # if there's an error generating the completion, we have to return an empty answer
         # and can only account for the time spent performing the failed generation
-        except Exception:
-            # logger.error(f"Error generating completion: {e}")
-            field_answers = {field_name: None for field_name in fields}
+        except Exception as e:
+            print(f"Error generating completion: {e}")
+            logger.error(f"Error generating completion: {e}")
+            field_answers = (
+                {"passed_operator": False}
+                if self.prompt_strategy.is_bool_prompt() or self.prompt_strategy.is_join_prompt()
+                else {field_name: None for field_name in fields}
+            )
             reasoning = None
             generation_stats = GenerationStats(
                 model_name=self.model_name,
@@ -381,40 +349,57 @@ class BaseGenerator(Generic[ContextType, InputType], ABC):
         # parse usage statistics and create the GenerationStats
         generation_stats = None
         if completion is not None:
-            usage = self._get_usage(completion, **kwargs)
-            # finish_reason = self._get_finish_reason(completion, **kwargs)
-            # answer_log_probs = self._get_answer_log_probs(completion, **kwargs)
+            usage = completion.usage.model_dump()
 
-            # get cost per input/output token for the model and parse number of input and output tokens
-            usd_per_input_token = MODEL_CARDS[self.model_name]["usd_per_input_token"]
+            # get cost per input/output token for the model
+            usd_per_input_token = MODEL_CARDS[self.model_name].get("usd_per_input_token", 0.0)
+            usd_per_audio_input_token = MODEL_CARDS[self.model_name].get("usd_per_audio_input_token", 0.0)
             usd_per_output_token = MODEL_CARDS[self.model_name]["usd_per_output_token"]
-            input_tokens = usage["input_tokens"]
-            output_tokens = usage["output_tokens"]
+
+            # TODO: for some models (e.g. GPT-5) we cannot separate text from image prompt tokens yet;
+            #       for now, we only use tokens from prompt_token_details if it's an audio prompt
+            # get output tokens (all text) and input tokens by modality
+            output_tokens = usage["completion_tokens"]
+            if self.prompt_strategy.is_audio_prompt():
+                input_audio_tokens = usage["prompt_tokens_details"].get("audio_tokens", 0)
+                input_text_tokens = usage["prompt_tokens_details"].get("text_tokens", 0)
+                input_image_tokens = 0
+            else:
+                input_audio_tokens = 0
+                input_text_tokens = usage["prompt_tokens"]
+                input_image_tokens = 0
+            input_tokens = input_audio_tokens + input_text_tokens + input_image_tokens
+
+            # compute the input and output token costs
+            total_input_cost = (input_text_tokens + input_image_tokens) * usd_per_input_token + input_audio_tokens * usd_per_audio_input_token
+            total_output_cost = output_tokens * usd_per_output_token
 
             generation_stats = GenerationStats(
                 model_name=self.model_name,
                 llm_call_duration_secs=end_time - start_time,
                 fn_call_duration_secs=0.0,
+                input_audio_tokens=input_audio_tokens,
+                input_text_tokens=input_text_tokens,
+                input_image_tokens=input_image_tokens,
                 total_input_tokens=input_tokens,
                 total_output_tokens=output_tokens,
-                total_input_cost=input_tokens * usd_per_input_token,
-                total_output_cost=output_tokens * usd_per_output_token,
-                cost_per_record=input_tokens * usd_per_input_token + output_tokens * usd_per_output_token,
+                total_input_cost=total_input_cost,
+                total_output_cost=total_output_cost,
+                cost_per_record=total_input_cost + total_output_cost,
                 total_llm_calls=1,
-                # "system_prompt": system_prompt,
-                # "prompt": prompt,
-                # "usage": usage,
-                # "finish_reason": finish_reason,
-                # "answer_log_probs": answer_log_probs,
-                # "answer": answer,
             )
 
         # pretty print prompt + full completion output for debugging
-        completion_text = self._get_completion_text(completion, **kwargs)
+        completion_text = completion.choices[0].message.content
         prompt = ""
         for message in messages:
             if message["role"] == "user":
-                prompt += message["content"] + "\n" if message["type"] == "text" else "<image>\n"
+                if message["type"] == "text":
+                    prompt += message["content"] + "\n"
+                elif message["type"] == "image":
+                    prompt += "<image>\n"
+                elif message["type"] == "input_audio":
+                    prompt += "<audio>\n"
         logger.debug(f"PROMPT:\n{prompt}")
         logger.debug(Fore.GREEN + f"{completion_text}\n" + Style.RESET_ALL)
 
@@ -422,17 +407,20 @@ class BaseGenerator(Generic[ContextType, InputType], ABC):
         reasoning = None
         try:
             reasoning = self._parse_reasoning(completion_text, **kwargs)
-        except Exception:
-            # logger.error(f"Error parsing reasoning and answers: {e}")
-            logger.debug("TODO: undo this")
+        except Exception as e:
+            logger.error(f"Error parsing reasoning and answers: {e}")
             pass
 
         # parse field answers
-        field_answers = None if fields is None else {field_name: None for field_name in fields}
+        field_answers = None 
+        if fields is not None and (self.prompt_strategy.is_bool_prompt() or self.prompt_strategy.is_join_prompt()):
+            field_answers = {"passed_operator": False}
+        elif fields is not None and not (self.prompt_strategy.is_bool_prompt() or self.prompt_strategy.is_join_prompt()):
+            field_answers = {field_name: None for field_name in fields}
         try:
             field_answers = self._parse_answer(completion_text, fields, json_output, **kwargs)
         except Exception as e:
-            # logger.error(f"Error parsing answers: {e}")
+            logger.error(f"Error parsing answers: {e}")
             os.makedirs("parse-answer-errors", exist_ok=True)
             ts = time.time()
             with open(f"parse-answer-errors/error-{ts}.txt", "w") as f:
@@ -448,162 +436,3 @@ class BaseGenerator(Generic[ContextType, InputType], ABC):
 
         logger.debug(f"Generated field answers: {field_answers}")
         return field_answers, reasoning, generation_stats, messages
-
-
-class OpenAIGenerator(BaseGenerator[str | list[str], str]):
-    """
-    Class for generating text using the OpenAI chat API.
-    """
-
-    def __init__(
-        self,
-        model: Model,
-        prompt_strategy: PromptStrategy,
-        cardinality: Cardinality = Cardinality.ONE_TO_ONE,
-        verbose: bool = False,
-    ):
-        # assert that model is an OpenAI model
-        assert model.is_openai_model()
-        super().__init__(model, prompt_strategy, cardinality, verbose, "developer")
-
-    def _get_client_or_model(self, **kwargs) -> OpenAI:
-        """Returns a client (or local model) which can be invoked to perform the generation."""
-        return APIClientFactory.get_client(APIClient.OPENAI, get_api_key("OPENAI_API_KEY"))
-
-    def _generate_completion(self, client: OpenAI, payload: dict, **kwargs) -> ChatCompletion:
-        """Generates a completion object using the client (or local model)."""
-        return client.chat.completions.create(**payload)
-
-    def _get_completion_text(self, completion: ChatCompletion, **kwargs) -> str:
-        """Extract the completion text from the completion object."""
-        return completion.choices[0].message.content
-
-    def _get_usage(self, completion: ChatCompletion, **kwargs) -> dict:
-        """Extract the usage statistics from the completion object."""
-        return {
-            "input_tokens": completion.usage.prompt_tokens,
-            "output_tokens": completion.usage.completion_tokens,
-        }
-
-    def _get_finish_reason(self, completion: ChatCompletion, **kwargs) -> str:
-        """Extract the finish reason from the completion object."""
-        return completion.choices[0].finish_reason
-
-    def _get_answer_log_probs(self, completion: ChatCompletion, **kwargs) -> list[float]:
-        """Extract the log probabilities from the completion object."""
-        return completion.choices[0].logprobs
-
-
-class TogetherGenerator(BaseGenerator[str | list[str], str]):
-    """
-    Class for generating text using the Together chat API.
-    """
-
-    def __init__(
-        self,
-        model: Model,
-        prompt_strategy: PromptStrategy,
-        cardinality: Cardinality = Cardinality.ONE_TO_ONE,
-        verbose: bool = False,
-    ):
-        # assert that model is a model offered by Together
-        assert model.is_together_model()
-        super().__init__(model, prompt_strategy, cardinality, verbose, "system")
-
-    def _generate_payload(self, messages: list[dict], **kwargs) -> dict:
-        """
-        Generates the payload which will be fed into the client (or local model).
-
-        Each message will be a dictionary with the following format:
-        {
-            "role": "user" | "system",
-            "type": "text" | "image",
-            "content": str
-        }
-
-        For LLAMA3, the payload needs to be in a {"role": <role>, "content": <content>} format.
-        """
-        # for other models, use our standard payload generation
-        if not self.model.is_llama_model():
-            return super()._generate_payload(messages, **kwargs)
-
-        # get basic parameters
-        model = self.model_name
-        temperature = kwargs.get("temperature", 0.0)
-
-        # construct messages in simple {"role": <role>, "content": <content>} format
-        chat_messages = []
-        for message in messages:
-            chat_messages.append({"role": message["role"], "content": message["content"]})
-
-        # construct and return payload
-        payload = {
-            "model": model,
-            "temperature": temperature,
-            "messages": chat_messages,
-        }
-
-        return payload
-
-    def _get_client_or_model(self, **kwargs) -> Together:
-        """Returns a client (or local model) which can be invoked to perform the generation."""
-        return APIClientFactory.get_client(APIClient.TOGETHER, get_api_key("TOGETHER_API_KEY"))
-
-    def _generate_completion(self, client: Together, payload: dict, **kwargs) -> ChatCompletionResponse:
-        """Generates a completion object using the client (or local model)."""
-        return client.chat.completions.create(**payload)
-
-    def _get_completion_text(self, completion: ChatCompletionResponse, **kwargs) -> str:
-        """Extract the completion text from the completion object."""
-        return completion.choices[0].message.content
-
-    def _get_usage(self, completion: ChatCompletionResponse, **kwargs) -> dict:
-        """Extract the usage statistics from the completion object."""
-        return {
-            "input_tokens": completion.usage.prompt_tokens,
-            "output_tokens": completion.usage.completion_tokens,
-        }
-
-    def _get_finish_reason(self, completion: ChatCompletionResponse, **kwargs) -> str:
-        """Extract the finish reason from the completion object."""
-        return completion.choices[0].finish_reason.value
-
-    def _get_answer_log_probs(self, completion: ChatCompletionResponse, **kwargs) -> list[float]:
-        """Extract the log probabilities from the completion object."""
-        return completion.choices[0].logprobs
-
-
-### CODE SYNTHESIS EXECUTION ###
-def code_execution(api: API, code: str, candidate_dict: dict[str, Any], verbose: bool = False):
-    inputs = {field_name: candidate_dict[field_name] for field_name in api.inputs}
-    response = api.api_execute(code, inputs)
-    pred = response["response"] if response["status"] and response["response"] else None
-    return pred
-
-
-def code_ensemble_execution(
-    api: API, code_ensemble: dict[str, str], candidate_dict: dict[str, Any], verbose: bool = True
-) -> GenerationOutput:
-    start_time = time.time()
-    try:
-        preds = list()
-        for _, code in code_ensemble.items():
-            pred = code_execution(api, code, candidate_dict)
-            preds.append(pred)
-
-        preds = [pred for pred in preds if pred is not None]
-
-        if len(preds) == 1:
-            majority_response = preds[0]
-            exec_stats = GenerationStats(fn_call_duration_secs=time.time() - start_time)
-            return majority_response, None, exec_stats
-
-        if len(preds) > 0:
-            majority_response = Counter(preds).most_common(1)[0][0]
-            exec_stats = GenerationStats(fn_call_duration_secs=time.time() - start_time)
-            return majority_response, None, exec_stats
-
-    except Exception:
-        pass
-
-    return None, None, GenerationStats(fn_call_duration_secs=time.time() - start_time)