first version FunSearch

.gitignore

@@ -0,0 +1,444 @@
+# Created by https://www.toptal.com/developers/gitignore/api/python,java,c++,pycharm,visualstudiocode,macos,linux,windows
+# Edit at https://www.toptal.com/developers/gitignore?templates=python,java,c++,pycharm,visualstudiocode,macos,linux,windows
+
+### C++ ###
+# Prerequisites
+*.d
+
+# Compiled Object files
+*.slo
+*.lo
+*.o
+*.obj
+
+# Precompiled Headers
+*.gch
+*.pch
+
+# Compiled Dynamic libraries
+*.so
+*.dylib
+*.dll
+
+# Fortran module files
+*.mod
+*.smod
+
+# Compiled Static libraries
+*.lai
+*.la
+*.a
+*.lib
+
+# Executables
+*.exe
+*.out
+*.app
+
+### Java ###
+# Compiled class file
+*.class
+
+# Log file
+*.log
+
+# BlueJ files
+*.ctxt
+
+# Mobile Tools for Java (J2ME)
+.mtj.tmp/
+
+# Package Files #
+*.jar
+*.war
+*.nar
+*.ear
+*.zip
+*.tar.gz
+*.rar
+
+# virtual machine crash logs, see http://www.java.com/en/download/help/error_hotspot.xml
+hs_err_pid*
+replay_pid*
+
+### Linux ###
+*~
+
+# temporary files which can be created if a process still has a handle open of a deleted file
+.fuse_hidden*
+
+# KDE directory preferences
+.directory
+
+# Linux trash folder which might appear on any partition or disk
+.Trash-*
+
+# .nfs files are created when an open file is removed but is still being accessed
+.nfs*
+
+### macOS ###
+# General
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# Icon must end with two \r
+Icon
+
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+### macOS Patch ###
+# iCloud generated files
+*.icloud
+
+### PyCharm ###
+# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider
+# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
+
+# User-specific stuff
+.idea/**/workspace.xml
+.idea/**/tasks.xml
+.idea/**/usage.statistics.xml
+.idea/**/dictionaries
+.idea/**/shelf
+
+# AWS User-specific
+.idea/**/aws.xml
+
+# Generated files
+.idea/**/contentModel.xml
+
+# Sensitive or high-churn files
+.idea/**/dataSources/
+.idea/**/dataSources.ids
+.idea/**/dataSources.local.xml
+.idea/**/sqlDataSources.xml
+.idea/**/dynamic.xml
+.idea/**/uiDesigner.xml
+.idea/**/dbnavigator.xml
+
+# Gradle
+.idea/**/gradle.xml
+.idea/**/libraries
+
+# Gradle and Maven with auto-import
+# When using Gradle or Maven with auto-import, you should exclude module files,
+# since they will be recreated, and may cause churn.  Uncomment if using
+# auto-import.
+# .idea/artifacts
+# .idea/compiler.xml
+# .idea/jarRepositories.xml
+# .idea/modules.xml
+# .idea/*.iml
+# .idea/modules
+# *.iml
+# *.ipr
+
+# CMake
+cmake-build-*/
+
+# Mongo Explorer plugin
+.idea/**/mongoSettings.xml
+
+# File-based project format
+*.iws
+
+# IntelliJ
+out/
+
+# mpeltonen/sbt-idea plugin
+.idea_modules/
+
+# JIRA plugin
+atlassian-ide-plugin.xml
+
+# Cursive Clojure plugin
+.idea/replstate.xml
+
+# SonarLint plugin
+.idea/sonarlint/
+
+# Crashlytics plugin (for Android Studio and IntelliJ)
+com_crashlytics_export_strings.xml
+crashlytics.properties
+crashlytics-build.properties
+fabric.properties
+
+# Editor-based Rest Client
+.idea/httpRequests
+
+# Android studio 3.1+ serialized cache file
+.idea/caches/build_file_checksums.ser
+
+### PyCharm Patch ###
+# Comment Reason: https://github.com/joeblau/gitignore.io/issues/186#issuecomment-215987721
+
+# *.iml
+# modules.xml
+# .idea/misc.xml
+# *.ipr
+
+# Sonarlint plugin
+# https://plugins.jetbrains.com/plugin/7973-sonarlint
+.idea/**/sonarlint/
+
+# SonarQube Plugin
+# https://plugins.jetbrains.com/plugin/7238-sonarqube-community-plugin
+.idea/**/sonarIssues.xml
+
+# Markdown Navigator plugin
+# https://plugins.jetbrains.com/plugin/7896-markdown-navigator-enhanced
+.idea/**/markdown-navigator.xml
+.idea/**/markdown-navigator-enh.xml
+.idea/**/markdown-navigator/
+
+# Cache file creation bug
+# See https://youtrack.jetbrains.com/issue/JBR-2257
+.idea/$CACHE_FILE$
+
+# CodeStream plugin
+# https://plugins.jetbrains.com/plugin/12206-codestream
+.idea/codestream.xml
+
+# Azure Toolkit for IntelliJ plugin
+# https://plugins.jetbrains.com/plugin/8053-azure-toolkit-for-intellij
+.idea/**/azureSettings.xml
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+### Python Patch ###
+# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
+poetry.toml
+
+# ruff
+.ruff_cache/
+
+# LSP config files
+pyrightconfig.json
+
+### VisualStudioCode ###
+.vscode/*
+!.vscode/settings.json
+!.vscode/tasks.json
+!.vscode/launch.json
+!.vscode/extensions.json
+!.vscode/*.code-snippets
+
+# Local History for Visual Studio Code
+.history/
+
+# Built Visual Studio Code Extensions
+*.vsix
+
+### VisualStudioCode Patch ###
+# Ignore all local history of files
+.history
+.ionide
+
+### Windows ###
+# Windows thumbnail cache files
+Thumbs.db
+Thumbs.db:encryptable
+ehthumbs.db
+ehthumbs_vista.db
+
+# Dump file
+*.stackdump
+
+# Folder config file
+[Dd]esktop.ini
+
+# Recycle Bin used on file shares
+$RECYCLE.BIN/
+
+# Windows Installer files
+*.cab
+*.msi
+*.msix
+*.msm
+*.msp
+
+# Windows shortcuts
+*.lnk
+
+# End of https://www.toptal.com/developers/gitignore/api/python,java,c++,pycharm,visualstudiocode,macos,linux,windows
+
+.*
+flow_modules/
+data/

EvaluatorFlowModule/EvaluatorFlow.py

@@ -0,0 +1,216 @@
+""" NOTE: THIS IS A BETA VERSION OF FUNSEARCH. NEW VERSION DOCUMENTATION WILL BE RELEASED SOON."""
+
+from aiflows.base_flows import AtomicFlow
+from typing import Dict, Any
+import os
+from aiflows.utils import logging
+import ast
+import signal
+from aiflows.interfaces.key_interface import KeyInterface
+log = logging.get_logger(f"aiflows.{__name__}")
+import threading
+from aiflows.messages import FlowMessage
+class TimeoutException(Exception):
+    pass
+
+def timeout_handler(signum, frame):
+    raise TimeoutException("Execution timed out")
+
+class EvaluatorFlow(AtomicFlow):
+    """ This class implements an EvaluatorFlow. It is a flow that evaluates a program (python code) using a given evaluator function. This code is an implementation of Funsearch (https://www.nature.com/articles/s41586-023-06924-6) and is heavily inspired by the original code (https://github.com/google-deepmind/funsearch)
+    
+    **Configuration Parameters**:
+    
+    - `name` (str): The name of the flow. Default: "EvaluatorFlow"
+    - `description` (str): A description of the flow. This description is used to generate the help message of the flow. Default: "A flow that evaluates code on tests"
+    - `py_file` (str): The python code containing the evaluation function. No default value. This MUST be passed as a parameter.
+    - `function_to_run_name` (str): The name of the function to run (the evaluation function) in the evaluator file.  No default value. This MUST be passed as a parameter.
+    - `test_inputs` (Dict[str,Any]): A dictionary of test inputs to evaluate the program. Default: {"test1": None, "test2": None}
+    - `timeout_seconds` (int): The maximum number of seconds to run the evaluation function before returning an error. Default: 10
+    - `run_error_score` (int): The score to return if the evaluation function fails to run. Default: -100
+    - `use_test_input_as_key` (bool): Whether to use the test input parameters as the key in the output dictionary. Default: False
+    
+    **Input Interface**:
+    
+    - `artifact` (str): The program/artifact to evaluate.
+    
+    **Output Interface**:
+    
+    - `scores_per_test` (Dict[str, Dict[str, Any]]): A dictionary of scores per test input.
+    
+    **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 __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        self.evaluator_py_file = self.flow_config["py_file"]
+        self.run_error_score = self.flow_config["run_error_score"]
+
+        # Create a local namespace for the class
+        self.local_namespace = {}
+        self.load_functions()
+        self.function_to_run_name = self.flow_config["function_to_run_name"]
+        assert self.function_to_run_name in self.local_namespace, f"Function {self.function_to_run_name} not found in {self.evaluator_py_file_path}"
+        self.function_to_run = self.local_namespace.get(self.function_to_run_name)
+
+        self.test_inputs = self.flow_config["test_inputs"]
+        self.timeout_seconds = self.flow_config["timeout_seconds"]
+        self.local_namespace = {}
+        
+        
+        select_island_id_with_default = lambda data_dict,**kwargs: {**data_dict,**{"island_id":  data_dict.get("island_id", None)}} 
+        
+        self.output_interface = KeyInterface(
+            additional_transformations= [select_island_id_with_default],
+            keys_to_select= ["scores_per_test"]
+        )
+
+    
+    def load_functions(self):
+        """ Load the functions from the evaluator py file with ast parsing"""
+        
+        file_content = self.evaluator_py_file
+        try:
+            # Parse the AST (Abstract Syntax Tree) of the file content
+            parsed_ast = ast.parse(file_content)
+
+            # Iterate over the parsed AST nodes
+            for node in parsed_ast.body:
+                # Check if the node is an import statement
+                if isinstance(node, ast.Import):
+                    # Execute the import statement in the global namespace
+                    exec(compile(ast.Module(body=[node],type_ignores=[]), '<ast>', 'exec'), self.local_namespace)
+                elif isinstance(node, ast.ImportFrom):
+                    # Execute the import-from statement in the global namespace
+                    exec(compile(ast.Module(body=[node],type_ignores=[]), '<ast>', 'exec'), self.local_namespace)
+
+            # Execute the content of the file in the global namespace
+            exec(file_content, self.local_namespace)
+        except Exception as e:
+            log.error(f"Error functions: {e}")
+            raise e
+        
+    def run_function_with_timeout(self, program: str, **kwargs):
+        """ Run the evaluation function with a timeout
+        
+        :param program: The program to evaluate
+        :type program: str
+        :param kwargs: The keyword arguments to pass to the evaluation function
+        :type kwargs: Dict[str, Any]
+        :return: A tuple (bool, result) where bool is True if the function ran successfully and result is the output of the function
+        :rtype: Tuple[bool, Any]
+        """
+        self.result = None
+        self.exception = None
+
+        # Function to run with a timeout
+        def target():
+            try:
+                result = self.function_to_run(program, **kwargs)
+                self.result = result
+            except Exception as e:
+                self.exception = e
+
+        # Create a separate thread to run the target function
+        thread = threading.Thread(target=target)
+        thread.start()
+
+        # Wait for the specified timeout
+        thread.join(self.timeout_seconds)
+
+        # If thread is still alive, it means the timeout has occurred
+        if thread.is_alive():
+            # Raise a TimeoutException
+            thread.terminate()
+            return False, f"Function execution timed out after {self.timeout_seconds} seconds"
+
+        # If thread has finished execution, check if there was an exception
+        if self.exception is not None:
+            return False, str(self.exception)
+
+        # If no exception, return the result
+        return True, self.result
+
+ 
+
+    def evaluate_program(self, program: str, **kwargs):
+        """ Evaluate the program using the evaluation function
+        
+        :param program: The program to evaluate
+        :type program: str
+        :param kwargs: The keyword arguments to pass to the evaluation function
+        :type kwargs: Dict[str, Any]
+        :return: A tuple (bool, result) where bool is True if the function ran successfully and result is the output of the function
+        :rtype: Tuple[bool, Any]
+        """
+        try:
+            runs_ok, test_output = self.run_function_with_timeout(program, **kwargs)
+            
+            return runs_ok, test_output
+
+        except Exception as e:
+            log.debug(f"Error defining runnin program: {e} (could be due to syntax error from LLM)")
+            return False, e
+
+    
+
+    def analyse(self, program: str):
+        """ Analyse the program on the test inputs
+        
+        :param program: The program to evaluate
+        :type program: str
+        :return: A dictionary of scores per test input
+        :rtype: Dict[str, Dict[str, Any]]
+        """
+        #Often happens that it returns a codeblock so remove it
+        if program.startswith("```python"):
+            program = program[9:]
+        if program.endswith("```"):
+            program = program[:-3]
+        
+        scores_per_test = {}
+        for key,test_input in self.test_inputs.items():
+            
+            test_input_key = str(test_input) if self.flow_config["use_test_input_as_key"] else key
+            
+            if test_input is None:
+                runs_ok,test_output = self.evaluate_program(program)
+            else:
+                runs_ok,test_output = self.evaluate_program(program, **test_input)  # Run the program
+            
+            if runs_ok and test_output is not None:  # and not utils.calls_ancestor(program) (TODO: check what they mean by this in the paper)
+                scores_per_test[test_input_key] = {"score": test_output, "feedback": "No feedback available."}
+                log.debug(f"Program run successfully for test case {test_input_key} with score: {test_output}")
+            else:
+                log.debug(f"Error running Program for test case {test_input_key}. Error is : {test_output} (could be due to syntax error from LLM)")
+                scores_per_test[test_input_key] = {"score": self.run_error_score, "feedback": str(test_output)}
+
+        return scores_per_test
+
+    def run(self, input_message: FlowMessage):
+        """ This method runs the flow. It's the main method of the flow.
+        
+        :param input_message: The input message
+        :type input_message: FlowMessage
+        """
+        input_data = input_message.data
+        
+        # Analyse the program
+        scores_per_test = self.analyse(input_data["artifact"])
+        # Prepare the response
+        response = {"scores_per_test": scores_per_test, "from": "EvaluatorFlow"}
+        
+        # Send back the response
+        reply = self.package_output_message(
+            input_message,
+            response
+        )
+        self.send_message(reply)

EvaluatorFlowModule/EvaluatorFlow.yaml

@@ -0,0 +1,18 @@
+_target_: flow_modules.aiflows.FunSearchFlowModule.EvaluatorFlow.instantiate_from_default_config
+name: EvaluatorFlow
+description: A flow that evaluates code on tests
+
+input_interface:
+  - "artifact"
+output_interface:
+  - "scores_per_test"
+
+py_file: ???
+function_to_run_name: ???
+test_inputs: 
+  test_1: null
+  test_2: null
+timeout_seconds: 10
+run_error_score: -100
+use_test_input_as_key: false
+    

EvaluatorFlowModule/__init__.py

@@ -0,0 +1,13 @@
+# ~~~ Specify the dependencies ~~~
+# e.g.,
+# dependencies = [
+#     {"url": "aiflows/AutoGPTFlowModule", "revision": "main"},
+# ]
+# Revision can correspond toa branch, commit hash or a absolute path to a local directory (ideal for development)
+# from aiflows import flow_verse
+
+# flow_verse.sync_dependencies(dependencies)
+
+# ~~~ Import of your flow class (if you have any) ~~~
+# from .NAMEOFYOURFLOW import NAMEOFYOURFLOWCLASS
+from .EvaluatorFlow import EvaluatorFlow

EvaluatorFlowModule/demo.yaml

@@ -0,0 +1,111 @@
+_target_: flow_modules.aiflows.FunSearchFlowModule.EvaluatorFlow.instantiate_from_default_config
+
+py_file: |2-
+  """Problem Description:
+  Serval has a string s that only consists of 0 and 1 of length n. The i-th character of s is denoted as s_i, where 1\leq i\leq n.
+  Serval can perform the following operation called Inversion Magic on the string s:
+  Choose an segment [l, r] (1\leq l\leq r\leq n). For l\leq i\leq r, change s_i into 1 if s_i is 0, and change s_i into 0 if s_i is 1.
+  For example, let s be 010100 and the segment [2,5] is chosen. The string s will be 001010 after performing the Inversion Magic.
+  Serval wants to make s a palindrome after performing Inversion Magic exactly once. Help him to determine whether it is possible.
+  A string is a palindrome iff it reads the same backwards as forwards. For example, 010010 is a palindrome but 10111 is not.
+
+  Input Description:
+  Input
+  Each test contains multiple test cases. The first line contains the number of test cases t (1\leq t\leq 10^4). The description of the test cases follows.
+  The first line of each test case contains a single integer n (2\leq n\leq 10^5) — the length of string s.
+  The second line of each test case contains a binary string s of length n. Only characters 0 and 1 can appear in s.
+  It's guaranteed that the sum of n over all test cases does not exceed 2\cdot 10^5.
+
+  Output Description:
+  Output
+  For each test case, print Yes if s can be a palindrome after performing Inversion Magic exactly once, and print No if not.
+  You can output Yes and No in any case (for example, strings yEs, yes, Yes and YES will be recognized as a positive response).
+
+  Public Tests:
+  Test 1:
+    Input: ['1', '4', '1001']
+    Output: 'YES'
+  Test 2:
+    Input: ['1', '5', '10010']
+    Output: 'YES'
+  Test 3:
+    Input: ['1', '7', '0111011']
+    Output: 'NO'
+
+  """
+
+
+  import ast
+  import itertools
+  import numpy as np
+  from typing import List
+
+  def solve(solve_function: str,input: List[str], expected_output: str) -> str:
+      """function used to run the solve function on input *kwargs and return the the predicted output
+      
+      :param solve_function: the function to run (the solve function below as a string)
+      :type solve_function: str
+      :param kwargs: the inputs to the solve function
+      :type kwargs: List[str]
+      """
+      local_namespace = {}
+      exec(solve_function,local_namespace)
+      found_name, program_name = get_function_name_from_code(solve_function)
+      
+      if not found_name:
+          raise ValueError(f"Function name not found in program: {solve_function}")
+      
+      solve_fn = local_namespace.get(program_name)
+      
+      prediction = solve_fn(input)
+      
+      prediction = prediction.split()
+      expected_output = expected_output.split()
+      
+      if len(prediction) != len(expected_output):
+          raise ValueError(f"Invalid Format of prediction")
+      
+      for i in range(len(prediction)):
+          if prediction[i] != expected_output[i]:
+              return False
+      
+      return True
+
+  def evaluate(solve_function: str, tests_inputs: List[str], expected_outputs: str) -> float:
+      """Returns the score of the solve function we're evolving based on the tests_inputs and expected_outputs.
+      Scores are between 0 and 1, unless the program fails to run, in which case the score is -1.
+      """
+      if solve(solve_function,tests_inputs,expected_outputs) == True:
+          return 1.0
+      return 0.0
+
+
+  def get_function_name_from_code(code):
+      tree = ast.parse(code)
+      for node in ast.walk(tree):
+          if isinstance(node, ast.FunctionDef):
+              return True, node.name
+
+      # something is wrong
+      return False, None
+
+
+function_to_run_name: evaluate
+test_inputs: 
+
+  test_1: 
+    tests_inputs: ['1', '4', '1001']
+    expected_outputs: 'YES'
+  test_2: 
+    tests_inputs: ['1', '5', '10010'] 
+    expected_outputs: 'YES'
+  test_3: 
+    tests_inputs: ['1', '7', '0111011'] 
+    expected_outputs: 'NO'
+  test_4: 
+    tests_inputs: ['3', '4', '1001', '5', '10010', '7', '0111011', '']
+    expected_outputs: 'YES\nYES\nNO\n'
+
+timeout_seconds: 10
+run_error_score: -1
+use_test_input_as_key: false

EvaluatorFlowModule/pip_requirements.txt

@@ -0,0 +1 @@
+astunparse==1.6.3

EvaluatorFlowModule/run.py

@@ -0,0 +1,62 @@
+import os
+
+import hydra
+
+import aiflows
+from aiflows.backends.api_info import ApiInfo
+from aiflows.utils.general_helpers import read_yaml_file, quick_load_api_keys
+
+from aiflows import logging
+from aiflows.flow_cache import CACHING_PARAMETERS, clear_cache
+
+from aiflows.utils import serving
+from aiflows.workers import run_dispatch_worker_thread
+from aiflows.messages import FlowMessage
+from aiflows.interfaces import KeyInterface
+from aiflows.utils.colink_utils import start_colink_server
+from aiflows import flow_verse
+
+
+dependencies = [
+    {
+        "url": "aiflows/FunSearchFlowModule",
+        "revision": os.path.abspath("../")
+    }
+]
+flow_verse.sync_dependencies(dependencies)
+
+logging.set_verbosity_debug()
+
+
+if __name__ == "__main__":
+    
+    cl = start_colink_server()
+    
+    serving.recursive_serve_flow(
+        cl=cl,
+        flow_class_name="flow_modules.aiflows.FunSearchFlowModule.EvaluatorFlow",
+        flow_endpoint="EvaluatorFlow",
+    )
+        
+    run_dispatch_worker_thread(cl)
+    
+    config_overrides = read_yaml_file(os.path.join(".", "demo.yaml"))
+
+    funsearch_proxy = serving.get_flow_instance(
+        cl=cl,
+        flow_endpoint="EvaluatorFlow",
+        config_overrides=config_overrides,
+    )
+    data = {
+        'artifact': \
+            'def solve_function(input) -> str:\n    """Attempt at solving the problem given the input input and returns the predicted output (see the top of the file for problem description)"""\n    return \'YES\'\n'
+        }
+        
+    input_message = funsearch_proxy.package_input_message(data = data)
+    
+    funsearch_proxy.send_message(input_message)
+    
+    future = funsearch_proxy.get_reply_future(input_message)
+    response = future.get_data()
+    print("~~~Response~~~")
+    print(response)

FunSearch.py

@@ -0,0 +1,446 @@
+from aiflows.base_flows import CompositeFlow
+from aiflows.utils import logging
+from aiflows.interfaces import KeyInterface
+from aiflows.messages import FlowMessage
+from typing import Dict, Any
+log = logging.get_logger(f"aiflows.{__name__}")
+
+class FunSearch(CompositeFlow):
+    """ This class implements FunSearch. This code is an implementation of Funsearch (https://www.nature.com/articles/s41586-023-06924-6) and is heavily inspired by the original code (https://github.com/google-deepmind/funsearch) . It's a Flow in charge of starting, stopping and managing (passing around messages) the FunSearch process. It passes messages around to the following subflows:
+    
+    - ProgramDBFlow: which is in charge of storing and retrieving programs.
+    - SamplerFlow: which is in charge of sampling programs.
+    - EvaluatorFlow: which is in charge of evaluating programs.
+    
+    *Configuration Parameters*:
+    
+    - `name` (str): The name of the flow. Default: "FunSearchFlow".
+    - `description` (str): The description of the flow. Default: "A flow implementing FunSearch"
+    - `subflows_config` (Dict[str,Any]): A dictionary of subflows configurations. Default:
+        - `ProgramDBFlow`: By default, it uses the `ProgramDBFlow` class and uses its default parameters.
+        - `SamplerFlow`: By default, it uses the `SamplerFlow` class and uses its default parameters.
+        - `EvaluatorFlow`: By default, it uses the `EvaluatorFlow` class and uses its default parameters.
+        
+    **Input Interface**:
+    
+    - `from` (str): The flow from which the message is coming from. It can be one of the following: "FunSearch", "SamplerFlow", "EvaluatorFlow", "ProgramDBFlow".
+    - `operation` (str): The operation to perform. It can be one of the following: "start", "stop", "get_prompt", "get_best_programs_per_island", "register_program".
+    - `content` (Dict[str,Any]): The content associated to an operation. Here is the expected content for each operation:
+        - "start":
+            - `num_samplers` (int): The number of samplers to start up. Note that it's still restricted by the number of workers available. Default: 1.
+        - "stop":
+            - No content. Pass either an empty dictionary or None. Works also with no content.
+        - "get_prompt":
+            - No content. Pass either an empty dictionary or None. Works also with no content.
+        - "get_best_programs_per_island":
+            - No content. Pass either an empty dictionary or None. Works also with no content.
+    
+    **Output Interface**:
+    
+    - `retrieved` (Dict[str,Any]): The retrieved data.
+    
+    **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 __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        
+        #next state per action
+        #this is a dictionary that maps the next state of the flow based on the action and the current state
+        self.next_state_per_action = {
+            "get_prompt": {
+                "FunSearch": "ProgramDBFlow",
+                "ProgramDBFlow": "SamplerFlow",
+            },
+            
+            "get_best_programs_per_island": {
+                "FunSearch": "ProgramDBFlow",
+                "ProgramDBFlow": "GenerateReply",
+            },
+            "register_program": {
+                "SamplerFlow": "EvaluatorFlow",
+                "EvaluatorFlow": "ProgramDBFlow",
+            },
+            "start":
+                {"FunSearch": "FunSearch"},
+            "stop":
+                {"FunSearch": "FunSearch"},
+        }
+        
+        #key interface to make a request for a prompt
+        self.make_request_for_prompt_data = KeyInterface(
+            keys_to_set= {"operation": "get_prompt", "content": {}, "from": "FunSearch"},
+            keys_to_select= ["operation", "content", "from"]
+        )
+    
+    def make_request_for_prompt(self):
+        """ This method makes a request for a prompt. It sends a message to itself with the operation "get_prompt" which will trigger the flow to call the `ProgramDBFlow` to get a prompt. """
+        
+        #Prepare data to make request for prompt
+        data = self.make_request_for_prompt_data({})
+        
+        #Package message to make request for prompt
+        msg = self.package_input_message(
+            data=data,
+            dst_flow="FunSearch"
+        )
+        #Send message to itself to start the process of getting a prompt
+        self.send_message(
+            msg
+        )
+        
+    def request_samplers(self,input_message: FlowMessage):
+        """ This method requests samplers. It sends a message to itself with the operation "get_prompt" which will trigger the flow to call the `ProgramDBFlow` to get a prompt.
+        
+        :param input_message: The input message that triggered the request for samplers.
+        :type input_message: FlowMessage
+        """
+        
+        #Get state associated with the message
+        message_state = self.pop_message_from_state(input_message.input_message_id)
+        #Get number of samplers to request
+        num_samplers = message_state["content"].get("num_samplers",1)
+        for i in range(num_samplers):
+            self.make_request_for_prompt()
+    
+    def get_next_state(self, input_message: FlowMessage):
+        """ This method determines the next state of the flow based on the input message. It will return the next state based on the current state and the message received.
+        
+        :param input_message: The input message that triggered the request for the next state.
+        :type input_message: FlowMessage
+        :return: The next state of the flow.
+        :rtype: str
+        """
+        #Get state associated with the message
+        message_state = self.get_message_from_state(input_message.input_message_id)
+        message_from = message_state["from"]
+        operation = message_state["operation"]
+        #Get next state based on the action and the current state
+        next_state = self.next_state_per_action[operation][message_from]
+        return next_state 
+    
+    def set_up_flow_state(self):
+        """ This method sets up the state of the flow. It's called at the beginning of the flow."""
+        super().set_up_flow_state()
+        #Dictonary containing state of message currently being handled by FunSearch
+        #Each message has its own state in the flow state
+        #Once a message is done being handled, it's removed from the state
+        self.flow_state["msg_requests"] = {}
+        #Flag to keep track if the first sample has been saved to the db
+        self.flow_state["first_sample_saved_to_db"] = False
+        #Flag to keep track if FunSearch is running
+        self.flow_state["funsearch_running"] = False
+        
+    def save_message_to_state(self,msg_id: str, message: FlowMessage):
+        """ This method saves a message to the state of the flow. It's used to keep track of state on a per message basis (i.e., state of the flow depending on the message received and id).
+        
+        :param msg_id: The id of the message to save.
+        :type msg_id: str
+        :param message: The message to save.
+        :type message: FlowMessage
+        """
+        self.flow_state["msg_requests"][msg_id] = {"og_message":  message}
+        
+    def rename_key_message_in_state(self, old_key: str, new_key: str):
+        """ This method renames a key in the state of the flow in the "msg_requests" dictonary. It's used to rename a key in the state of the flow (i.e., rename a message id).
+        
+        :param old_key: The old key to rename.
+        :type old_key: str
+        :param new_key: The new key to rename to.
+        :type new_key: str
+        """
+        self.flow_state["msg_requests"][new_key]  = self.flow_state["msg_requests"].pop(old_key) 
+
+    def message_in_state(self,msg_id: str) -> bool:
+        """ This method checks if a message is in the state of the flow (in "msg_requests" dictionary). It returns True if the message is in the state, otherwise it returns False.
+        
+        :param msg_id: The id of the message to check if it's in the state.
+        :type msg_id: str
+        :return: True if the message is in the state, otherwise False.
+        :rtype: bool
+        """
+        
+        return msg_id in self.flow_state["msg_requests"].keys()
+
+    def get_message_from_state(self, msg_id: str) -> Dict[str,Any]:
+        """ This method returns the state associated with a message id in the state of the flow (in "msg_requests" dictionary).
+        
+        :param msg_id: The id of the message to get the state from.
+        :type msg_id: str
+        :return: The state associated with the message id.
+        :rtype: Dict[str,Any]
+        """
+        return self.flow_state["msg_requests"][msg_id]
+
+    def pop_message_from_state(self, msg_id: str) -> Dict[str,Any]:
+        """ This method pops a message from the state of the flow (in "msg_requests" dictionary). It the state associate to a message and removes it from the state.
+        
+        :param msg_id: The id of the message to pop from the state.
+        :type msg_id: str
+        :return: The state associated with the message id.
+        :rtype: Dict[str,Any]
+        """
+        return self.flow_state["msg_requests"].pop(msg_id)
+    
+    def merge_message_request_state(self,id: str, new_states: Dict[str,Any]):
+        """ This method merges new states to a message in the state of the flow (in "msg_requests" dictionary). It merges new states to a message in the state.
+        
+        :param id: The id of the message to merge new states to.
+        :type id: str
+        :param new_states: The new states to merge to the message.
+        :type new_states: Dict[str,Any]
+        """
+        self.flow_state["msg_requests"][id] = {**self.flow_state["msg_requests"][id], **new_states}
+    
+    def register_data_to_state(self, input_message: FlowMessage):
+        """This method registers the input message data to the flow state. It's called everytime a new input message is received.
+        
+        :param input_message: The input message
+        :type input_message: FlowMessage
+        """
+
+        #Determine Who the message is from (should be either FunSearch, SamplerFlow, EvaluatorFlow, or ProgramDBFlow)
+        msg_from = input_message.data.get("from", "FunSearch")
+        #Check if this a first request or part of a message that is being handled (it's part of message being handled if message is in the state)
+        msg_id = input_message.input_message_id   
+        msg_in_state = self.message_in_state(msg_id)
+        
+        #If message is not in state, save it to state
+        if not msg_in_state:
+            self.save_message_to_state(msg_id, input_message)
+        
+        #Get the state associated to the message
+        message_state = self.get_message_from_state(msg_id)
+        
+        #Determine what to do based on who the message is from
+        
+        if msg_from == "FunSearch":
+            #Calls From FunSearch expect operation and content
+            operation = input_message.data["operation"]
+            content = input_message.data.get("content",{})
+            to_add_to_state = {
+                "content": content,
+                "operation": operation
+            }
+            #save operation and content to state   
+            self.merge_message_request_state(msg_id, to_add_to_state)
+            
+        elif msg_from == "SamplerFlow":
+            #Calls From SamplerFlow expect api_output, merge it to state
+            to_add_to_state = {
+                "content": {
+                    **message_state.get("content",{}),
+                    **{"artifact": input_message.data["api_output"]}
+                },
+                "operation": "register_program"
+            }
+            self.merge_message_request_state(msg_id, to_add_to_state)
+                        
+        elif msg_from == "EvaluatorFlow":
+            #Calls From EvaluatorFlow expect scores_per_test, merge it to state
+            message_state = self.get_message_from_state(msg_id)
+            to_add_to_state = {
+                "content": {
+                    **message_state.get("content",{}),
+                    **{"scores_per_test": input_message.data["scores_per_test"]}
+                }
+            }
+            self.merge_message_request_state(msg_id, to_add_to_state)
+                        
+        elif msg_from == "ProgramDBFlow":
+            #Calls From ProgramDBFlow expect retrieved, merge it to state
+            to_add_to_state = {
+                "retrieved": input_message.data["retrieved"],
+            }
+            
+            #if message from ProgramDBFlow is associate to a "get_prompt" operation, 
+            # save island_id to state
+            if message_state["operation"] == "get_prompt":
+                island_id = input_message.data["retrieved"]["island_id"]
+                to_add_to_state["content"] = {
+                    **message_state.get("content",{}),
+                    **{"island_id": island_id}
+                }
+                
+            self.merge_message_request_state(msg_id, to_add_to_state) 
+        
+        #save from to state
+        self.merge_message_request_state(msg_id, {"from": msg_from})
+
+    def call_program_db(self, input_message):
+        """ This method calls the ProgramDBFlow. It sends a message to the ProgramDBFlow with the data of the input message.
+        
+        :param input_message: The input message to send to the ProgramDBFlow.
+        :type input_message: FlowMessage
+        """
+        
+        #Fetch state associated with the message
+        msg_id = input_message.input_message_id
+        message_state = self.get_message_from_state(input_message.input_message_id)
+
+        #Get operation and content from state to send to ProgramDBFlow
+        operation = message_state["operation"]
+        content = message_state.get("content", {})
+        
+        data = {
+            "operation": operation,
+            "content": content
+        }
+        #package message to send to ProgramDBFlow
+        msg = self.package_input_message(
+            data = data,
+            dst_flow = "ProgramDBFlow"
+        )      
+        
+        #If operation is "register_program",
+        # pop message from state (because inital message has been fully handled) and set first_sample_saved_to_db to True
+        #Send a message to register program without expecting a reply (no need to wait for a reply, just save to db and move on)
+        if data["operation"] == "register_program":
+            self.pop_message_from_state(msg_id)
+            
+            self.flow_state["first_sample_saved_to_db"] = True
+            
+            self.subflows["ProgramDBFlow"].send_message(
+                msg
+            )
+        
+        # If operation is "get_prompt" or "get_best_programs_per_island"
+        # rename key in state to new message id (in order to be able to track of the message in state when the reply arrives)
+        elif data["operation"] in ["get_prompt","get_best_programs_per_island"]:
+            self.rename_key_message_in_state(msg_id, msg.message_id)  
+            #if no sample has been saved to db, Send input message back to itself (to try again, hopefully this time a sample will be saved to db)
+            if not self.flow_state["first_sample_saved_to_db"]:
+                #send back to itself message (to try again)
+                self.send_message(
+                    input_message
+                )
+            #If a sample has been saved to db, send message to ProgramDBFlow to fetch prompt or best programs per island
+            else:
+                self.subflows["ProgramDBFlow"].get_reply(
+                    msg
+                )
+        #If operation is not "register_program", "get_prompt" or "get_best_programs_per_island"
+        else:
+            log.error("No operation found, input_message received: \n" + str(input_message))
+        
+    def call_evaluator(self, input_message):
+        """ This method calls the EvaluatorFlow. It sends a message to the EvaluatorFlow with the data of the input message.
+        
+        :param input_message: The input message to send to the EvaluatorFlow.
+        :type input_message: FlowMessage
+        """
+    
+        #Fetch state associated with the message
+        msg_id = input_message.input_message_id 
+        message_state = self.get_message_from_state(msg_id)
+        
+        #Get data to send to EvaluatorFlow (artifact generated by Sampler to be evaluated)
+        data = {
+            "artifact": message_state["content"]["artifact"]
+        }
+            
+        msg = self.package_input_message(
+            data = data,
+            dst_flow = "EvaluatorFlow"
+        )
+        # rename key in state to new message id (in order to be able to track of the message in state when the reply arrives)
+        self.rename_key_message_in_state(msg_id, msg.message_id)  
+        #Send message to EvaluatorFlow and expect a reply to be sent back to FunSearch's input message queue    
+        self.subflows["EvaluatorFlow"].get_reply(
+            msg
+        )
+        
+    def call_sampler(self, input_message):
+        """ This method calls the SamplerFlow. It sends a message to the SamplerFlow with the data of the input message.
+        
+        :param input_message: The input message to send to the SamplerFlow.
+        :type input_message: FlowMessage
+        """
+        
+        #Fetch state associated with the message
+        msg_id = input_message.input_message_id
+        message_state = self.get_message_from_state(msg_id)
+        
+        #Get data to send to SamplerFlow (prompt to generate a program)
+        data = {
+            **message_state["retrieved"],
+        }
+        msg = self.package_input_message(
+            data = data,
+            dst_flow = "SamplerFlow"
+        )
+        # rename key in state to new message id (in order to be able to track of the message in state when the reply arrives)
+        self.rename_key_message_in_state(msg_id, msg.message_id)        
+        
+        #send message to SamplerFlow and expect a reply to be sent back to FunSearch's input message queue
+        self.subflows["SamplerFlow"].get_reply(
+            msg
+        )
+        #If FunSearch is running, make a new request for a prompt (to keep the process going)
+        if self.flow_state["funsearch_running"]:
+            self.make_request_for_prompt()
+        
+        
+    def generate_reply(self, input_message: FlowMessage):
+        """ This method generates a reply to a message sent to user. It packages the output message and sends it.
+        
+        :param input_message: The input message to generate a reply to.
+        :type input_message: FlowMessage
+        """
+        
+        #Fetch state associated with the message
+        msg_id = input_message.input_message_id
+        message_state = self.pop_message_from_state(msg_id)
+        #Prepare response to send to user (due to a call to get_best_programs_per_island)
+        response = {
+            "retrieved": message_state["retrieved"]
+        }
+        reply = self.package_output_message(
+            message_state["og_message"],
+            response
+        )
+        
+        self.send_message(
+            reply
+        )
+        
+    def run(self,input_message: FlowMessage):
+        """ This method runs the flow. It's the main method of the flow. It's called when the flow is executed.
+        
+        :input_message: The input message of the flow
+        :type input_message: Message
+        """
+        self.register_data_to_state(input_message)
+        
+        next_state = self.get_next_state(input_message)
+        
+        if next_state == "ProgramDBFlow":
+            self.call_program_db(input_message)
+            
+        elif next_state == "EvaluatorFlow":
+            self.call_evaluator(input_message)
+            
+        elif next_state == "SamplerFlow":
+            self.call_sampler(input_message)
+        
+        elif next_state == "GenerateReply":
+            self.generate_reply(input_message)
+            
+        elif next_state == "FunSearch":
+            #If operation is "start", set funsearch_running to True and make a request for a prompt
+            if input_message.data["operation"] == "start":
+                self.flow_state["funsearch_running"] = True
+                self.request_samplers(input_message)
+            #If operation is "stop", set funsearch_running to False (will stop the process of generating new samples)
+            elif input_message.data["operation"] == "stop":
+                self.flow_state["funsearch_running"] = False
+        else:
+            log.error("No next state found, input_message received: \n" + str(input_message))    

FunSearch.yaml

@@ -0,0 +1,37 @@
+name: FunSearchFlow
+description: "A flow implementing FunSearch"
+_target_: flow_modules.aiflows.FunSearchFlowModule.FunSearch.instantiate_from_default_config
+user_id: local
+input_interface:
+  - "action"
+  - "content"
+  - "from"
+
+output_interface:
+  - "retrieved"
+### Subflows specification
+subflows_config:
+  SamplerFlow:
+    flow_class_name: flow_modules.aiflows.FunSearchFlowModule.SamplerFlow
+    flow_endpoint: SamplerFlow
+    parallel_dispatch: True
+    singleton: False
+    user_id: local
+    name: "Sampler Flow"
+    description: "A flow that queries an LLM model to generate prompts"
+
+  EvaluatorFlow:
+    flow_class_name: flow_modules.aiflows.FunSearchFlowModule.EvaluatorFlow
+    flow_endpoint: EvaluatorFlow
+    user_id: local
+    parallel_dispatch: True
+    name: "A flow that evaluates code on tests"
+    description: "A flow that evaluates code on tests"
+
+  ProgramDBFlow:
+    flow_class_name: flow_modules.aiflows.FunSearchFlowModule.ProgramDBFlow
+    flow_endpoint: ProgramDBFlow
+    singleton: True 
+    user_id: local
+    name: "ProgramDB"
+    description: "A flow that registers samples and evaluations in a database"

Loader.py

@@ -0,0 +1,55 @@
+import ast
+import os
+import yaml
+
+class Loader:
+    def __init__(self, file_path, target_name):
+        self.py_file_path = file_path
+        self.target_name = target_name
+
+        if not os.path.exists(file_path):
+            raise ValueError(f"File {file_path} does not exist")
+
+        with open(file_path, 'r') as file:
+            self.source_code = file.read()
+
+    def load_target(self):
+        if self.py_file_path.endswith('.yaml'):
+            return self.load_yaml()
+        else:
+            return self.load_code()
+    
+    def load_full_file(self):
+        return self.source_code
+
+    def load_code(self):
+        # Parse the source code into an abstract syntax tree (AST)
+        tree = ast.parse(self.source_code)
+
+        # Find the target node (FunctionDef, ClassDef, or variable)
+        target_node = None
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and node.name == self.target_name:
+                target_node = node
+                break
+            elif isinstance(node, ast.Assign):
+                for target in node.targets:
+                    if isinstance(target, ast.Name) and target.id == self.target_name:
+                        target_node = node
+                        break
+
+        if target_node is not None:
+            # Extract the source code of the target
+            target_source_code = ast.unparse(target_node)
+            return target_source_code
+        else:
+            raise ValueError(f"Target '{self.target_name}' not found in the module.")
+
+    def load_yaml(self):
+        try:
+            with open(self.py_file_path, 'r') as yaml_file:
+                yaml_content = yaml.safe_load(yaml_file)
+            return yaml_content
+        except yaml.YAMLError as e:
+            raise ValueError(f"Error loading YAML file: {e}")
+

ProgramDBFlowModule/Cluster.py

@@ -0,0 +1,68 @@
+import numpy as np
+from .artifacts import AbstractArtifact
+from .utils import _softmax
+class Cluster:
+    """ An implementation of a Cluster of an island.  This code is an implementation of Funsearch (https://www.nature.com/articles/s41586-023-06924-6) and is heavily inspired by the original code (https://github.com/google-deepmind/funsearch)
+
+    **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 __init__(self,score: float,first_program: AbstractArtifact,epsilon=1e-6,sample_with_replacement=False, default_program_temperature=0.1):
+        self.score: float = score
+        self.programs: list[AbstractArtifact] = [first_program]
+        self.lengths = np.array([len(str(first_program))],dtype=np.float32)
+        self.epsilon = epsilon
+        self.sample_with_replacement = sample_with_replacement
+        self.default_program_temperature = default_program_temperature
+        
+        
+    def compute_length_probs(self,program_temperature: float):
+        """ Compute the probability of each program given the length of the program. The probability is computed as the softmax of the negative length of the program. The temperature of the softmax is controlled by the program_temperature parameter.
+        
+        :param program_temperature: The temperature of the softmax
+        :type program_temperature: float
+        :return: The probability of each program given the length of the program
+        :rtype: np.array
+        """
+        min_length = np.min(self.lengths)
+        max_length = np.max(self.lengths)
+        
+
+        length_logits = (self.lengths - min_length)/(max_length  + self.epsilon)
+        
+        probs =  _softmax(-length_logits,program_temperature)
+        return probs
+    
+    def register_program(self,program: str):
+        """ Register a program on the cluster.
+        
+        :param program: The program to register
+        :type program: str
+        """
+        self.programs.append(program)
+        self.lengths = np.append(self.lengths,len(str(program)))
+        
+    def sample_program(self,program_temperature=None):
+        """ Sample a program from the cluster given the program temperature.
+        
+        :param program_temperature: The temperature of the program
+        :type program_temperature: float, optional
+        :return: The sampled program
+        :rtype: str
+        """
+        if program_temperature is None:
+            program_temperature = self.default_program_temperature
+        
+        
+        probs = self.compute_length_probs(program_temperature)
+        #sample an index of probs randomly givent the probs
+        index = np.random.choice(len(probs),p=probs,replace=self.sample_with_replacement)
+        
+        return self.programs[index]

ProgramDBFlowModule/Island.py

@@ -0,0 +1,168 @@
+#Heavily Inpired by https://github.com/google-deepmind/funsearch/tree/main
+
+
+import numpy as np
+from typing import Callable
+from .Cluster import Cluster
+from .Program import Program
+import ast
+import astunparse
+from typing import  Optional,Dict,Any
+from .artifacts import AbstractArtifact
+from collections.abc import Mapping, Sequence
+from copy import deepcopy
+from .Program import ProgramVisitor,Program,text_to_artifact
+import dataclasses
+import scipy
+from .utils import _softmax
+ScoresPerTest = Mapping
+
+
+
+class Island:
+    """ An implementation of an Island of the ProgramDB.  This code is an implementation of Funsearch (https://www.nature.com/articles/s41586-023-06924-6) and is heavily inspired by the original code (https://github.com/google-deepmind/funsearch)
+
+    **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 __init__(self,
+                 artifact_to_evolve_name: str,
+                 artifacts_per_prompt: int,
+                 temperature: float,
+                 temperature_period: int,
+                 template: Program,
+                 reduce_score_method:  Optional[Callable] = np.mean,
+                 sample_with_replacement: Optional[bool] = False):
+        
+        self.artifact_to_evolve_name:str = artifact_to_evolve_name
+        self.artifacts_per_prompt: int = artifacts_per_prompt
+        
+        self.temperature: float = temperature
+        self.temperature_period: int = temperature_period
+        
+        self.clusters: Dict[str,Cluster] = {}
+        self.template: Program = template
+        
+        self.total_programs_on_island: int = 0
+        self.reduce_score_method: Callable = reduce_score_method
+        self.sample_with_replacement: bool = sample_with_replacement
+        
+    def register_program(self, program: AbstractArtifact ,scores_per_test: ScoresPerTest):
+        """ Register a program on the island.
+        
+        :param program: The program to register
+        :type program: AbstractArtifact
+        :param scores_per_test: The scores per test of the program
+        :type scores_per_test: Dict[str,Any]
+        """
+        scores_per_test_key = " ".join([str(key) + ":" + str(score) for key,score in scores_per_test.items()])
+        
+        scores_per_test_values = np.array([ score_per_test["score"] for score_per_test in scores_per_test.values()])
+        
+        if scores_per_test_key not in self.clusters:
+            score = self.reduce_score_method(scores_per_test_values)
+            self.clusters[scores_per_test_key] = Cluster(score = score,first_program=program,sample_with_replacement=self.sample_with_replacement)
+        
+        else:
+            self.clusters[scores_per_test_key].register_program(program)
+            
+        self.total_programs_on_island += 1
+    
+    def pick_clusters(self):
+        """ Pick the clusters to generate the prompt
+        
+        :return: The clusters and their names
+        :rtype: Tuple[List[Cluster],List[str]]
+        """
+        cluster_keys = list(self.clusters.keys())
+        clusters = [self.clusters[key] for key in cluster_keys]
+        
+        cluster_scores = np.array([cluster.score for cluster in clusters])
+        cluster_temperature = self.temperature * (1 - (self.total_programs_on_island % self.temperature_period) / self.temperature_period)
+        probs = _softmax(cluster_scores, cluster_temperature)
+
+        #can occur at the begining when there are not many clusters
+        functions_per_prompt = min(self.artifacts_per_prompt,len(self.clusters))
+        select_cluster_ids = np.random.choice(len(cluster_scores),size=functions_per_prompt,p = probs,replace=self.sample_with_replacement)
+        
+        return [clusters[cluster_id] for cluster_id in select_cluster_ids], [cluster_keys[cluster_id] for cluster_id in select_cluster_ids]
+    
+    def _get_versioned_artifact_name(self,i):
+        """ Get the versioned artifact name
+        
+        :param i: The version of the artifact
+        :type i: int
+        :return: The versioned artifact name
+        :rtype: str
+        """
+        return self.artifact_to_evolve_name  + "_v" +str(i)
+    
+    def _generate_prompt(self,implementations: Sequence[AbstractArtifact], chosen_cluster_names: Sequence[str]):
+        """ Generate the prompt
+        
+        :param implementations: The implementations
+        :type implementations: Sequence[AbstractArtifact]
+        :param chosen_cluster_names: The chosen cluster names
+        :type chosen_cluster_names: Sequence[str]
+        :return: The prompt
+        :rtype: str
+        """
+        implementations = deepcopy(implementations)
+        
+        versioned_artifacts: list[AbstractArtifact] = []
+        
+        for i,implementation in enumerate(implementations):
+            new_artifact_name = self._get_versioned_artifact_name(i)
+            implementation.name = new_artifact_name
+            score_per_test = " ".join(chosen_cluster_names[i].split(" "))
+            implementation.docstring = f'Scores per test: {score_per_test}'
+            if i >= 1:
+                implementation.docstring += f'\nImproved version of {self._get_versioned_artifact_name(i-1)}'
+            implementation = implementation.rename_artifact_calls(source_name = self.artifact_to_evolve_name, target_name = new_artifact_name)
+            versioned_artifacts.append(text_to_artifact(implementation))
+       
+        #Create the header of the function to be generated by the LLM
+        next_version = len(implementations)
+        new_artifact_name = self._get_versioned_artifact_name(next_version)
+        
+        docstring = f'Improved version of {self._get_versioned_artifact_name(next_version-1)}'
+        
+        header = dataclasses.replace(
+            implementations[-1],
+            name=new_artifact_name,
+            body='',
+            docstring=docstring
+        )
+        
+        versioned_artifacts.append(header)
+        #rename the call to the target function
+        prompt = dataclasses.replace(self.template, artifacts=versioned_artifacts)
+        return str(prompt)
+        
+    def get_prompt(self):
+        """ Get the prompt
+        
+        :return: The prompt
+        :rtype: str
+        """
+        chosen_clusters, chosen_cluster_names = self.pick_clusters()
+        
+        scores = [cluster.score for cluster in chosen_clusters]
+        
+        indices = np.argsort(scores)
+        
+        sorted_implementations = [chosen_clusters[i].sample_program() for i in indices]
+        sorted_cluster_names = [chosen_cluster_names[i] for i in indices]
+        
+        version_generated = len(sorted_implementations)
+        return self._generate_prompt(sorted_implementations,sorted_cluster_names),version_generated
+        
+            
+        

ProgramDBFlowModule/Program.py

@@ -0,0 +1,110 @@
+import dataclasses
+from typing import Optional
+from .artifacts import AbstractArtifact
+import ast
+from .artifacts import FunctionArtifact
+
+@dataclasses.dataclass(frozen=True)
+class Program:
+  """A parsed Python program."""
+
+  # `preface` is everything from the beginning of the code till the first
+  # function is found.
+  preface: str
+  artifacts: list[AbstractArtifact]
+
+  def __str__(self) -> str:
+    program = f'{self.preface}\n' if self.preface else ''
+    program += '\n'.join([str(f) for f in self.artifacts])
+    return program
+
+  def find_artifact_index(self, artifact_name: str) -> int:
+    """Returns the index of input function name."""
+    
+    artifact_names = [a.name for a in self.artifacts]
+    count = artifact_names.count(artifact_name)
+    if count == 0:
+      raise ValueError(
+          f'artifact {artifact_name} does not exist in program:\n{str(self)}'
+      )
+    if count > 1:
+      raise ValueError(
+          f'artifact {artifact_name} exists more than once in program:\n'
+          f'{str(self)}'
+      )
+    index = artifact_names.index(artifact_name)
+    return index
+
+  def get_artifact(self, artifact_name: str) -> AbstractArtifact:
+    index = self.find_artifact_index(artifact_name)
+    return self.artifacts[index]
+  
+# TODO: Do this for various types of artifacts (only for functions rn)
+class ProgramVisitor(ast.NodeVisitor):
+    """Parses code to collect all required information to produce a `Program`.
+
+    Note that we do not store function decorators.
+    """
+
+    def __init__(self, sourcecode: str):
+        self._codelines: list[str] = sourcecode.splitlines()
+
+        self._preface: str = ''
+        self._artifacts: list[AbstractArtifact] = []
+        self._current_artifact: Optional[str] = None
+
+    def visit_FunctionDef(self,
+                        node: ast.FunctionDef):
+        """Collects all information about the function being parsed."""
+        if node.col_offset == 0:  # We only care about first level functions.
+            self._current_function = node.name
+            if not self._artifacts:
+                self._preface = '\n'.join(self._codelines[:node.lineno - 1])
+                
+            function_end_line = node.end_lineno
+            body_start_line = node.body[0].lineno - 1
+            # Extract the docstring.
+            docstring = None
+            if isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value,ast.Str):
+                docstring = f'  """{ast.literal_eval(ast.unparse(node.body[0]))}"""'
+                if len(node.body) > 1:
+                    body_start_line = node.body[1].lineno - 1
+                else:
+                    body_start_line = function_end_line
+
+            self._artifacts.append(FunctionArtifact(
+                name=node.name,
+                args=ast.unparse(node.args),
+                return_type=ast.unparse(node.returns) if node.returns else None,
+                docstring=docstring,
+                body='\n'.join(self._codelines[body_start_line:function_end_line]),
+            ))
+        self.generic_visit(node)
+
+    def return_program(self) -> Program:
+        return Program(preface=self._preface, artifacts=self._artifacts)
+      
+
+def text_to_program(text: str) -> Program:
+  """Returns Program object by parsing input text using Python AST."""
+  # We assume that the program is composed of some preface (e.g. imports,
+  # classes, assignments, ...) followed by a sequence of functions.
+  
+  #Often happens that it returns a codeblock so remove it
+  if text.startswith("```python"):
+      text = text[9:]
+  if text.endswith("```"):
+      text = text[:-3]
+  
+  tree = ast.parse(text)
+  visitor = ProgramVisitor(text)
+  visitor.visit(tree)
+  return visitor.return_program()
+
+def text_to_artifact(text: str) -> AbstractArtifact:
+  """Returns Function object by parsing input text using Python AST."""
+  program = text_to_program(text)
+  if len(program.artifacts) != 1:
+    raise ValueError(f'Only one artifact expected, got {len(program.artifacts)}'
+                        f':\n{program.functions}')
+  return program.artifacts[0]

ProgramDBFlowModule/ProgramDBFlow.py

@@ -0,0 +1,271 @@
+""" NOTE: THIS IS A BETA VERSION OF FUNSEARCH. NEW VERSION DOCUMENTATION WILL BE RELEASED SOON."""
+
+
+from aiflows.base_flows import AtomicFlow
+from .Island import Island,ScoresPerTest
+import numpy as np
+from typing import Callable,Dict,Union, Any,Optional, List
+import time
+from aiflows.utils import logging
+from .artifacts import AbstractArtifact
+from .Program import Program,text_to_artifact
+import ast
+import os
+from aiflows.messages import FlowMessage
+log = logging.get_logger(f"aiflows.{__name__}")
+
+class ProgramDBFlow(AtomicFlow):
+    """ This class implements a ProgramDBFlow. It's a flow that stores programs and their scores in a database. It can also query the database for the best programs or generate a prompt containing stored programs in order to evolve them with a SamplerFlow. This code is an implementation of Funsearch (https://www.nature.com/articles/s41586-023-06924-6) and is heavily inspired by the original code (https://github.com/google-deepmind/funsearch)
+    
+    **Configuration Parameters**:
+    
+    - `name` (str): The name of the flow. Default: "ProgramDBFlow"
+    - `description` (str): A description of the flow. This description is used to generate the help message of the flow. Default: " A flow that saves programs in a database of islands"
+    - `artifact_to_evolve_name` (str): The name of the artifact/program to evolve. Default: "solve_function"
+    - `evaluate_function` (str): The function used to evaluate the program. No Default value. This MUST be passed as a parameter.
+    - `evaluate_file_full_content` (str): The full content of the file containing the evaluation function. No Default value. This MUST be passed as a parameter.
+    - `num_islands`: The number of islands to use. Default: 3
+    - `reset_period`: The period in seconds to reset the islands. Default: 3600
+    - `artifacts_per_prompt`: The number of previous artifacts/programs to include in a prompt. Default: 2
+    - `temperature`: The temperature of the island. Default: 0.1
+    - `temperature_period`: The period in seconds to change the temperature. Default: 30000
+    - `sample_with_replacement`: Whether to sample with replacement. Default: False
+    - `portion_of_islands_to_reset`: The portion of islands to reset. Default: 0.5
+    - `template` (dict): The template to use for a program. Default: {"preface": ""}
+
+    **Input Interface**:
+    
+    - `operation` (str): The operation to perform. It can be one of the following: ["register_program","get_prompt","get_best_programs_per_island"]
+    
+    **Output Interface**:
+    
+    - `retrieved` (Any): The retrieved data. It can be one of the following:
+        - If the operation is "get_prompt", it can be a dictionary with the following keys
+            - `code` (str): The code of the prompt
+            - `version_generated` (int): The version of the prompt generated
+            - `island_id` (int): The id of the island that generated the prompt
+            - `header` (str): The header of the prompt
+        - If the operation is "register_program", it can be a string with the message "Program registered" or "Program failed to register"
+        - If the operation is "get_best_programs_per_island", it can be a dictionary with the following keys:
+            - `best_island_programs` (List[Dict[str,Any]]): A list of dictionaries with the following keys:
+                - `rank` (int): The rank of the program (1 is the best)
+                - `score` (float): The score of the program
+                - `program` (str): The program
+                - `island_id` (int): The id of the island that generated the program
+                
+    **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 __init__(self,
+                 **kwargs
+                ):
+        super().__init__(**kwargs)
+        
+        #Unpack config (for clarity)
+        self.artifact_to_evolve_name: str = self.flow_config["artifact_to_evolve_name"]
+        
+        self.artifacts_per_prompt: int = self.flow_config["artifacts_per_prompt"]
+        self.temperature: float = self.flow_config["temperature"]
+        self.temperature_period: int = self.flow_config["temperature_period"]
+        self.reduce_score_method: Callable = np.mean 
+        self.sample_with_replacement: bool = self.flow_config["sample_with_replacement"]
+        self.num_islands: int = self.flow_config["num_islands"]
+        self.portion_of_islands_to_reset: float = self.flow_config["portion_of_islands_to_reset"]
+        self.reset_period: float = self.flow_config["reset_period"]
+        
+        self.evaluate_function = self.flow_config["evaluate_function"]
+        self.evaluate_file_full_content = self.flow_config["evaluate_file_full_content"] 
+        
+        
+        
+        assert self.portion_of_islands_to_reset <= 1.0 and self.portion_of_islands_to_reset >= 0.0, "portion_of_islands_to_reset must be between 0 and 1"
+        
+        self.islands_to_reset = int(round(self.portion_of_islands_to_reset * self.num_islands)) #round to nearest integer
+        
+          
+    def set_up_flow_state(self):
+        """ This method sets up the state of the flow and clears the previous messages."""
+        super().set_up_flow_state()
+        
+        preface = \
+            self.flow_config["template"]["preface"] + "\n\n" "#function used to evaluate the program:\n" + self.flow_config["evaluate_function"] + "\n\n"
+            
+        self.template: Program = Program(preface =preface,artifacts = [])
+        
+        # ~~~instantiate isladns~~~
+        self.flow_state["islands"] = [
+            Island(
+                artifact_to_evolve_name =self.flow_config["artifact_to_evolve_name"],
+                artifacts_per_prompt = self.flow_config["artifacts_per_prompt"],
+                temperature = self.flow_config["temperature"],
+                temperature_period = self.flow_config["temperature_period"],
+                reduce_score_method = np.mean,
+                sample_with_replacement = self.flow_config["sample_with_replacement"],
+                template=self.template
+                )
+        for _ in range(self.flow_config["num_islands"])
+        ]
+        
+        self.flow_state["last_reset_time"] = time.time()
+        self.flow_state["best_score_per_island"] = [float("-inf") for _ in range(self.flow_config["num_islands"])]
+        self.flow_state["best_program_per_island"] = [None for _ in range(self.flow_config["num_islands"])]
+        self.flow_state["best_scores_per_test_per_island"] = [None for _ in range(self.flow_config["num_islands"])]
+        self.flow_state["first_program_registered"] = False
+          
+    def get_prompt(self):
+        """ This method gets a prompt from an island. It returns the code, the version generated and the island id."""
+        island_id = np.random.choice(len(self.flow_state["islands"]))
+        code, version_generated = self.flow_state["islands"][island_id].get_prompt()
+        return code,version_generated,island_id
+    
+    def reset_islands(self):
+        """ This method resets the islands. It resets the worst islands and copies the best programs to the worst islands."""
+        # gaussian noise to break ties
+        sorted_island_ids = np.argsort(
+            np.array(self.flow_state["best_score_per_island"]) +
+            (np.random.randn(len(self.flow_state["best_score_per_island"])) * 1e-6)
+        )
+        
+        reset_island_ids = sorted_island_ids[:self.islands_to_reset]
+        keep_island_ids = sorted_island_ids[self.islands_to_reset:]
+       
+        for island_id in reset_island_ids:
+            self.flow_state["islands"][island_id] = Island(
+                               artifact_to_evolve_name =self.artifact_to_evolve_name,
+                               artifacts_per_prompt = self.artifacts_per_prompt,
+                               temperature = self.temperature,
+                               temperature_period = self.temperature_period,
+                               reduce_score_method = np.mean,
+                               sample_with_replacement = self.sample_with_replacement,
+                               template=self.template
+                               )
+            
+            self.flow_state["best_score_per_island"][island_id] = float("-inf")
+            founder_island_id = np.random.choice(keep_island_ids)
+            founder = self.flow_state["best_score_per_island"][founder_island_id]
+            founder_scores = self.flow_state["best_scores_per_test_per_island"][founder_island_id]
+            self._register_program_in_island(program=founder,island_id=island_id,scores_per_test=founder_scores)
+        
+    
+    
+    def register_program(self, program: AbstractArtifact ,island_id: int,scores_per_test: ScoresPerTest):
+        """ This method registers a program in an island. It also updates the best program if needed.
+        
+        :param program: The program to register
+        :type program: AbstractArtifact
+        :param island_id: The id of the island to register the program
+        :type island_id: int
+        :param scores_per_test: The scores per test of the program
+        :type scores_per_test: ScoresPerTest
+        """
+        if not program.calls_ancestor(artifact_to_evolve=self.artifact_to_evolve_name):
+            #program added at the beggining, so add to all islands
+            if island_id is None:
+                for id in range(self.num_islands):
+                    self._register_program_in_island(program=program,island_id=id,scores_per_test=scores_per_test)
+                    
+            else:
+                self._register_program_in_island(program=program,island_id=island_id,scores_per_test=scores_per_test)
+            
+        #reset islands if needed
+        if time.time() - self.flow_state["last_reset_time"]> self.reset_period:
+            self.reset_islands()
+            self.flow_state["last_reset_time"] = time.time()
+    
+    def _register_program_in_island(self,program: AbstractArtifact, scores_per_test: ScoresPerTest, island_id: Optional[int] = None):
+        """ This method registers a program in an island. It also updates the best program if needed.
+        
+        :param program: The program to register
+        :type program: AbstractArtifact
+        :param scores_per_test: The scores per test of the program
+        :type scores_per_test: ScoresPerTest
+        :param island_id: The id of the island to register the program
+        :type island_id: Optional[int]
+        """
+        self.flow_state["islands"][island_id].register_program(program,scores_per_test)
+        
+        scores_per_test_values = np.array([ score_per_test["score"] for score_per_test in scores_per_test.values()])
+        score = self.reduce_score_method(scores_per_test_values)
+        
+        if score > self.flow_state["best_score_per_island"][island_id]:
+            self.flow_state["best_score_per_island"][island_id] = score
+            self.flow_state["best_program_per_island"][island_id] = str(program)
+            self.flow_state["best_scores_per_test_per_island"][island_id] = scores_per_test
+    
+    def get_best_programs(self) -> List[Dict[str,Any]]:
+        """ This method returns the best programs per island."""
+        sorted_island_ids = np.argsort(np.array(self.flow_state["best_score_per_island"]))
+        return {
+        "best_island_programs": [
+            {
+                "rank": self.num_islands - rank,
+                "score": self.flow_state["best_score_per_island"][island_id],
+                "program": self.flow_state["best_program_per_island"][island_id],
+                "island_id": int(island_id),
+                }
+            for rank,island_id in enumerate(sorted_island_ids)
+            ]
+        }
+    
+
+    def run(self,input_message: FlowMessage):
+        """ This method runs the flow. It performs the operation requested in the input message."""
+        input_data = input_message.data
+        operation = input_data["operation"]
+        content = input_data["content"]
+        
+        possible_operations = [
+            "register_program",
+            "get_prompt",
+            "get_best_programs_per_island",
+        ]
+        
+        if operation not in possible_operations:
+            raise ValueError(f"operation must be one of the following: {possible_operations}")
+        
+        response = {}
+        if operation == "get_prompt":
+            response["retrieved"] = False
+            
+            if not self.flow_state["first_program_registered"]:
+                response["retrieved"] = False
+            
+            else:
+                code,version_generated,island_id = self.get_prompt()
+                
+                response["retrieved"] = {
+                    "code":code,
+                    "version_generated": version_generated,
+                    "island_id":island_id,
+                    "header": self.evaluate_file_full_content
+                }
+            
+        elif operation == "register_program":
+            try:
+                
+                artifact = text_to_artifact(content["artifact"])
+                island_id = content.get("island_id",None)
+                scores_per_test = content["scores_per_test"]
+                if scores_per_test is not None:
+                    self.register_program(program=artifact,island_id=island_id,scores_per_test=scores_per_test)
+                response["retrieved"] = "Program registered"
+                self.flow_state["first_program_registered"] = True
+            except:
+                response["retrieved"] = "Program failed to register"
+        else:
+            response["retrieved"] = self.get_best_programs()
+        
+        response["from"] = "ProgramDBFlow"
+        reply = self.package_output_message(
+            input_message,
+            response
+        )
+        
+        self.send_message(reply)

ProgramDBFlowModule/ProgramDBFlow.yaml

@@ -0,0 +1,21 @@
+#default configuration of your flow (if needed) goes here
+_target_: flow_modules.aiflows.FunSearchFlowModule.ProgramDBFlow.instantiate_from_default_config
+input_interface:
+  - "operation"
+output_interface:
+  - "retrieved"
+
+name: ProgramDBFlow
+description: A flow that saves programs in a database of islands
+artifact_to_evolve_name: solve_function 
+evaluate_function: ???
+evaluate_file_full_content: ???
+num_islands: 3
+reset_period : 14400 #4 hours in seconds
+artifacts_per_prompt: 2
+temperature: 0.1
+temperature_period: 30000
+sample_with_replacement:  False
+portion_of_islands_to_reset:  0.5
+template:
+  preface: ""

ProgramDBFlowModule/__init__.py

@@ -0,0 +1,13 @@
+# ~~~ Specify the dependencies ~~~
+# e.g.,
+# dependencies = [
+#     {"url": "aiflows/AutoGPTFlowModule", "revision": "main"},
+# ]
+# Revision can correspond toa branch, commit hash or a absolute path to a local directory (ideal for development)
+# from aiflows import flow_verse
+
+# flow_verse.sync_dependencies(dependencies)
+
+# ~~~ Import of your flow class (if you have any) ~~~
+# from .NAMEOFYOURFLOW import NAMEOFYOURFLOWCLASS
+from .ProgramDBFlow import ProgramDBFlow

ProgramDBFlowModule/artifacts/__init__.py

@@ -0,0 +1,2 @@
+from .abstract import AbstractArtifact
+from .function import FunctionArtifact

ProgramDBFlowModule/artifacts/abstract.py

@@ -0,0 +1,49 @@
+"""
+This code is an implementation of Funsearch (https://www.nature.com/articles/s41586-023-06924-6) and is heavily inspired by the original code (https://github.com/google-deepmind/funsearch)
+
+    **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}
+    }
+"""
+from typing import Any, Optional
+import dataclasses
+
+
+@dataclasses.dataclass
+class AbstractArtifact:
+    """ Abstract class for artifacts."""
+    
+    name: str
+    args: str
+    body: str
+    return_type: Optional[str] = None
+    docstring: Optional[str] = None
+
+    def __str__(self)->str:
+        raise NotImplementedError()
+    
+    def __setattr__(self, name: str, value: str) -> None:
+        # Ensure there aren't leading & trailing new lines in `body`.
+        if name == 'body':
+            value = value.strip('\n')
+        # Ensure there aren't leading & trailing quotes in `docstring``.
+        if name == 'docstring' and value is not None:
+            if '"""' in value:
+                value = value.strip()
+                value = value.replace('"""', '')
+        super().__setattr__(name, value)
+    
+    def rename_artifact_calls(self, source_name, target_name) -> str:
+        raise NotImplementedError
+    
+    def text_to_artifact(self):
+        raise NotImplementedError
+    
+    def calls_ancestor(self,artifact_to_evolve: str) -> bool:
+        raise NotImplementedError

ProgramDBFlowModule/artifacts/function.py

@@ -0,0 +1,116 @@
+"""
+This code is an implementation of Funsearch (https://www.nature.com/articles/s41586-023-06924-6) and is heavily inspired by the original code (https://github.com/google-deepmind/funsearch)
+
+    **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}
+    }
+"""
+
+from . import AbstractArtifact
+import dataclasses
+import tokenize
+import io
+from collections.abc import Iterator, MutableSet, Sequence
+
+@dataclasses.dataclass
+class FunctionArtifact(AbstractArtifact):
+    def __str__(self) -> str:
+        
+        return_type = f' -> {self.return_type}' if self.return_type else ''
+        function = f'def {self.name}({self.args}){return_type}:\n'
+        
+        if self.docstring:
+        # self.docstring is already indented on every line except the first one.
+        # Here, we assume the indentation is always two spaces.
+            new_line = '\n' if self.body else ''
+            function += f'    """{self.docstring}"""{new_line}'
+            
+        # self.body is already indented.
+        function += self.body + '\n\n'
+        return function
+    
+    @staticmethod    
+    def _tokenize(code: str) -> Iterator[tokenize.TokenInfo]:
+        """Transforms `code` into Python tokens."""
+        code_bytes = code.encode()
+        code_io = io.BytesIO(code_bytes)
+        return tokenize.tokenize(code_io.readline)
+
+    @staticmethod
+    def _untokenize(tokens: Sequence[tokenize.TokenInfo]) -> str:
+        """Transforms a list of Python tokens into code."""
+        code_bytes = tokenize.untokenize(tokens)
+        return code_bytes.decode()
+    
+    def _get_artifacts_called(self) -> MutableSet[str]:
+        """Returns the set of all functions called in function."""
+        code = str(self.body)
+        return set(token.string for token, is_call in
+                    self._yield_token_and_is_call(code) if is_call)
+    
+    def calls_ancestor(self,artifact_to_evolve: str) -> bool:
+        """Returns whether the generated function is calling an earlier version."""
+        
+        for name in self._get_artifacts_called():
+            # In `program` passed into this function the most recently generated
+            # function has already been renamed to `function_to_evolve` (wihout the
+            # suffix). Therefore any function call starting with `function_to_evolve_v`
+            # is a call to an ancestor function.
+            if name.startswith(f'{artifact_to_evolve}_v') and not name.startswith(self.name):
+                return True
+        return False
+
+    
+    def _yield_token_and_is_call(cls,code: str) -> Iterator[tuple[tokenize.TokenInfo, bool]]:
+        """Yields each token with a bool indicating whether it is a function call."""
+        
+        tokens = cls._tokenize(code)
+        prev_token = None
+        is_attribute_access = False
+        for token in tokens:
+            if (prev_token and  # If the previous token exists and
+                prev_token.type == tokenize.NAME and  # it is a Python identifier
+                token.type == tokenize.OP and  # and the current token is a delimiter
+                token.string == "("
+                ):  # and in particular it is '('.
+                yield prev_token, not is_attribute_access
+                is_attribute_access = False
+            else:
+                if prev_token:
+                    is_attribute_access = (
+                        prev_token.type == tokenize.OP and prev_token.string == '.'
+                    )
+                    yield prev_token, False
+                    
+            prev_token = token
+        if prev_token:
+            yield prev_token, False
+
+    
+    def rename_artifact_calls(self, source_name, target_name) -> str:
+        implementation = str(self)
+        
+        if source_name not in implementation:
+            return implementation
+        
+        modified_tokens = []
+        for token, is_call in self._yield_token_and_is_call(implementation):
+            if is_call and token.string == source_name:
+                # Replace the function name token
+                modified_token = tokenize.TokenInfo(
+                    type=token.type,
+                    string=target_name,
+                    start=token.start,
+                    end=token.end,
+                    line=token.line,
+                )
+                modified_tokens.append(modified_token)
+            else:
+                modified_tokens.append(token)
+        return self._untokenize(modified_tokens)

ProgramDBFlowModule/demo.yaml

@@ -0,0 +1,24 @@
+#default configuration of your flow (if needed) goes here
+_target_: flow_modules.aiflows.FunSearchFlowModule.ProgramDBFlow.instantiate_from_default_config
+input_interface:
+  - "operation"
+output_interface:
+  - "retrieved"
+
+name: ProgramDBFlow
+description: A flow that saves programs in a database of islands
+artifact_to_evolve_name: solve_function 
+evaluate_function: |2-
+  def evaluate(solve_function: str, tests_inputs: List[str], expected_outputs: str) -> float:\n    """Returns the score of the solve function we\'re evolving based on the tests_inputs and expected_outputs.\n    Scores are between 0 and 1, unless the program fails to run, in which case the score is -1.\n    """\n    if solve(solve_function, tests_inputs, expected_outputs) == True:\n        return 1.0\n    return 0.0
+
+evaluate_file_full_content: |2-
+  """Problem Description:\nServal has a string s that only consists of 0 and 1 of length n. The i-th character of s is denoted as s_i, where 1\\leq i\\leq n.\nServal can perform the following operation called Inversion Magic on the string s:\nChoose an segment [l, r] (1\\leq l\\leq r\\leq n). For l\\leq i\\leq r, change s_i into 1 if s_i is 0, and change s_i into 0 if s_i is 1.\nFor example, let s be 010100 and the segment [2,5] is chosen. The string s will be 001010 after performing the Inversion Magic.\nServal wants to make s a palindrome after performing Inversion Magic exactly once. Help him to determine whether it is possible.\nA string is a palindrome iff it reads the same backwards as forwards. For example, 010010 is a palindrome but 10111 is not.\n\nInput Description:\nInput\nEach test contains multiple test cases. The first line contains the number of test cases t (1\\leq t\\leq 10^4). The description of the test cases follows.\nThe first line of each test case contains a single integer n (2\\leq n\\leq 10^5) — the length of string s.\nThe second line of each test case contains a binary string s of length n. Only characters 0 and 1 can appear in s.\nIt\'s guaranteed that the sum of n over all test cases does not exceed 2\\cdot 10^5.\n\nOutput Description:\nOutput\nFor each test case, print Yes if s can be a palindrome after performing Inversion Magic exactly once, and print No if not.\nYou can output Yes and No in any case (for example, strings yEs, yes, Yes and YES will be recognized as a positive response).\n\nPublic Tests:\nTest 1:\n  Input: [\'1\', \'4\', \'1001\']\n  Output: \'YES\'\nTest 2:\n  Input: [\'1\', \'5\', \'10010\']\n  Output: \'YES\'\nTest 3:\n  Input: [\'1\', \'7\', \'0111011\']\n  Output: \'NO\'\n\n"""\n\n\nimport ast\nimport itertools\nimport numpy as np\nfrom typing import List\n\ndef solve(solve_function: str,input: List[str], expected_output: str) -> str:\n    """function used to run the solve function on input *kwargs and return the the predicted output\n    \n    :param solve_function: the function to run (the solve function below as a string)\n    :type solve_function: str\n    :param kwargs: the inputs to the solve function\n    :type kwargs: List[str]\n    """\n    local_namespace = {}\n    exec(solve_function,local_namespace)\n    found_name, program_name = get_function_name_from_code(solve_function)\n    \n    if not found_name:\n        raise ValueError(f"Function name not found in program: {solve_function}")\n    \n    solve_fn = local_namespace.get(program_name)\n    \n    prediction = solve_fn(input)\n    \n    prediction = prediction.split()\n    expected_output = expected_output.split()\n    \n    if len(prediction) != len(expected_output):\n        raise ValueError(f"Invalid Format of prediction")\n    \n    for i in range(len(prediction)):\n        if prediction[i] != expected_output[i]:\n            return False\n    \n    return True\n\ndef evaluate(solve_function: str, tests_inputs: List[str], expected_outputs: str) -> float:\n    """Returns the score of the solve function we\'re evolving based on the tests_inputs and expected_outputs.\n    Scores are between 0 and 1, unless the program fails to run, in which case the score is -1.\n    """\n    if solve(solve_function,tests_inputs,expected_outputs) == True:\n        return 1.0\n    return 0.0\n\n\ndef get_function_name_from_code(code):\n    tree = ast.parse(code)\n    for node in ast.walk(tree):\n        if isinstance(node, ast.FunctionDef):\n            return True, node.name\n\n    # something is wrong\n    return False, None\n\n\n\n\n    \n    
+num_islands: 3
+reset_period : 14400 #4 hours in seconds
+artifacts_per_prompt: 2
+temperature: 0.1
+temperature_period: 30000
+sample_with_replacement:  False
+portion_of_islands_to_reset:  0.5
+template:
+  preface: ""

ProgramDBFlowModule/pip_requirements.txt

@@ -0,0 +1,2 @@
+numpy==1.26.2
+SciPy==1.11.4

ProgramDBFlowModule/run.py

@@ -0,0 +1,88 @@
+import os
+
+import hydra
+
+import aiflows
+from aiflows.backends.api_info import ApiInfo
+from aiflows.utils.general_helpers import read_yaml_file, quick_load_api_keys
+
+from aiflows import logging
+from aiflows.flow_cache import CACHING_PARAMETERS, clear_cache
+
+from aiflows.utils import serving
+from aiflows.workers import run_dispatch_worker_thread
+from aiflows.messages import FlowMessage
+from aiflows.interfaces import KeyInterface
+from aiflows.utils.colink_utils import start_colink_server
+from aiflows import flow_verse
+
+
+dependencies = [
+    {
+        "url": "aiflows/FunSearchFlowModule",
+        "revision": os.path.abspath("../")
+    }
+]
+flow_verse.sync_dependencies(dependencies)
+
+logging.set_verbosity_debug()
+
+
+if __name__ == "__main__":
+    
+    cl = start_colink_server()
+    
+    serving.recursive_serve_flow(
+        cl=cl,
+        flow_class_name="flow_modules.aiflows.FunSearchFlowModule.ProgramDBFlow",
+        flow_endpoint="ProgramDBFlow",
+    )
+        
+    run_dispatch_worker_thread(cl)
+    
+    config_overrides = read_yaml_file(os.path.join(".", "demo.yaml"))
+
+    funsearch_proxy = serving.get_flow_instance(
+        cl=cl,
+        flow_endpoint="ProgramDBFlow",
+        config_overrides=config_overrides,
+    )
+    
+    data = {
+        'operation': 'register_program',
+        'content': {
+            'artifact': 'def solve_function(input) -> str:\n    """Attempt at solving the problem given the input input and returns the predicted output (see the top of the file for problem description)"""\n    return \'YES\'\n',
+            'scores_per_test': 
+                {
+                    'test_1': 
+                        {'score': 1.0, 'feedback': 'No feedback available.'},
+                    'test_2': 
+                        {'score': 1.0, 'feedback': 'No feedback available.'},
+                    'test_3': {'score': 0.0, 'feedback': 'No feedback available.'},
+                    'test_4': {'score': -1, 'feedback': 'Invalid Format of prediction'}
+                }
+        }
+    }
+    
+    input_message = funsearch_proxy.package_input_message(data = data)
+    funsearch_proxy.send_message(input_message)
+    
+    data = {'operation': 'get_prompt', 'content': {}}
+    input_message = funsearch_proxy.package_input_message(data = data)
+    
+    example_of_prompt = funsearch_proxy.get_reply_future(input_message).get_data()
+    
+    data = {
+        "operation": "get_best_programs_per_island",
+        "content": {}
+    }
+   
+    input_message = funsearch_proxy.package_input_message(data = data)
+    
+    best_pg_per_island = funsearch_proxy.get_reply_future(input_message).get_data()
+    
+    print("~~~Example of Prompt~~~")
+    print(example_of_prompt)
+    
+    print("~~~Best Programs Per Island~~~")
+    print(best_pg_per_island)

ProgramDBFlowModule/utils.py

@@ -0,0 +1,55 @@
+""" This code is an implementation of Funsearch (https://www.nature.com/articles/s41586-023-06924-6) and is heavily inspired by the original code (https://github.com/google-deepmind/funsearch)
+
+    **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}
+    }
+"""
+
+import numpy as np
+import scipy
+def add_gaussian_noise(array,mean=0,std=1):
+    return array + np.random.normal(mean,std,array.shape)
+
+def get_versioned_function_name(function_name,i):
+    pass
+
+def rename_function(implementation,function_name):
+    pass
+
+def string_concatenation(strings,split_char='\n\n'):
+    pass
+
+def make_header_like(implementation,function_name):
+    pass
+
+def rename_function_calls(implementation,function_name):
+    pass
+
+def extract_template_from_program(program):
+    pass
+
+
+def _softmax(logits: np.ndarray, temperature: float, epsilon = 1e-6) -> np.ndarray:
+    """Returns the tempered softmax of 1D finite `logits`."""
+    if not np.all(np.isfinite(logits)):
+        non_finites = set(logits[~np.isfinite(logits)])
+        raise ValueError(f'`logits` contains non-finite value(s): {non_finites}')
+    if not np.issubdtype(logits.dtype, np.floating):
+        logits = np.array(logits, dtype=np.float32)
+
+    result = scipy.special.softmax(logits / temperature, axis=-1)
+    
+    #Non zero mass to prevent zero probabilities
+    result += epsilon  # Add epsilon to prevent zeros
+    result /= np.sum(result, axis=-1, keepdims=True)  # Normalize
+
+    # Ensure that probabilities sum to 1 to prevent error in `np.random.choice`.
+    index = np.argmax(result)
+    result[index] = 1 - np.sum(result[0:index]) - np.sum(result[index+1:])
+    return result

README.md

@@ -1,3 +1,675 @@
 ---
 license: mit
 ---
+# Table of Contents
+
+* [FunSearch](#FunSearch)
+  * [FunSearch](#FunSearch.FunSearch)
+    * [make\_request\_for\_prompt](#FunSearch.FunSearch.make_request_for_prompt)
+    * [request\_samplers](#FunSearch.FunSearch.request_samplers)
+    * [get\_next\_state](#FunSearch.FunSearch.get_next_state)
+    * [set\_up\_flow\_state](#FunSearch.FunSearch.set_up_flow_state)
+    * [save\_message\_to\_state](#FunSearch.FunSearch.save_message_to_state)
+    * [rename\_key\_message\_in\_state](#FunSearch.FunSearch.rename_key_message_in_state)
+    * [message\_in\_state](#FunSearch.FunSearch.message_in_state)
+    * [get\_message\_from\_state](#FunSearch.FunSearch.get_message_from_state)
+    * [pop\_message\_from\_state](#FunSearch.FunSearch.pop_message_from_state)
+    * [merge\_message\_request\_state](#FunSearch.FunSearch.merge_message_request_state)
+    * [register\_data\_to\_state](#FunSearch.FunSearch.register_data_to_state)
+    * [call\_program\_db](#FunSearch.FunSearch.call_program_db)
+    * [call\_evaluator](#FunSearch.FunSearch.call_evaluator)
+    * [call\_sampler](#FunSearch.FunSearch.call_sampler)
+    * [generate\_reply](#FunSearch.FunSearch.generate_reply)
+    * [run](#FunSearch.FunSearch.run)
+* [ProgramDBFlowModule](#ProgramDBFlowModule)
+* [ProgramDBFlowModule.ProgramDBFlow](#ProgramDBFlowModule.ProgramDBFlow)
+  * [ProgramDBFlow](#ProgramDBFlowModule.ProgramDBFlow.ProgramDBFlow)
+    * [set\_up\_flow\_state](#ProgramDBFlowModule.ProgramDBFlow.ProgramDBFlow.set_up_flow_state)
+    * [get\_prompt](#ProgramDBFlowModule.ProgramDBFlow.ProgramDBFlow.get_prompt)
+    * [reset\_islands](#ProgramDBFlowModule.ProgramDBFlow.ProgramDBFlow.reset_islands)
+    * [register\_program](#ProgramDBFlowModule.ProgramDBFlow.ProgramDBFlow.register_program)
+    * [get\_best\_programs](#ProgramDBFlowModule.ProgramDBFlow.ProgramDBFlow.get_best_programs)
+    * [run](#ProgramDBFlowModule.ProgramDBFlow.ProgramDBFlow.run)
+* [SamplerFlowModule](#SamplerFlowModule)
+* [SamplerFlowModule.SamplerFlow](#SamplerFlowModule.SamplerFlow)
+  * [SamplerFlow](#SamplerFlowModule.SamplerFlow.SamplerFlow)
+    * [run](#SamplerFlowModule.SamplerFlow.SamplerFlow.run)
+* [EvaluatorFlowModule](#EvaluatorFlowModule)
+* [EvaluatorFlowModule.EvaluatorFlow](#EvaluatorFlowModule.EvaluatorFlow)
+  * [EvaluatorFlow](#EvaluatorFlowModule.EvaluatorFlow.EvaluatorFlow)
+    * [load\_functions](#EvaluatorFlowModule.EvaluatorFlow.EvaluatorFlow.load_functions)
+    * [run\_function\_with\_timeout](#EvaluatorFlowModule.EvaluatorFlow.EvaluatorFlow.run_function_with_timeout)
+    * [evaluate\_program](#EvaluatorFlowModule.EvaluatorFlow.EvaluatorFlow.evaluate_program)
+    * [analyse](#EvaluatorFlowModule.EvaluatorFlow.EvaluatorFlow.analyse)
+    * [run](#EvaluatorFlowModule.EvaluatorFlow.EvaluatorFlow.run)
+
+<a id="FunSearch"></a>
+
+# FunSearch
+
+<a id="FunSearch.FunSearch"></a>
+
+## FunSearch Objects
+
+```python
+class FunSearch(CompositeFlow)
+```
+
+This class implements FunSearch. This code is an implementation of Funsearch (https://www.nature.com/articles/s41586-023-06924-6) and is heavily inspired by the original code (https://github.com/google-deepmind/funsearch) . It's a Flow in charge of starting, stopping and managing (passing around messages) the FunSearch process. It passes messages around to the following subflows:
+
+- ProgramDBFlow: which is in charge of storing and retrieving programs.
+- SamplerFlow: which is in charge of sampling programs.
+- EvaluatorFlow: which is in charge of evaluating programs.
+
+*Configuration Parameters*:
+
+- `name` (str): The name of the flow. Default: "FunSearchFlow".
+- `description` (str): The description of the flow. Default: "A flow implementing FunSearch"
+- `subflows_config` (Dict[str,Any]): A dictionary of subflows configurations. Default:
+    - `ProgramDBFlow`: By default, it uses the `ProgramDBFlow` class and uses its default parameters.
+    - `SamplerFlow`: By default, it uses the `SamplerFlow` class and uses its default parameters.
+    - `EvaluatorFlow`: By default, it uses the `EvaluatorFlow` class and uses its default parameters.
+
+**Input Interface**:
+
+- `from` (str): The flow from which the message is coming from. It can be one of the following: "FunSearch", "SamplerFlow", "EvaluatorFlow", "ProgramDBFlow".
+- `operation` (str): The operation to perform. It can be one of the following: "start", "stop", "get_prompt", "get_best_programs_per_island", "register_program".
+- `content` (Dict[str,Any]): The content associated to an operation. Here is the expected content for each operation:
+    - "start":
+        - `num_samplers` (int): The number of samplers to start up. Note that it's still restricted by the number of workers available. Default: 1.
+    - "stop":
+        - No content. Pass either an empty dictionary or None. Works also with no content.
+    - "get_prompt":
+        - No content. Pass either an empty dictionary or None. Works also with no content.
+    - "get_best_programs_per_island":
+        - No content. Pass either an empty dictionary or None. Works also with no content.
+
+**Output Interface**:
+
+- `retrieved` (Dict[str,Any]): The retrieved data.
+
+**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}
+}
+
+<a id="FunSearch.FunSearch.make_request_for_prompt"></a>
+
+#### make\_request\_for\_prompt
+
+```python
+def make_request_for_prompt()
+```
+
+This method makes a request for a prompt. It sends a message to itself with the operation "get_prompt" which will trigger the flow to call the `ProgramDBFlow` to get a prompt.
+
+<a id="FunSearch.FunSearch.request_samplers"></a>
+
+#### request\_samplers
+
+```python
+def request_samplers(input_message: FlowMessage)
+```
+
+This method requests samplers. It sends a message to itself with the operation "get_prompt" which will trigger the flow to call the `ProgramDBFlow` to get a prompt.
+
+**Arguments**:
+
+- `input_message` (`FlowMessage`): The input message that triggered the request for samplers.
+
+<a id="FunSearch.FunSearch.get_next_state"></a>
+
+#### get\_next\_state
+
+```python
+def get_next_state(input_message: FlowMessage)
+```
+
+This method determines the next state of the flow based on the input message. It will return the next state based on the current state and the message received.
+
+**Arguments**:
+
+- `input_message` (`FlowMessage`): The input message that triggered the request for the next state.
+
+**Returns**:
+
+`str`: The next state of the flow.
+
+<a id="FunSearch.FunSearch.set_up_flow_state"></a>
+
+#### set\_up\_flow\_state
+
+```python
+def set_up_flow_state()
+```
+
+This method sets up the state of the flow. It's called at the beginning of the flow.
+
+<a id="FunSearch.FunSearch.save_message_to_state"></a>
+
+#### save\_message\_to\_state
+
+```python
+def save_message_to_state(msg_id: str, message: FlowMessage)
+```
+
+This method saves a message to the state of the flow. It's used to keep track of state on a per message basis (i.e., state of the flow depending on the message received and id).
+
+**Arguments**:
+
+- `msg_id` (`str`): The id of the message to save.
+- `message` (`FlowMessage`): The message to save.
+
+<a id="FunSearch.FunSearch.rename_key_message_in_state"></a>
+
+#### rename\_key\_message\_in\_state
+
+```python
+def rename_key_message_in_state(old_key: str, new_key: str)
+```
+
+This method renames a key in the state of the flow in the "msg_requests" dictonary. It's used to rename a key in the state of the flow (i.e., rename a message id).
+
+**Arguments**:
+
+- `old_key` (`str`): The old key to rename.
+- `new_key` (`str`): The new key to rename to.
+
+<a id="FunSearch.FunSearch.message_in_state"></a>
+
+#### message\_in\_state
+
+```python
+def message_in_state(msg_id: str) -> bool
+```
+
+This method checks if a message is in the state of the flow (in "msg_requests" dictionary). It returns True if the message is in the state, otherwise it returns False.
+
+**Arguments**:
+
+- `msg_id` (`str`): The id of the message to check if it's in the state.
+
+**Returns**:
+
+`bool`: True if the message is in the state, otherwise False.
+
+<a id="FunSearch.FunSearch.get_message_from_state"></a>
+
+#### get\_message\_from\_state
+
+```python
+def get_message_from_state(msg_id: str) -> Dict[str, Any]
+```
+
+This method returns the state associated with a message id in the state of the flow (in "msg_requests" dictionary).
+
+**Arguments**:
+
+- `msg_id` (`str`): The id of the message to get the state from.
+
+**Returns**:
+
+`Dict[str,Any]`: The state associated with the message id.
+
+<a id="FunSearch.FunSearch.pop_message_from_state"></a>
+
+#### pop\_message\_from\_state
+
+```python
+def pop_message_from_state(msg_id: str) -> Dict[str, Any]
+```
+
+This method pops a message from the state of the flow (in "msg_requests" dictionary). It the state associate to a message and removes it from the state.
+
+**Arguments**:
+
+- `msg_id` (`str`): The id of the message to pop from the state.
+
+**Returns**:
+
+`Dict[str,Any]`: The state associated with the message id.
+
+<a id="FunSearch.FunSearch.merge_message_request_state"></a>
+
+#### merge\_message\_request\_state
+
+```python
+def merge_message_request_state(id: str, new_states: Dict[str, Any])
+```
+
+This method merges new states to a message in the state of the flow (in "msg_requests" dictionary). It merges new states to a message in the state.
+
+**Arguments**:
+
+- `id` (`str`): The id of the message to merge new states to.
+- `new_states` (`Dict[str,Any]`): The new states to merge to the message.
+
+<a id="FunSearch.FunSearch.register_data_to_state"></a>
+
+#### register\_data\_to\_state
+
+```python
+def register_data_to_state(input_message: FlowMessage)
+```
+
+This method registers the input message data to the flow state. It's called everytime a new input message is received.
+
+**Arguments**:
+
+- `input_message` (`FlowMessage`): The input message
+
+<a id="FunSearch.FunSearch.call_program_db"></a>
+
+#### call\_program\_db
+
+```python
+def call_program_db(input_message)
+```
+
+This method calls the ProgramDBFlow. It sends a message to the ProgramDBFlow with the data of the input message.
+
+**Arguments**:
+
+- `input_message` (`FlowMessage`): The input message to send to the ProgramDBFlow.
+
+<a id="FunSearch.FunSearch.call_evaluator"></a>
+
+#### call\_evaluator
+
+```python
+def call_evaluator(input_message)
+```
+
+This method calls the EvaluatorFlow. It sends a message to the EvaluatorFlow with the data of the input message.
+
+**Arguments**:
+
+- `input_message` (`FlowMessage`): The input message to send to the EvaluatorFlow.
+
+<a id="FunSearch.FunSearch.call_sampler"></a>
+
+#### call\_sampler
+
+```python
+def call_sampler(input_message)
+```
+
+This method calls the SamplerFlow. It sends a message to the SamplerFlow with the data of the input message.
+
+**Arguments**:
+
+- `input_message` (`FlowMessage`): The input message to send to the SamplerFlow.
+
+<a id="FunSearch.FunSearch.generate_reply"></a>
+
+#### generate\_reply
+
+```python
+def generate_reply(input_message: FlowMessage)
+```
+
+This method generates a reply to a message sent to user. It packages the output message and sends it.
+
+**Arguments**:
+
+- `input_message` (`FlowMessage`): The input message to generate a reply to.
+
+<a id="FunSearch.FunSearch.run"></a>
+
+#### run
+
+```python
+def run(input_message: FlowMessage)
+```
+
+This method runs the flow. It's the main method of the flow. It's called when the flow is executed.
+
+<a id="ProgramDBFlowModule"></a>
+
+# ProgramDBFlowModule
+
+<a id="ProgramDBFlowModule.ProgramDBFlow.ProgramDBFlow"></a>
+
+## ProgramDBFlow Objects
+
+```python
+class ProgramDBFlow(AtomicFlow)
+```
+
+This class implements a ProgramDBFlow. It's a flow that stores programs and their scores in a database. It can also query the database for the best programs or generate a prompt containing stored programs in order to evolve them with a SamplerFlow. This code is an implementation of Funsearch (https://www.nature.com/articles/s41586-023-06924-6) and is heavily inspired by the original code (https://github.com/google-deepmind/funsearch)
+
+**Configuration Parameters**:
+
+- `name` (str): The name of the flow. Default: "ProgramDBFlow"
+- `description` (str): A description of the flow. This description is used to generate the help message of the flow. Default: " A flow that saves programs in a database of islands"
+- `artifact_to_evolve_name` (str): The name of the artifact/program to evolve. Default: "solve_function"
+- `evaluate_function` (str): The function used to evaluate the program. No Default value. This MUST be passed as a parameter.
+- `evaluate_file_full_content` (str): The full content of the file containing the evaluation function. No Default value. This MUST be passed as a parameter.
+- `num_islands`: The number of islands to use. Default: 3
+- `reset_period`: The period in seconds to reset the islands. Default: 3600
+- `artifacts_per_prompt`: The number of previous artifacts/programs to include in a prompt. Default: 2
+- `temperature`: The temperature of the island. Default: 0.1
+- `temperature_period`: The period in seconds to change the temperature. Default: 30000
+- `sample_with_replacement`: Whether to sample with replacement. Default: False
+- `portion_of_islands_to_reset`: The portion of islands to reset. Default: 0.5
+- `template` (dict): The template to use for a program. Default: {"preface": ""}
+
+**Input Interface**:
+
+- `operation` (str): The operation to perform. It can be one of the following: ["register_program","get_prompt","get_best_programs_per_island"]
+
+**Output Interface**:
+
+- `retrieved` (Any): The retrieved data. It can be one of the following:
+    - If the operation is "get_prompt", it can be a dictionary with the following keys
+        - `code` (str): The code of the prompt
+        - `version_generated` (int): The version of the prompt generated
+        - `island_id` (int): The id of the island that generated the prompt
+        - `header` (str): The header of the prompt
+    - If the operation is "register_program", it can be a string with the message "Program registered" or "Program failed to register"
+    - If the operation is "get_best_programs_per_island", it can be a dictionary with the following keys:
+        - `best_island_programs` (List[Dict[str,Any]]): A list of dictionaries with the following keys:
+            - `rank` (int): The rank of the program (1 is the best)
+            - `score` (float): The score of the program
+            - `program` (str): The program
+            - `island_id` (int): The id of the island that generated the program
+
+**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}
+}
+
+<a id="ProgramDBFlowModule.ProgramDBFlow.ProgramDBFlow.set_up_flow_state"></a>
+
+#### set\_up\_flow\_state
+
+```python
+def set_up_flow_state()
+```
+
+This method sets up the state of the flow and clears the previous messages.
+
+<a id="ProgramDBFlowModule.ProgramDBFlow.ProgramDBFlow.get_prompt"></a>
+
+#### get\_prompt
+
+```python
+def get_prompt()
+```
+
+This method gets a prompt from an island. It returns the code, the version generated and the island id.
+
+<a id="ProgramDBFlowModule.ProgramDBFlow.ProgramDBFlow.reset_islands"></a>
+
+#### reset\_islands
+
+```python
+def reset_islands()
+```
+
+This method resets the islands. It resets the worst islands and copies the best programs to the worst islands.
+
+<a id="ProgramDBFlowModule.ProgramDBFlow.ProgramDBFlow.register_program"></a>
+
+#### register\_program
+
+```python
+def register_program(program: AbstractArtifact, island_id: int,
+                     scores_per_test: ScoresPerTest)
+```
+
+This method registers a program in an island. It also updates the best program if needed.
+
+**Arguments**:
+
+- `program` (`AbstractArtifact`): The program to register
+- `island_id` (`int`): The id of the island to register the program
+- `scores_per_test` (`ScoresPerTest`): The scores per test of the program
+
+<a id="ProgramDBFlowModule.ProgramDBFlow.ProgramDBFlow.get_best_programs"></a>
+
+#### get\_best\_programs
+
+```python
+def get_best_programs() -> List[Dict[str, Any]]
+```
+
+This method returns the best programs per island.
+
+<a id="ProgramDBFlowModule.ProgramDBFlow.ProgramDBFlow.run"></a>
+
+#### run
+
+```python
+def run(input_message: FlowMessage)
+```
+
+This method runs the flow. It performs the operation requested in the input message.
+
+<a id="SamplerFlowModule"></a>
+
+# SamplerFlowModule
+
+<a id="SamplerFlowModule.SamplerFlow"></a>
+
+# SamplerFlowModule.SamplerFlow
+
+<a id="SamplerFlowModule.SamplerFlow.SamplerFlow"></a>
+
+## SamplerFlow Objects
+
+```python
+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}
+}
+
+<a id="SamplerFlowModule.SamplerFlow.SamplerFlow.run"></a>
+
+#### run
+
+```python
+def run(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.
+
+**Returns**:
+
+`Any`: The output of the backend.
+
+<a id="EvaluatorFlowModule"></a>
+
+# EvaluatorFlowModule
+
+<a id="EvaluatorFlowModule.EvaluatorFlow"></a>
+
+# EvaluatorFlowModule.EvaluatorFlow
+
+
+<a id="EvaluatorFlowModule.EvaluatorFlow.EvaluatorFlow"></a>
+
+## EvaluatorFlow Objects
+
+```python
+class EvaluatorFlow(AtomicFlow)
+```
+
+This class implements an EvaluatorFlow. It is a flow that evaluates a program (python code) using a given evaluator function. This code is an implementation of Funsearch (https://www.nature.com/articles/s41586-023-06924-6) and is heavily inspired by the original code (https://github.com/google-deepmind/funsearch)
+
+**Configuration Parameters**:
+
+- `name` (str): The name of the flow. Default: "EvaluatorFlow"
+- `description` (str): A description of the flow. This description is used to generate the help message of the flow. Default: "A flow that evaluates code on tests"
+- `py_file` (str): The python code containing the evaluation function. No default value. This MUST be passed as a parameter.
+- `function_to_run_name` (str): The name of the function to run (the evaluation function) in the evaluator file.  No default value. This MUST be passed as a parameter.
+- `test_inputs` (Dict[str,Any]): A dictionary of test inputs to evaluate the program. Default: {"test1": None, "test2": None}
+- `timeout_seconds` (int): The maximum number of seconds to run the evaluation function before returning an error. Default: 10
+- `run_error_score` (int): The score to return if the evaluation function fails to run. Default: -100
+- `use_test_input_as_key` (bool): Whether to use the test input parameters as the key in the output dictionary. Default: False
+
+**Input Interface**:
+
+- `artifact` (str): The program/artifact to evaluate.
+
+**Output Interface**:
+
+- `scores_per_test` (Dict[str, Dict[str, Any]]): A dictionary of scores per test input.
+
+**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}
+}
+
+<a id="EvaluatorFlowModule.EvaluatorFlow.EvaluatorFlow.load_functions"></a>
+
+#### load\_functions
+
+```python
+def load_functions()
+```
+
+Load the functions from the evaluator py file with ast parsing
+
+<a id="EvaluatorFlowModule.EvaluatorFlow.EvaluatorFlow.run_function_with_timeout"></a>
+
+#### run\_function\_with\_timeout
+
+```python
+def run_function_with_timeout(program: str, **kwargs)
+```
+
+Run the evaluation function with a timeout
+
+**Arguments**:
+
+- `program` (`str`): The program to evaluate
+- `kwargs` (`Dict[str, Any]`): The keyword arguments to pass to the evaluation function
+
+**Returns**:
+
+`Tuple[bool, Any]`: A tuple (bool, result) where bool is True if the function ran successfully and result is the output of the function
+
+<a id="EvaluatorFlowModule.EvaluatorFlow.EvaluatorFlow.evaluate_program"></a>
+
+#### evaluate\_program
+
+```python
+def evaluate_program(program: str, **kwargs)
+```
+
+Evaluate the program using the evaluation function
+
+**Arguments**:
+
+- `program` (`str`): The program to evaluate
+- `kwargs` (`Dict[str, Any]`): The keyword arguments to pass to the evaluation function
+
+**Returns**:
+
+`Tuple[bool, Any]`: A tuple (bool, result) where bool is True if the function ran successfully and result is the output of the function
+
+<a id="EvaluatorFlowModule.EvaluatorFlow.EvaluatorFlow.analyse"></a>
+
+#### analyse
+
+```python
+def analyse(program: str)
+```
+
+Analyse the program on the test inputs
+
+**Arguments**:
+
+- `program` (`str`): The program to evaluate
+
+**Returns**:
+
+`Dict[str, Dict[str, Any]]`: A dictionary of scores per test input
+
+<a id="EvaluatorFlowModule.EvaluatorFlow.EvaluatorFlow.run"></a>
+
+#### run
+
+```python
+def run(input_message: FlowMessage)
+```
+
+This method runs the flow. It's the main method of the flow.
+
+**Arguments**:
+
+- `input_message` (`FlowMessage`): The input message
+

SamplerFlowModule/SamplerFlow.py

@@ -0,0 +1,87 @@
+""" 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)

SamplerFlowModule/SamplerFlow.yaml

@@ -0,0 +1,99 @@
+input_interface:
+  - "action"
+  - "content"
+
+
+_target_: flow_modules.aiflows.FunSearchFlowModule.SamplerFlow.instantiate_from_default_config
+name: SamplerFlowModule
+description: A flow that queries an LLM model to generate prompts for the Sampler flow
+
+input_interface_initialized:
+  - header
+  - code
+
+input_interface_non_initialized:
+  - header
+  - code
+
+output_interface:
+  - "api_output"
+  - "from"
+
+backend:
+  api_infos: ???
+  model_name: gpt-4
+  frequency_penalty: 0
+  max_tokens: 2000
+  n: 1
+  presence_penalty: 0
+  stream: false
+  temperature: 1.0
+  top_p: 0.4
+
+previous_messages:
+  last_k: 1
+  first_k: 1
+
+system_message_prompt_template:
+  template: |2-
+    
+    Role: You are an expert coder and algorithm designer who provides solutions.
+
+
+    Goal: Write evolve functions that optimize the score of an evaluation function (provided). Aim to achieve higher scores than previous functions. Write a new function that optimizes the evaluation function and scores better than the given functions. Develop non-trivial functions and actual algorithms.
+
+
+    Resources:
+
+    You will be given code related to evolve and functions.
+    You will be given a few functions attempting to optimize the evaluation function.
+    The function you're evolving (the last one passed) is the only one you have to write. Make sure to include only nested functions or recursive calls.
+    
+    
+    Naming:
+
+    Evaluation function is named {{evaluate_name}}.
+    Evolve functions have a naming structure similar to: {{evolve_name}}_v.
+
+
+    Important Constraints:
+
+    Only write one evolve function solution: Write only the uncompleted function given to you (i.e., only {{evolve_name}}_v{i}, where {i} is the function with the largest number in the prompt, and i is at most {{artifacts_per_prompt}}). Stop writing once you're done with your first solution.
+    Output must be a function that is parsable by the AST library: Write code only within functions. No textual comments, no code blocks (like ```python), or no explanations. Provide only Python code encapsulated in a single function. Your output should be parsable by the AST Python library. Do not include starting or ending comments like "Here's an improved version."
+    Follow the naming structure: Keep the function name as it is in the input.
+    Keep the same function format: Maintain the same input and output format in your solution.
+    Self-contained solution: You can't call other functions (only recursive calls or calls to nested functions).
+    Do not use the evaluate function in the solution: Avoid using the evaluate function in your solution.
+    
+  partial_variables:
+    evaluate_name: ???
+    evolve_name: ???
+    artifacts_per_prompt: ???
+
+
+init_human_message_prompt_template:
+  _target_: aiflows.prompt_template.JinjaPrompt
+  template: |2-
+    functions for evaluation:
+    {{header}}
+
+    function to evolve and to evaluate:
+    {{code}}
+  input_variables: ["header","code"]
+  partial_variables: {}
+
+human_message_prompt_template:
+  _target_: aiflows.prompt_template.JinjaPrompt
+  template: |2-
+    full code:
+    {{header}}
+
+    function to evolve and to evaluate:
+    {{code}}
+  input_variables: ["header","code"]
+  partial_variables: {}
+  
+
+
+  
+

SamplerFlowModule/__init__.py

@@ -0,0 +1,13 @@
+# ~~~ Specify the dependencies ~~~
+
+dependencies = [
+    {"url": "aiflows/ChatFlowModule", "revision": "main"}
+]
+# Revision can correspond toa branch, commit hash or a absolute path to a local directory (ideal for development)
+from aiflows import flow_verse
+
+flow_verse.sync_dependencies(dependencies)
+
+# ~~~ Import of your flow class (if you have any) ~~~
+# from .NAMEOFYOURFLOW import NAMEOFYOURFLOWCLASS
+from .SamplerFlow import SamplerFlow

SamplerFlowModule/demo.yaml

@@ -0,0 +1,10 @@
+_target_: flow_modules.aiflows.FunSearchFlowModule.SamplerFlow.instantiate_from_default_config
+
+backend:
+  api_infos: ???
+
+system_message_prompt_template:
+  partial_variables:
+    evaluate_name: evaluate
+    evolve_name: solve_function
+    artifacts_per_prompt: 2

SamplerFlowModule/run.py

@@ -0,0 +1,181 @@
+import os
+
+import hydra
+
+import aiflows
+from aiflows.backends.api_info import ApiInfo
+from aiflows.utils.general_helpers import read_yaml_file, quick_load_api_keys
+
+from aiflows import logging
+from aiflows.flow_cache import CACHING_PARAMETERS, clear_cache
+
+from aiflows.utils import serving
+from aiflows.workers import run_dispatch_worker_thread
+from aiflows.messages import FlowMessage
+from aiflows.interfaces import KeyInterface
+from aiflows.utils.colink_utils import start_colink_server
+from aiflows import flow_verse
+
+
+dependencies = [
+    {
+        "url": "aiflows/FunSearchFlowModule",
+        "revision": os.path.abspath("../")
+    }
+]
+flow_verse.sync_dependencies(dependencies)
+
+logging.set_verbosity_debug()
+
+
+if __name__ == "__main__":
+    
+    cl = start_colink_server()
+    
+    serving.recursive_serve_flow(
+        cl=cl,
+        flow_class_name="flow_modules.aiflows.FunSearchFlowModule.SamplerFlow",
+        flow_endpoint="SamplerFlow",
+    )
+        
+    run_dispatch_worker_thread(cl)
+    
+    config_overrides = read_yaml_file(os.path.join(".", "demo.yaml"))
+
+    api_information = [ApiInfo(backend_used="openai",
+                              api_key = os.getenv("OPENAI_API_KEY"))]
+    
+    quick_load_api_keys(config_overrides, api_information, key="api_infos")
+
+    funsearch_proxy = serving.get_flow_instance(
+        cl=cl,
+        flow_endpoint="SamplerFlow",
+        config_overrides=config_overrides,
+    )
+    
+    code = \
+"""
+#function used to evaluate the program:
+def evaluate(solve_function: str, tests_inputs: List[str], expected_outputs: str) -> float:
+    \"\"\"Returns the score of the solve function we're evolving based on the tests_inputs and expected_outputs.
+    Scores are between 0 and 1, unless the program fails to run, in which case the score is -1.
+    \"\"\"
+    if solve(solve_function, tests_inputs, expected_outputs) == True:
+        return 1.0
+    return 0.0
+
+
+def solve_function_v0(input) -> str:
+    \"\"\"Scores per test: test_1:{'score': 1.0, 'feedback': 'No feedback available.'} test_2:{'score': 1.0, 'feedback': 'No feedback available.'} test_3:{'score': 0.0, 'feedback': 'No feedback available.'} test_4:{'score': -1, 'feedback': 'Invalid Format of prediction'}\"\"\"
+    return 'YES'
+
+
+def solve_function_v1(input) -> str:
+    \"\"\"Improved version of solve_function_v0\"\"\"
+"""
+
+    header = \
+"""
+\"\"\"Problem Description:
+Serval has a string s that only consists of 0 and 1 of length n. The i-th character of s is denoted as s_i, where 1\leq i\leq n.
+Serval can perform the following operation called Inversion Magic on the string s:
+Choose an segment [l, r] (1\leq l\leq r\leq n). For l\leq i\leq r, change s_i into 1 if s_i is 0, and change s_i into 0 if s_i is 1.
+For example, let s be 010100 and the segment [2,5] is chosen. The string s will be 001010 after performing the Inversion Magic.
+Serval wants to make s a palindrome after performing Inversion Magic exactly once. Help him to determine whether it is possible.
+A string is a palindrome iff it reads the same backwards as forwards. For example, 010010 is a palindrome but 10111 is not.
+
+Input Description:
+Input
+Each test contains multiple test cases. The first line contains the number of test cases t (1\leq t\leq 10^4). The description of the test cases follows.
+The first line of each test case contains a single integer n (2\leq n\leq 10^5) — the length of string s.
+The second line of each test case contains a binary string s of length n. Only characters 0 and 1 can appear in s.
+It's guaranteed that the sum of n over all test cases does not exceed 2\cdot 10^5.
+
+Output Description:
+Output
+For each test case, print Yes if s can be a palindrome after performing Inversion Magic exactly once, and print No if not.
+You can output Yes and No in any case (for example, strings yEs, yes, Yes and YES will be recognized as a positive response).
+
+Public Tests:
+Test 1:
+  Input: ['1', '4', '1001']
+  Output: 'YES'
+Test 2:
+  Input: ['1', '5', '10010']
+  Output: 'YES'
+Test 3:
+  Input: ['1', '7', '0111011']
+  Output: 'NO'
+
+\"\"\"
+
+
+import ast
+import itertools
+import numpy as np
+from typing import List
+
+def solve(solve_function: str,input: List[str], expected_output: str) -> str:
+    \"\"\"function used to run the solve function on input *kwargs and return the the predicted output
+    
+    :param solve_function: the function to run (the solve function below as a string)
+    :type solve_function: str
+    :param kwargs: the inputs to the solve function
+    :type kwargs: List[str]
+    \"\"\"
+    local_namespace = {}
+    exec(solve_function,local_namespace)
+    found_name, program_name = get_function_name_from_code(solve_function)
+    
+    if not found_name:
+        raise ValueError(f"Function name not found in program: {solve_function}")
+    
+    solve_fn = local_namespace.get(program_name)
+    
+    prediction = solve_fn(input)
+    
+    prediction = prediction.split()
+    expected_output = expected_output.split()
+    
+    if len(prediction) != len(expected_output):
+        raise ValueError(f"Invalid Format of prediction")
+    
+    for i in range(len(prediction)):
+        if prediction[i] != expected_output[i]:
+            return False
+    
+    return True
+
+def evaluate(solve_function: str, tests_inputs: List[str], expected_outputs: str) -> float:
+    \"\"\"Returns the score of the solve function we're evolving based on the tests_inputs and expected_outputs.
+    Scores are between 0 and 1, unless the program fails to run, in which case the score is -1.
+    \"\"\"
+    if solve(solve_function,tests_inputs,expected_outputs) == True:
+        return 1.0
+    return 0.0
+
+
+def get_function_name_from_code(code):
+    tree = ast.parse(code)
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef):
+            return True, node.name
+
+    # something is wrong
+    return False, None
+
+"""
+    
+    data = {
+        'code': code,
+        'header': header    
+    }
+        
+    input_message = funsearch_proxy.package_input_message(data = data)
+    
+    funsearch_proxy.send_message(input_message)
+    
+    future = funsearch_proxy.get_reply_future(input_message)
+    response = future.get_data()
+    print("~~~Response~~~")
+    print(response)

__init__.py

@@ -0,0 +1,4 @@
+from .FunSearch import FunSearch 
+from .SamplerFlowModule.SamplerFlow import SamplerFlow
+from .EvaluatorFlowModule.EvaluatorFlow import EvaluatorFlow
+from .ProgramDBFlowModule.ProgramDBFlow import ProgramDBFlow

cf_functions.py

@@ -0,0 +1,60 @@
+
+import ast
+import itertools
+import numpy as np
+from typing import List
+
+def solve(solve_function: str,input: List[str], expected_output: str) -> str:
+    """function used to run the solve function on input *kwargs and return the the predicted output
+    
+    :param solve_function: the function to run (the solve function below as a string)
+    :type solve_function: str
+    :param kwargs: the inputs to the solve function
+    :type kwargs: List[str]
+    """
+    local_namespace = {}
+    exec(solve_function,local_namespace)
+    found_name, program_name = get_function_name_from_code(solve_function)
+    
+    if not found_name:
+        raise ValueError(f"Function name not found in program: {solve_function}")
+    
+    solve_fn = local_namespace.get(program_name)
+    
+    prediction = solve_fn(input)
+    
+    prediction = prediction.split()
+    expected_output = expected_output.split()
+    
+    if len(prediction) != len(expected_output):
+        raise ValueError(f"Invalid Format of prediction")
+    
+    for i in range(len(prediction)):
+        if prediction[i] != expected_output[i]:
+            return False
+    
+    return True
+
+def evaluate(solve_function: str, tests_inputs: List[str], expected_outputs: str) -> float:
+    """Returns the score of the solve function we're evolving based on the tests_inputs and expected_outputs.
+    Scores are between 0 and 1, unless the program fails to run, in which case the score is -1.
+    """
+    if solve(solve_function,tests_inputs,expected_outputs) == True:
+        return 1.0
+    return 0.0
+
+
+def get_function_name_from_code(code):
+    tree = ast.parse(code)
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef):
+            return True, node.name
+
+    # something is wrong
+    return False, None
+
+
+
+
+    
+    

demo.yaml

@@ -0,0 +1,34 @@
+name: FunSearchFlow
+description: "A flow implementing FunSearch Asynchronous Search"
+
+### Subflows specification
+subflows_config:
+  SamplerFlow:
+    # _target_: flow_modules.SamplerFlowModule.SamplerFlow.instantiate_from_default_config
+    name: "Sampler Flow"
+    description: "A flow that queries an LLM model to generate prompts"
+    system_message_prompt_template:
+      partial_variables: ???
+    backend:
+      api_infos: ???
+  EvaluatorFlow:
+    name: "A flow that evaluates code on tests"
+    description: "A flow that evaluates code on tests"
+    singleton: False
+    run_error_score: ???
+    py_file: ???
+    function_to_run_name: ???
+    test_inputs: ???
+    timeout_seconds: 10
+    use_test_input_as_key: ???
+
+  ProgramDBFlow:
+    # _target_: flow_modules.ProgramDBFlowModule.ProgramDBFlow.instantiate_from_default_config
+    name: "ProgramDB"
+    description: "A flow that registers samples and evaluations in a database"
+    evaluate_function: ???
+    evaluate_file_full_content: ???
+    artifact_to_evolve_name: solve_function
+    num_islands: 3
+    template:
+      preface: ""

pip_requirements.txt

@@ -0,0 +1,4 @@
+astunparse==1.6.3
+numpy==1.26.2
+SciPy==1.11.4
+pandas

run.py

@@ -0,0 +1,241 @@
+import os
+
+import hydra
+
+import aiflows
+from aiflows.backends.api_info import ApiInfo
+from aiflows.utils.general_helpers import read_yaml_file, quick_load_api_keys
+
+from aiflows import logging
+from aiflows.flow_cache import CACHING_PARAMETERS, clear_cache
+
+from aiflows.utils import serving
+from aiflows.workers import run_dispatch_worker_thread
+from aiflows.messages import FlowMessage
+from aiflows.interfaces import KeyInterface
+from aiflows.utils.colink_utils import start_colink_server
+from aiflows import flow_verse
+import pandas as pd
+import sys
+from copy import deepcopy
+import requests
+import time
+
+dependencies = [
+    {
+        "url": "aiflows/FunSearchFlowModule",
+        "revision": "../FunSearchFlowModule"
+    }
+]
+flow_verse.sync_dependencies(dependencies)
+from flow_modules.aiflows.FunSearchFlowModule.Loader import Loader
+
+logging.set_verbosity_debug()
+
+
+def load_problem(id, ds_location = "./data/codeforces.jsonl.gz"):
+    def make_problem_descriptions_str(row):
+        def write_public_tests_individual_io_str(row):
+            public_tests = row.public_tests_individual_io
+            tests = ""
+            for i,test in enumerate(public_tests):
+                input = test[0]
+                output = test[1]
+                tests += f"Test {i+1}:\n  Input: {input}\n  Output: \'{output}\'\n"
+            return tests
+        
+        problem_descritption = row.problem_description
+        input_descriptions = row.input_description
+        ouput_descriptions = row.output_description
+        public_tests = write_public_tests_individual_io_str(row)
+        
+        problem_description_str = f"Problem Description:\n{problem_descritption}\n\n"
+        input_description_str = f"Input Description:\n{input_descriptions}\n\n"
+        output_description_str = f"Output Description:\n{ouput_descriptions}\n\n"
+        public_tests_str = f"Public Tests:\n{public_tests}\n"
+        
+        final_str = problem_description_str + input_description_str + output_description_str +public_tests_str
+        return final_str
+    
+    df = pd.read_json(ds_location, lines=True, compression='gzip')
+    row = df[df.id == id].iloc[0]
+
+    assert row.non_unique_output == False, "Problem has non unique output. Not supported yet"
+
+    problem_description = make_problem_descriptions_str(row)
+    public_test = row.public_tests_individual_io
+    tests = {}
+    test_counter = 1
+
+    for public_test in public_test:
+        tests["test_"+str(test_counter)] = {"tests_inputs": public_test[0], "expected_outputs": public_test[1]}
+        test_counter += 1
+        
+    for hidden_test in row.hidden_tests_io:
+        tests["test_"+str(test_counter)] = {"tests_inputs": hidden_test[0], "expected_outputs": hidden_test[1]}
+        test_counter += 1
+    
+    return  tests, problem_description 
+
+def download_codeforces_data(data_folder_path,file_name):
+    print("Downloading data....")
+    os.makedirs(data_folder_path, exist_ok=True)
+    url = "https://github.com/epfl-dlab/cc_flows/raw/main/data/codeforces/codeforces.jsonl.gz"
+    response = requests.get(url, stream=True)
+    
+    if response.status_code == 200:
+        with open(os.path.join(data_folder_path,file_name), 'wb') as file:
+            for chunk in response:
+                file.write(chunk)
+        print("Download complete")
+    else:
+        print("Failed to download data", response.status_code)
+        
+
+def get_configs(problem_id, ds_location = "./data/codeforces.jsonl.gz"):
+    tests,  problem_description = load_problem(problem_id,ds_location)
+    
+    path = os.path.join(".", "demo.yaml")
+    funsearch_cfg = read_yaml_file(path)
+       
+    evaluate_function_file_path: str = "./cf_functions.py"
+    evaluate_function_name: str = "evaluate"
+    evolve_function_name:str = "solve_function"
+     
+    loader = Loader(file_path = evaluate_function_file_path, target_name = evaluate_function_name)
+    evaluate_function: str= loader.load_target()
+    evaluate_file_full_content = loader.load_full_file()
+    
+    evaluate_file_full_content = f"\"\"\"{problem_description}\"\"\"\n\n" + evaluate_file_full_content
+    
+    #~~~~~ ProgramDBFlow Overrides ~~~~~~~~
+    funsearch_cfg["subflows_config"]["ProgramDBFlow"]["evaluate_function"] = evaluate_function
+    funsearch_cfg["subflows_config"]["ProgramDBFlow"]["evaluate_file_full_content"] = evaluate_file_full_content
+    funsearch_cfg["subflows_config"]["ProgramDBFlow"]["artifact_to_evolve_name"] = evolve_function_name
+    
+    if len(tests) > 0:
+        first_test = tests["test_1"]
+
+        dummy_solution = f"def {evolve_function_name}(input) -> str:" +\
+            "\n    \"\"\"Attempt at solving the problem given the input input and returns the predicted output (see the top of the file for problem description)\"\"\"" +\
+            f"\n    return \'{first_test['expected_outputs']}\'\n"
+
+    
+    else:
+        dummy_solution = f"def {evolve_function_name}(input) -> str:" +\
+            "\n    \"\"\"Attempt at solving the problem given the input input and returns the predicted output (see the top of the file for problem description)\"\"\"" +\
+            f"\n    return 0.0\"\"\n"
+    
+    #~~~~~~~~~~Evaluator overrides~~~~~~~~~~~~
+    funsearch_cfg["subflows_config"]["EvaluatorFlow"]["py_file"] = evaluate_file_full_content
+    funsearch_cfg["subflows_config"]["EvaluatorFlow"]["run_error_score"] = -1
+    funsearch_cfg["subflows_config"]["EvaluatorFlow"]["function_to_run_name"] = evaluate_function_name
+    funsearch_cfg["subflows_config"]["EvaluatorFlow"]["test_inputs"] = tests
+    #Hides test inputs from LLM (necessary for hidden tests. Makes same setup as in a real contest.)
+    funsearch_cfg["subflows_config"]["EvaluatorFlow"]["use_test_input_as_key"] = False
+   
+    
+    #~~~~~~~~~~Sampler overrides~~~~~~~~~~~~  
+    funsearch_cfg["subflows_config"]["SamplerFlow"]["system_message_prompt_template"]["partial_variables"] = \
+        {
+            "evaluate_name": evaluate_function_name,
+            "evolve_name": evolve_function_name,
+            "artifacts_per_prompt": 2
+        }
+        
+        
+    return funsearch_cfg, dummy_solution
+
+
+FLOW_MODULES_PATH = "./"
+
+
+if __name__ == "__main__":
+    
+    cl = start_colink_server()
+    
+    problem_id = "1789B" #put the problem id here
+    
+    if not os.path.exists("./data/codeforces.jsonl.gz"):
+        download_codeforces_data("./data", "codeforces.jsonl.gz")
+   
+    funsearch_cfg, dummy_solution = get_configs(problem_id)
+    #Serve Program Database and get its flow type explicitly
+    api_information = [ApiInfo(backend_used="openai",
+                              api_key = os.getenv("OPENAI_API_KEY"))]
+    
+    serving.recursive_serve_flow(
+        cl=cl,
+        flow_class_name="flow_modules.aiflows.FunSearchFlowModule.FunSearch",
+        flow_endpoint="FunSearch",
+    )
+    
+    # #Serve the rest
+    # serving.recursive_serve_flow(
+    #     cl=cl,
+    #     flow_type="FunSearch_served",
+    #     default_config=funsearch_cfg,
+    #     default_state=None,
+    #     default_dispatch_point="coflows_dispatch",
+    # )
+    n_workers = 10
+    for i in range(n_workers):
+        run_dispatch_worker_thread(cl)
+    
+    quick_load_api_keys(funsearch_cfg, api_information, key="api_infos")
+    config_overrides = None
+    #Mount ProgramDBFlow first to get it's flow ref
+    funsearch_proxy = serving.get_flow_instance(
+        cl=cl,
+        flow_endpoint="FunSearch",
+        config_overrides=funsearch_cfg,
+    )
+    
+    data = {
+        "from": "SamplerFlow",
+        "operation": "register_program",
+        "api_output": dummy_solution
+    }
+    
+    input_message = funsearch_proxy.package_input_message(data = data)
+    
+    funsearch_proxy.send_message(input_message)
+    
+    
+    data = {
+        "from": "FunSearch",
+        "operation": "start",
+        "content": {"num_samplers": 5},
+    }
+    
+    input_message = funsearch_proxy.package_input_message(data = data)
+    
+    funsearch_proxy.send_message(input_message)
+
+    data = {
+        "from": "FunSearch",
+        "operation": "stop",
+        "content": {},
+    }
+    
+    input_message = funsearch_proxy.package_input_message(data = data)
+    
+    funsearch_proxy.send_message(input_message)
+    
+    
+    wait_time = 1000
+    print(f"Waiting {wait_time} seconds  before requesting result...")
+    time.sleep(wait_time)
+
+    data = {
+        "from": "FunSearch",
+        "operation": "get_best_programs_per_island",
+        "content": {}
+    }
+    
+    input_message = funsearch_proxy.package_input_message(data = data)
+    
+    future = funsearch_proxy.get_reply_future(input_message)
+    print("waiting for response....")
+    response = future.get_data()
+    print(response)