| """ NOTE: THIS IS A BETA VERSION OF FUNSEARCH. NEW VERSION DOCUMENTATION WILL BE RELEASED SOON.""" | |
| from typing import Dict, Any, List | |
| from aiflows.interfaces import KeyInterface | |
| from aiflows.utils import logging | |
| from flow_modules.aiflows.ChatFlowModule import ChatAtomicFlow | |
| log = logging.get_logger(f"aiflows.{__name__}") | |
| class SamplerFlow(ChatAtomicFlow): | |
| """ This class implements a SamplerFlow. It is a flow that queries a LLM to generate a response to a given input. This class is a child of ChatAtomicFlow. | |
| and expects the same parameters as ChatAtomicFlow (see https://huggingface.co/aiflows/ChatFlowModule). | |
| **Configuration Parameters**: | |
| - `name` (str): The name of the flow. Default: "SamplerFlowModule" | |
| - `description` (str): A description of the flow. Default: "A flow that queries an LLM model to generate prompts for the Sampler flow" | |
| - `backend` Dict[str,Any]: The backend of the flow. Used to call models via an API. | |
| See litellm's supported models and APIs here: https://docs.litellm.ai/docs/providers. | |
| The default parameters of the backend are all defined at aiflows.backends.llm_lite.LiteLLMBackend | |
| (also see the defaul parameters of litellm's completion parameters: https://docs.litellm.ai/docs/completion/input#input-params-1). | |
| Except for the following parameters who are overwritten by the ChatAtomicFlow in ChatAtomicFlow.yaml: | |
| - `model_name` (Union[Dict[str,str],str]): The name of the model to use. Default: "gpt-4" | |
| When using multiple API providers, the model_name can be a dictionary of the form | |
| {"provider_name": "model_name"}. E.g. {"openai": "gpt-3.5-turbo", "azure": "azure/gpt-3.5-turbo"} | |
| Default: "gpt-3.5-turbo" (the name needs to follow the name of the model in litellm https://docs.litellm.ai/docs/providers). | |
| - `n` (int) : The number of answers to generate. Default: 1 | |
| - `max_tokens` (int): The maximum number of tokens to generate. Default: 2000 | |
| - `temperature` float: The temperature of the generation. Default: 1.0 | |
| - `top_p` float: An alternative to sampling with temperature. It instructs the model to consider the results of | |
| the tokens with top_p probability. Default: 0.4 | |
| - `frequency_penalty` (number): It is used to penalize new tokens based on their frequency in the text so far. Default: 0.0 | |
| - `presence_penalty` (number): It is used to penalize new tokens based on their existence in the text so far. Default: 0.0 | |
| - `stream` (bool): Whether to stream the response or not. Default: false | |
| - `system_message_prompt_template` (Dict[str,Any]): The template of the system message. It is used to generate the system message. Default: See SamplerFlow.yaml for default. | |
| - `init_human_message_prompt_template` (Dict[str,Any]): The prompt template of the human/user message used to initialize the conversation | |
| (first time in). It is used to generate the human message. It's passed as the user message to the LLM. Default: See SamplerFlow.yaml for default. | |
| - `human_message_prompt_template` (Dict[str,Any]): The prompt template of the human/user message (message used everytime the except the first time in). | |
| It's passed as the user message to the LLM. Default: See SamplerFlow.yaml for default. | |
| - `previous_messages` (Dict[str,Any]): Defines which previous messages to include in the input of the LLM. Note that if `first_k`and `last_k` are both none, | |
| all the messages of the flows's history are added to the input of the LLM. Default: | |
| - `first_k` (int): If defined, adds the first_k earliest messages of the flow's chat history to the input of the LLM. Default: 1 | |
| - `last_k` (int): If defined, adds the last_k latest messages of the flow's chat history to the input of the LLM. Default: 1 | |
| *Input Interface Initialized (Expected input the first time in flow)*: | |
| - `header` (str): A header message to include in prompt | |
| - `code` (str): The "example" samples to generate our new sample from. | |
| *Input Interface (Expected input the after the first time in flow)*: | |
| - `header` (str): A header message to include in prompt | |
| - `code` (str): The "example" samples to generate our new sample from. | |
| *Output Interface*: | |
| - `api_output` (str): The output of the API call. It is the response of the LLM to the input. | |
| - `from` (str): The name of the flow that generated the output. It's always "SamplerFlow" | |
| **Citation**: | |
| @Article{FunSearch2023, | |
| author = {Romera-Paredes, Bernardino and Barekatain, Mohammadamin and Novikov, Alexander and Balog, Matej and Kumar, M. Pawan and Dupont, Emilien and Ruiz, Francisco J. R. and Ellenberg, Jordan and Wang, Pengming and Fawzi, Omar and Kohli, Pushmeet and Fawzi, Alhussein}, | |
| journal = {Nature}, | |
| title = {Mathematical discoveries from program search with large language models}, | |
| year = {2023}, | |
| doi = {10.1038/s41586-023-06924-6} | |
| } | |
| """ | |
| def run(self,input_message): | |
| """ This method calls the backend of the flow (so queries the LLM). It calls the backend with the previous messages of the flow. | |
| :return: The output of the backend. | |
| :rtype: Any | |
| """ | |
| input_data = input_message.data | |
| response = self.query_llm(input_data=input_data) | |
| reply_message = self.package_output_message( | |
| input_message, | |
| response = {"api_output": response, "from": "SamplerFlow"} | |
| ) | |
| self.send_message(reply_message) |