looking good

Author: TLSDC

src/agentlab/llm/chat_api.py

@@ -261,7 +261,7 @@ def __init__(
             **client_args,
         )
 
-    def __call__(self, messages: list[dict]) -> dict:
+    def __call__(self, messages: list[dict], n_samples: int = 1) -> dict:
         # Initialize retry tracking attributes
         self.retries = 0
         self.success = False
@@ -275,6 +275,7 @@ def __call__(self, messages: list[dict]) -> dict:
                 completion = self.client.chat.completions.create(
                     model=self.model_name,
                     messages=messages,
+                    n=n_samples,
                     temperature=self.temperature,
                     max_tokens=self.max_tokens,
                 )
@@ -305,7 +306,10 @@ def __call__(self, messages: list[dict]) -> dict:
         ):
             tracking.TRACKER.instance(input_tokens, output_tokens, cost)
 
-        return AIMessage(completion.choices[0].message.content)
+        if n_samples == 1:
+            return AIMessage(completion.choices[0].message.content)
+        else:
+            return [AIMessage(c.message.content) for c in completion.choices]
 
     def get_stats(self):
         return {

src/agentlab/llm/llm_utils.py

@@ -90,6 +90,60 @@ def retry(
     raise ParseError(f"Could not parse a valid value after {n_retry} retries.")
 
 
+def retry_multiple(
+    chat: "ChatModel",
+    messages: "Discussion",
+    n_retry: int,
+    parser: callable,
+    log: bool = True,
+    num_samples: int = 1,
+):
+    """Retry querying the chat models with the response from the parser until it
+    returns a valid value.
+
+    If the answer is not valid, it will retry and append to the chat the  retry
+    message.  It will stop after `n_retry`.
+
+    Note, each retry has to resend the whole prompt to the API. This can be slow
+    and expensive.
+
+    Args:
+        chat (ChatModel): a ChatModel object taking a list of messages and
+            returning a list of answers, all in OpenAI format.
+        messages (list): the list of messages so far. This list will be modified with
+            the new messages and the retry messages.
+        n_retry (int): the maximum number of sequential retries.
+        parser (callable): a function taking a message and retruning a parsed value,
+            or raising a ParseError
+        log (bool): whether to log the retry messages.
+
+    Returns:
+        dict: the parsed value, with a string at key "action".
+
+    Raises:
+        ParseError: if the parser could not parse the response after n_retry retries.
+    """
+    tries = 0
+    while tries < n_retry:
+        answer_list = chat(messages, num_samples=num_samples)
+        # TODO: could we change this to not use inplace modifications ?
+        messages.append(answer)
+        parsed_answers = []
+        errors = []
+        for answer in answer_list:
+            try:
+                parsed_answers.append(parser(answer["content"]))
+            except ParseError as parsing_error:
+                errors.append(str(parsing_error))
+            tries += 1
+            if log:
+                msg = f"Query failed. Retrying {tries}/{n_retry}.\n[LLM]:\n{answer['content']}\n[User]:\n{str(errors)}"
+                logging.info(msg)
+            messages.append(dict(role="user", content=str(errors)))
+
+    raise ParseError(f"Could not parse a valid value after {n_retry} retries.")
+
+
 def truncate_tokens(text, max_tokens=8000, start=0, model_name="gpt-4"):
     """Use tiktoken to truncate a text to a maximum number of tokens."""
     enc = tiktoken.encoding_for_model(model_name)