examples/FunSearch/FunSearch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FunSearch Demo:\n",
    "\n",
    "**In this notebook you will:**\n",
    "- Learn how to use FunSearch on the task of solving codeforces problems [Section 1](#section-1-solving-codeforces-problems-with-funsearch)\n",
    "- Learn the various ways to interact with FunSearch [Section 1](#section-1-solving-codeforces-problems-with-funsearch)\n",
    "- Learn how to adapt FunSearch to your own problem solving needs [Section 2](#section-2-adapting-funsearch-to-your-own-problem)\n",
    "\n",
    "\n",
    "## What is FunSearch?\n",
    "\n",
    "FunSearch, born from a paper featured in Nature by DeepMind, can be succinctly characterized as follows, mirroring their website's description:\n",
    "\n",
    "> 'FunSearch, a method to search for new solutions in mathematics and computer science. FunSearch works by pairing a pre-trained LLM, whose goal is to provide creative solutions in the form of computer code, with an automated “evaluator”, which guards against hallucinations and incorrect ideas. By iterating back-and-forth between these two components, initial solutions “evolve” into new knowledge. The system searches for “functions” written in computer code; hence the name FunSearch.'  - [[1]](https://deepmind.google/discover/blog/funsearch-making-new-discoveries-in-mathematical-sciences-using-large-language-models/) [[2]](https://www.nature.com/articles/s41586-023-06924-6)\n",
    "\n",
    "**Pre-requisites:**\n",
    "\n",
    "We strongly recommend reading the original paper [[2]](https://www.nature.com/articles/s41586-023-06924-6) for a more in-depth understanding of FunSearch.\n",
    "\n",
    "[1] DeepMind, 2023. [\"FunSearch: Making new discoveries in mathematical sciences using Large Language Models\"](https://deepmind.google/discover/blog/funsearch-making-new-discoveries-in-mathematical-sciences-using-large-language-models/)\n",
    "\n",
    "[2] [Romera-Paredes, Bernardino, Mohammadamin Barekatain, Alexander Novikov, Matej Balog, M. Pawan Kumar, Emilien Dupont, Francisco JR Ruiz et al. \"Mathematical discoveries from program search with large language models.\" Nature 625, no. 7995 (2024): 468-475](https://www.nature.com/articles/s41586-023-06924-6)\n",
    "\n",
    "⚠️ **Warning**:\n",
    "\n",
    " This notebook utilizes GPT-4, which incurs a usage cost. Please ensure you're aware of any associated charges before running the code in this notebook.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports, Dependencies, and Starting a Local Server (No action required, just run the cells)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Fetching 36 files: 100%|██████████| 36/36 [00:00<00:00, 154233.85it/s]\n",
      "Fetching 36 files: 100%|██████████| 36/36 [00:00<00:00, 144079.15it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: astunparse==1.6.3 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from -r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 1)) (1.6.3)\n",
      "Requirement already satisfied: numpy==1.26.2 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from -r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 2)) (1.26.2)\n",
      "Requirement already satisfied: SciPy==1.11.4 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from -r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 3)) (1.11.4)\n",
      "Requirement already satisfied: aiflows>=1.1.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from -r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.1.1)\n",
      "Requirement already satisfied: pandas in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from -r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 5)) (2.2.1)\n",
      "Requirement already satisfied: wheel<1.0,>=0.23.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from astunparse==1.6.3->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 1)) (0.41.2)\n",
      "Requirement already satisfied: six<2.0,>=1.6.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from astunparse==1.6.3->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 1)) (1.16.0)\n",
      "Requirement already satisfied: colorama==0.4.6 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.4.6)\n",
      "Requirement already satisfied: pytest==7.3.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (7.3.1)\n",
      "Requirement already satisfied: pytest-cov==4.1.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (4.1.0)\n",
      "Requirement already satisfied: hydra-core==1.3.2 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.3.2)\n",
      "Requirement already satisfied: hydra-colorlog==1.1.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.1.0)\n",
      "Requirement already satisfied: wrapt-timeout-decorator==1.3.12.2 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.3.12.2)\n",
      "Requirement already satisfied: diskcache==5.6.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (5.6.1)\n",
      "Requirement already satisfied: openai==1.0.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.0.0)\n",
      "Requirement already satisfied: huggingface-hub==0.19.4 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.19.4)\n",
      "Requirement already satisfied: jsonlines==3.1.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (3.1.0)\n",
      "Requirement already satisfied: jinja2==3.1.2 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (3.1.2)\n",
      "Requirement already satisfied: mock==5.0.2 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (5.0.2)\n",
      "Requirement already satisfied: rich==12.6.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (12.6.0)\n",
      "Requirement already satisfied: litellm==1.0.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.0.0)\n",
      "Requirement already satisfied: colink==0.3.7 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.3.7)\n",
      "Requirement already satisfied: termcolor==2.4.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (2.4.0)\n",
      "Requirement already satisfied: streamlit==1.32.2 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.32.2)\n",
      "Requirement already satisfied: grpcio>=1.27.2 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from colink==0.3.7->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.62.1)\n",
      "Requirement already satisfied: protobuf<5.0dev,>=3.19.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from colink==0.3.7->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (4.25.3)\n",
      "Requirement already satisfied: coincurve>=18.0.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from colink==0.3.7->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (19.0.1)\n",
      "Requirement already satisfied: cryptography>=39.0.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from colink==0.3.7->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (42.0.5)\n",
      "Requirement already satisfied: pika>=1.2.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from colink==0.3.7->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.3.2)\n",
      "Requirement already satisfied: pyjwt>=2.6.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from colink==0.3.7->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (2.8.0)\n",
      "Requirement already satisfied: requests>=2.28.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from colink==0.3.7->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (2.31.0)\n",
      "Requirement already satisfied: requests-toolbelt>=0.10.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from colink==0.3.7->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.0.0)\n",
      "Requirement already satisfied: redis>=4.5.4 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from colink==0.3.7->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (5.0.3)\n",
      "Requirement already satisfied: filelock in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from huggingface-hub==0.19.4->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (3.13.1)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from huggingface-hub==0.19.4->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (2024.2.0)\n",
      "Requirement already satisfied: tqdm>=4.42.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from huggingface-hub==0.19.4->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (4.66.2)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from huggingface-hub==0.19.4->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (6.0.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from huggingface-hub==0.19.4->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (4.10.0)\n",
      "Requirement already satisfied: packaging>=20.9 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from huggingface-hub==0.19.4->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (23.2)\n",
      "Requirement already satisfied: colorlog in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from hydra-colorlog==1.1.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (6.8.2)\n",
      "Requirement already satisfied: omegaconf<2.4,>=2.2 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from hydra-core==1.3.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (2.3.0)\n",
      "Requirement already satisfied: antlr4-python3-runtime==4.9.* in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from hydra-core==1.3.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (4.9.3)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from jinja2==3.1.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (2.1.5)\n",
      "Requirement already satisfied: attrs>=19.2.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from jsonlines==3.1.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (23.2.0)\n",
      "Requirement already satisfied: aiohttp in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from litellm==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (3.9.3)\n",
      "Requirement already satisfied: appdirs<2.0.0,>=1.4.4 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from litellm==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.4.4)\n",
      "Requirement already satisfied: certifi<2024.0.0,>=2023.7.22 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from litellm==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (2023.11.17)\n",
      "Requirement already satisfied: click in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from litellm==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (8.1.7)\n",
      "Requirement already satisfied: importlib-metadata>=6.8.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from litellm==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (6.11.0)\n",
      "Requirement already satisfied: python-dotenv>=0.2.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from litellm==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.0.1)\n",
      "Requirement already satisfied: tiktoken>=0.4.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from litellm==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.6.0)\n",
      "Requirement already satisfied: tokenizers in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from litellm==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.15.2)\n",
      "Requirement already satisfied: anyio<4,>=3.5.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from openai==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (3.7.1)\n",
      "Requirement already satisfied: distro<2,>=1.7.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from openai==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.9.0)\n",
      "Requirement already satisfied: httpx<1,>=0.23.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from openai==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.27.0)\n",
      "Requirement already satisfied: pydantic<3,>=1.9.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from openai==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.10.14)\n",
      "Requirement already satisfied: iniconfig in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from pytest==7.3.1->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (2.0.0)\n",
      "Requirement already satisfied: pluggy<2.0,>=0.12 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from pytest==7.3.1->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.4.0)\n",
      "Requirement already satisfied: coverage>=5.2.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from coverage[toml]>=5.2.1->pytest-cov==4.1.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (7.4.4)\n",
      "Requirement already satisfied: commonmark<0.10.0,>=0.9.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from rich==12.6.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.9.1)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.6.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from rich==12.6.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (2.17.2)\n",
      "Requirement already satisfied: altair<6,>=4.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (5.2.0)\n",
      "Requirement already satisfied: blinker<2,>=1.0.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.7.0)\n",
      "Requirement already satisfied: cachetools<6,>=4.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (5.3.3)\n",
      "Requirement already satisfied: pillow<11,>=7.1.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (10.2.0)\n",
      "Requirement already satisfied: pyarrow>=7.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (15.0.2)\n",
      "Requirement already satisfied: tenacity<9,>=8.1.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (8.2.3)\n",
      "Requirement already satisfied: toml<2,>=0.10.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.10.2)\n",
      "Requirement already satisfied: gitpython!=3.1.19,<4,>=3.0.7 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (3.1.42)\n",
      "Requirement already satisfied: pydeck<1,>=0.8.0b4 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.8.1b0)\n",
      "Requirement already satisfied: tornado<7,>=6.0.3 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (6.4)\n",
      "Requirement already satisfied: cli-exit-tools in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from wrapt-timeout-decorator==1.3.12.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.2.6)\n",
      "Requirement already satisfied: lib-detect-testenv in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from wrapt-timeout-decorator==1.3.12.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (2.0.8)\n",
      "Requirement already satisfied: multiprocess in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from wrapt-timeout-decorator==1.3.12.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.70.16)\n",
      "Requirement already satisfied: wrapt in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from wrapt-timeout-decorator==1.3.12.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.16.0)\n",
      "Requirement already satisfied: dill in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from wrapt-timeout-decorator==1.3.12.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.3.8)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from pandas->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 5)) (2.9.0)\n",
      "Requirement already satisfied: pytz>=2020.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from pandas->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 5)) (2024.1)\n",
      "Requirement already satisfied: tzdata>=2022.7 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from pandas->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 5)) (2024.1)\n",
      "Requirement already satisfied: jsonschema>=3.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from altair<6,>=4.0->streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (4.21.1)\n",
      "Requirement already satisfied: toolz in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from altair<6,>=4.0->streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.12.1)\n",
      "Requirement already satisfied: idna>=2.8 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from anyio<4,>=3.5.0->openai==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (3.6)\n",
      "Requirement already satisfied: sniffio>=1.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from anyio<4,>=3.5.0->openai==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.3.1)\n",
      "Requirement already satisfied: asn1crypto in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from coincurve>=18.0.0->colink==0.3.7->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.5.1)\n",
      "Requirement already satisfied: cffi>=1.3.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from coincurve>=18.0.0->colink==0.3.7->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.16.0)\n",
      "Requirement already satisfied: gitdb<5,>=4.0.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from gitpython!=3.1.19,<4,>=3.0.7->streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (4.0.11)\n",
      "Requirement already satisfied: httpcore==1.* in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from httpx<1,>=0.23.0->openai==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.0.4)\n",
      "Requirement already satisfied: h11<0.15,>=0.13 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from httpcore==1.*->httpx<1,>=0.23.0->openai==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.14.0)\n",
      "Requirement already satisfied: zipp>=0.5 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from importlib-metadata>=6.8.0->litellm==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (3.17.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from requests>=2.28.1->colink==0.3.7->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (3.3.2)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from requests>=2.28.1->colink==0.3.7->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (2.2.1)\n",
      "Requirement already satisfied: regex>=2022.1.18 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from tiktoken>=0.4.0->litellm==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (2023.12.25)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiohttp->litellm==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.3.1)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiohttp->litellm==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.4.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiohttp->litellm==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (6.0.5)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from aiohttp->litellm==1.0.0->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (1.9.4)\n",
      "Requirement already satisfied: pycparser in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from cffi>=1.3.0->coincurve>=18.0.0->colink==0.3.7->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (2.21)\n",
      "Requirement already satisfied: smmap<6,>=3.0.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from gitdb<5,>=4.0.1->gitpython!=3.1.19,<4,>=3.0.7->streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (5.0.1)\n",
      "Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from jsonschema>=3.0->altair<6,>=4.0->streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (2023.12.1)\n",
      "Requirement already satisfied: referencing>=0.28.4 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from jsonschema>=3.0->altair<6,>=4.0->streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.34.0)\n",
      "Requirement already satisfied: rpds-py>=0.7.1 in /Users/nicolasbaldwin/opt/miniconda3/envs/mockenv/lib/python3.11/site-packages (from jsonschema>=3.0->altair<6,>=4.0->streamlit==1.32.2->aiflows>=1.1.1->-r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt (line 4)) (0.18.0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Fetching 9 files: 100%|██████████| 9/9 [00:00<00:00, 48897.33it/s]\n",
      "Fetching 9 files: 100%|██████████| 9/9 [00:00<00:00, 64860.37it/s]\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import hydra\n",
    "import aiflows\n",
    "from aiflows.backends.api_info import ApiInfo\n",
    "from aiflows.utils.general_helpers import read_yaml_file, quick_load_api_keys\n",
    "from aiflows.utils import serving\n",
    "from aiflows.workers import run_dispatch_worker_thread\n",
    "from aiflows.utils.colink_utils import start_colink_server\n",
    "from aiflows import flow_verse\n",
    "from aiflows.utils import logging\n",
    "#setting verbosity to warning because of the large amount of logs (seems to create lagging in notebook)\n",
    "logging.set_verbosity_warning()\n",
    "#If you want to see more logs, you can set the verbosity to info by uncommenting the following line\n",
    "#logging.set_verbosity_info()\n",
    "import json\n",
    "from aiflows.utils import serving\n",
    "import os\n",
    "\n",
    "#LoadFunSearch\n",
    "dependencies = [\n",
    "    {\n",
    "        \"url\": \"aiflows/FunSearchFlowModule\",\n",
    "        \"revision\": \"main\"\n",
    "    }\n",
    "]\n",
    "flow_verse.sync_dependencies(dependencies)\n",
    "\n",
    "!pip install -r flow_modules/aiflows/FunSearchFlowModule/pip_requirements.txt\n",
    "\n",
    "from utils import make_best_program_per_island_nice, dict_to_yaml, compile_and_writefile, download_codeforces_data, get_configs_cf\n",
    "import pandas as pd\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Start Local Colink Server\n",
    "cl = start_colink_server()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 1: Solving Codeforces Problems with FunSearch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Downloading Codeforces Dataset\n",
    "\n",
    "Let's start by downloading the codeforces dataset and taking a look at the first few rows. The full details of the dataset can be found [here](https://github.com/epfl-dlab/cc_flows?tab=readme-ov-file#codeforces)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(\"./data/codeforces.jsonl.gz\"):\n",
    "    download_codeforces_data(\"./data\", \"codeforces.jsonl.gz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Taking a Look at DS\n",
    "df = pd.read_json(\"./data/codeforces.jsonl.gz\", lines=True, compression='gzip')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>contest</th>\n",
       "      <th>problem_name</th>\n",
       "      <th>problem_fullname</th>\n",
       "      <th>problem_url</th>\n",
       "      <th>id</th>\n",
       "      <th>difficulty</th>\n",
       "      <th>tags</th>\n",
       "      <th>header</th>\n",
       "      <th>problem_description</th>\n",
       "      <th>input_description</th>\n",
       "      <th>...</th>\n",
       "      <th>hidden_tests_io_truncated</th>\n",
       "      <th>num_tests_non_truncated</th>\n",
       "      <th>solutions</th>\n",
       "      <th>solution_url</th>\n",
       "      <th>working_solution</th>\n",
       "      <th>id_hash</th>\n",
       "      <th>time_limit</th>\n",
       "      <th>memory_limit</th>\n",
       "      <th>interpolated_date</th>\n",
       "      <th>release_date</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1600</td>\n",
       "      <td>J</td>\n",
       "      <td>J. Robot Factory</td>\n",
       "      <td>https://codeforces.com/contest/1600/problem/J</td>\n",
       "      <td>1600J</td>\n",
       "      <td>1400</td>\n",
       "      <td>[bitmasks, dfs and similar]</td>\n",
       "      <td>J. Robot Factorytime limit per test1 secondmem...</td>\n",
       "      <td>You have received data from a Bubble bot. You ...</td>\n",
       "      <td>Input\\nThe first line has two numbers which ar...</td>\n",
       "      <td>...</td>\n",
       "      <td>[[[917 992, 13 15 9 14 15 15 11 12 11 14 15 9 ...</td>\n",
       "      <td>2</td>\n",
       "      <td>[{'verdict': 'Accepted', 'language': 'PyPy 3-6...</td>\n",
       "      <td>https://codeforces.com/contest/1600/submission...</td>\n",
       "      <td>import sys\\ninput = sys.stdin.readline\\n \\n \\n...</td>\n",
       "      <td>0c63a6eae6aefee1587c318ed6fbab140ab440b07575e4...</td>\n",
       "      <td>1 second</td>\n",
       "      <td>256 megabytes</td>\n",
       "      <td>None</td>\n",
       "      <td>2021-10-09</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1638</td>\n",
       "      <td>B</td>\n",
       "      <td>B. Odd Swap Sort</td>\n",
       "      <td>https://codeforces.com/contest/1638/problem/B</td>\n",
       "      <td>1638B</td>\n",
       "      <td>1100</td>\n",
       "      <td>[data structures, math, sortings]</td>\n",
       "      <td>B. Odd Swap Sorttime limit per test1 secondmem...</td>\n",
       "      <td>You are given an array a_1, a_2, \\dots, a_n. Y...</td>\n",
       "      <td>Input\\nEach test contains multiple test cases....</td>\n",
       "      <td>...</td>\n",
       "      <td>[[[33954, 1, 2, 1, 1, 2, 2 4, 2, 2 2, 2, 1 2, ...</td>\n",
       "      <td>1</td>\n",
       "      <td>[{'verdict': 'Accepted', 'language': 'PyPy 3-6...</td>\n",
       "      <td>https://codeforces.com/contest/1638/submission...</td>\n",
       "      <td>for s in[*open(0)][2::2]:\\n a=[[],[]]\\n for x ...</td>\n",
       "      <td>16faeda04f84a9f7f54640aa51249fcdc9cf264617373c...</td>\n",
       "      <td>1 second</td>\n",
       "      <td>256 megabytes</td>\n",
       "      <td>None</td>\n",
       "      <td>2022-02-14</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1638</td>\n",
       "      <td>A</td>\n",
       "      <td>A. Reverse</td>\n",
       "      <td>https://codeforces.com/contest/1638/problem/A</td>\n",
       "      <td>1638A</td>\n",
       "      <td>800</td>\n",
       "      <td>[constructive algorithms, greedy, math]</td>\n",
       "      <td>A. Reversetime limit per test1 secondmemory li...</td>\n",
       "      <td>You are given a permutation p_1, p_2, \\ldots, ...</td>\n",
       "      <td>Input\\nEach test contains multiple test cases....</td>\n",
       "      <td>...</td>\n",
       "      <td>[[[500, 1, 1, 2, 1 2, 2, 2 1, 3, 1 2 3, 3, 1 3...</td>\n",
       "      <td>1</td>\n",
       "      <td>[{'verdict': 'Accepted', 'language': 'Python 3...</td>\n",
       "      <td>https://codeforces.com/contest/1638/submission...</td>\n",
       "      <td>for s in[*open(0)][2::2]:\\n a=[*map(int,s.spli...</td>\n",
       "      <td>f031d386d7a22ba8571134c260aebd7b20f6cb9795b897...</td>\n",
       "      <td>1 second</td>\n",
       "      <td>256 megabytes</td>\n",
       "      <td>None</td>\n",
       "      <td>2022-02-14</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3 rows × 26 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   contest problem_name  problem_fullname  \\\n",
       "0     1600            J  J. Robot Factory   \n",
       "1     1638            B  B. Odd Swap Sort   \n",
       "2     1638            A        A. Reverse   \n",
       "\n",
       "                                     problem_url     id  difficulty  \\\n",
       "0  https://codeforces.com/contest/1600/problem/J  1600J        1400   \n",
       "1  https://codeforces.com/contest/1638/problem/B  1638B        1100   \n",
       "2  https://codeforces.com/contest/1638/problem/A  1638A         800   \n",
       "\n",
       "                                      tags  \\\n",
       "0              [bitmasks, dfs and similar]   \n",
       "1        [data structures, math, sortings]   \n",
       "2  [constructive algorithms, greedy, math]   \n",
       "\n",
       "                                              header  \\\n",
       "0  J. Robot Factorytime limit per test1 secondmem...   \n",
       "1  B. Odd Swap Sorttime limit per test1 secondmem...   \n",
       "2  A. Reversetime limit per test1 secondmemory li...   \n",
       "\n",
       "                                 problem_description  \\\n",
       "0  You have received data from a Bubble bot. You ...   \n",
       "1  You are given an array a_1, a_2, \\dots, a_n. Y...   \n",
       "2  You are given a permutation p_1, p_2, \\ldots, ...   \n",
       "\n",
       "                                   input_description  ...  \\\n",
       "0  Input\\nThe first line has two numbers which ar...  ...   \n",
       "1  Input\\nEach test contains multiple test cases....  ...   \n",
       "2  Input\\nEach test contains multiple test cases....  ...   \n",
       "\n",
       "                           hidden_tests_io_truncated  num_tests_non_truncated  \\\n",
       "0  [[[917 992, 13 15 9 14 15 15 11 12 11 14 15 9 ...                        2   \n",
       "1  [[[33954, 1, 2, 1, 1, 2, 2 4, 2, 2 2, 2, 1 2, ...                        1   \n",
       "2  [[[500, 1, 1, 2, 1 2, 2, 2 1, 3, 1 2 3, 3, 1 3...                        1   \n",
       "\n",
       "                                           solutions  \\\n",
       "0  [{'verdict': 'Accepted', 'language': 'PyPy 3-6...   \n",
       "1  [{'verdict': 'Accepted', 'language': 'PyPy 3-6...   \n",
       "2  [{'verdict': 'Accepted', 'language': 'Python 3...   \n",
       "\n",
       "                                        solution_url  \\\n",
       "0  https://codeforces.com/contest/1600/submission...   \n",
       "1  https://codeforces.com/contest/1638/submission...   \n",
       "2  https://codeforces.com/contest/1638/submission...   \n",
       "\n",
       "                                    working_solution  \\\n",
       "0  import sys\\ninput = sys.stdin.readline\\n \\n \\n...   \n",
       "1  for s in[*open(0)][2::2]:\\n a=[[],[]]\\n for x ...   \n",
       "2  for s in[*open(0)][2::2]:\\n a=[*map(int,s.spli...   \n",
       "\n",
       "                                             id_hash time_limit  \\\n",
       "0  0c63a6eae6aefee1587c318ed6fbab140ab440b07575e4...   1 second   \n",
       "1  16faeda04f84a9f7f54640aa51249fcdc9cf264617373c...   1 second   \n",
       "2  f031d386d7a22ba8571134c260aebd7b20f6cb9795b897...   1 second   \n",
       "\n",
       "    memory_limit interpolated_date release_date  \n",
       "0  256 megabytes              None   2021-10-09  \n",
       "1  256 megabytes              None   2022-02-14  \n",
       "2  256 megabytes              None   2022-02-14  \n",
       "\n",
       "[3 rows x 26 columns]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each Row in the dataset is a codeforces problem with a unique id in the `id` column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>contest</th>\n",
       "      <th>problem_name</th>\n",
       "      <th>problem_fullname</th>\n",
       "      <th>problem_url</th>\n",
       "      <th>id</th>\n",
       "      <th>difficulty</th>\n",
       "      <th>tags</th>\n",
       "      <th>header</th>\n",
       "      <th>problem_description</th>\n",
       "      <th>input_description</th>\n",
       "      <th>...</th>\n",
       "      <th>hidden_tests_io_truncated</th>\n",
       "      <th>num_tests_non_truncated</th>\n",
       "      <th>solutions</th>\n",
       "      <th>solution_url</th>\n",
       "      <th>working_solution</th>\n",
       "      <th>id_hash</th>\n",
       "      <th>time_limit</th>\n",
       "      <th>memory_limit</th>\n",
       "      <th>interpolated_date</th>\n",
       "      <th>release_date</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>142</th>\n",
       "      <td>1789</td>\n",
       "      <td>B</td>\n",
       "      <td>B. Serval and Inversion Magic</td>\n",
       "      <td>https://codeforces.com/contest/1789/problem/B</td>\n",
       "      <td>1789B</td>\n",
       "      <td>800</td>\n",
       "      <td>[brute force, implementation, strings, two poi...</td>\n",
       "      <td>B. Serval and Inversion Magictime limit per te...</td>\n",
       "      <td>Serval has a string s that only consists of 0 ...</td>\n",
       "      <td>Input\\nEach test contains multiple test cases....</td>\n",
       "      <td>...</td>\n",
       "      <td>[[[10000, 2, 00, 2, 01, 2, 10, 2, 11, 3, 000, ...</td>\n",
       "      <td>1</td>\n",
       "      <td>[{'verdict': 'Accepted', 'language': 'PyPy 3-6...</td>\n",
       "      <td>https://codeforces.com/contest/1789/submission...</td>\n",
       "      <td>class solve:\\n    def __init__(self):\\n       ...</td>\n",
       "      <td>c5c18430aba9033582c658845be4146bcbf13ccf522fa2...</td>\n",
       "      <td>1 second</td>\n",
       "      <td>256 megabytes</td>\n",
       "      <td>None</td>\n",
       "      <td>2023-02-25</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1 rows × 26 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     contest problem_name               problem_fullname  \\\n",
       "142     1789            B  B. Serval and Inversion Magic   \n",
       "\n",
       "                                       problem_url     id  difficulty  \\\n",
       "142  https://codeforces.com/contest/1789/problem/B  1789B         800   \n",
       "\n",
       "                                                  tags  \\\n",
       "142  [brute force, implementation, strings, two poi...   \n",
       "\n",
       "                                                header  \\\n",
       "142  B. Serval and Inversion Magictime limit per te...   \n",
       "\n",
       "                                   problem_description  \\\n",
       "142  Serval has a string s that only consists of 0 ...   \n",
       "\n",
       "                                     input_description  ...  \\\n",
       "142  Input\\nEach test contains multiple test cases....  ...   \n",
       "\n",
       "                             hidden_tests_io_truncated  \\\n",
       "142  [[[10000, 2, 00, 2, 01, 2, 10, 2, 11, 3, 000, ...   \n",
       "\n",
       "     num_tests_non_truncated  \\\n",
       "142                        1   \n",
       "\n",
       "                                             solutions  \\\n",
       "142  [{'verdict': 'Accepted', 'language': 'PyPy 3-6...   \n",
       "\n",
       "                                          solution_url  \\\n",
       "142  https://codeforces.com/contest/1789/submission...   \n",
       "\n",
       "                                      working_solution  \\\n",
       "142  class solve:\\n    def __init__(self):\\n       ...   \n",
       "\n",
       "                                               id_hash time_limit  \\\n",
       "142  c5c18430aba9033582c658845be4146bcbf13ccf522fa2...   1 second   \n",
       "\n",
       "      memory_limit interpolated_date release_date  \n",
       "142  256 megabytes              None   2023-02-25  \n",
       "\n",
       "[1 rows x 26 columns]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Problem Statement:\n",
      "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.\n",
      "Serval can perform the following operation called Inversion Magic on the string s:\n",
      "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.\n",
      "For example, let s be 010100 and the segment [2,5] is chosen. The string s will be 001010 after performing the Inversion Magic.\n",
      "Serval wants to make s a palindrome after performing Inversion Magic exactly once. Help him to determine whether it is possible.\n",
      "A string is a palindrome iff it reads the same backwards as forwards. For example, 010010 is a palindrome but 10111 is not.\n"
     ]
    }
   ],
   "source": [
    "problem_id = \"1789B\"\n",
    "prob = df[df.id == problem_id]\n",
    "display(df[df.id == problem_id])\n",
    "print(f\"Problem Statement:\\n{prob.problem_description.values[0]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solving a Codeforces Problem with FunSearch\n",
    "\n",
    "\n",
    "Let's employ FunSearch to tackle a Codeforces problem. \n",
    "\n",
    "#### Serving, and Getting an Instance of FunSearch\n",
    "\n",
    "To facilitate this process, we've developed a helper function named `get_configs_cf`, which is responsible for setting up the configuration files and providing an initial \"dummy solution\" for FunSearch to address a designated Codeforces problem -— in this case, problem ID 1789B. We'll delve into the details of how to configure these configuration files in Section 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "problem_id = \"1789B\"\n",
    "funsearch_cfg, dummy_solution = get_configs_cf(problem_id,ds_location=\"./data/codeforces.jsonl.gz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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",
      "\n"
     ]
    }
   ],
   "source": [
    "print(dummy_solution)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"name\": \"FunSearchFlow\",\n",
      "    \"description\": \"A flow implementing FunSearch Asynchronous Search\",\n",
      "    \"subflows_config\": {\n",
      "        \"SamplerFlow\": {\n",
      "            \"name\": \"Sampler Flow\",\n",
      "            \"description\": \"A flow that queries an LLM model to generate prompts\",\n",
      "            \"system_message_prompt_template\": {\n",
      "                \"partial_variables\": {\n",
      "                    \"evaluate_name\": \"evaluate\",\n",
      "                    \"evolve_name\": \"solve_function\",\n",
      "                    \"artifacts_per_prompt\": 2\n",
      "                }\n",
      "            },\n",
      "            \"backend\": {\n",
      "                \"api_infos\": \"???\"\n",
      "            }\n",
      "        },\n",
      "        \"EvaluatorFlow\": {\n",
      "            \"name\": \"A flow that evaluates code on tests\",\n",
      "            \"description\": \"A flow that evaluates code on tests\",\n",
      "            \"singleton\": false,\n",
      "            \"run_error_score\": -1,\n",
      "            \"py_file\": \"\\\"\\\"\\\"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) \\u2014 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    \",\n",
      "            \"function_to_run_name\": \"evaluate\",\n",
      "            \"test_inputs\": {\n",
      "                \"test_1\": {\n",
      "                    \"tests_inputs\": [\n",
      "                        \"1\",\n",
      "                        \"4\",\n",
      "                        \"1001\"\n",
      "                    ],\n",
      "                    \"expected_outputs\": \"YES\"\n",
      "                },\n",
      "                \"test_2\": {\n",
      "                    \"tests_inputs\": [\n",
      "                        \"1\",\n",
      "                        \"5\",\n",
      "                        \"10010\"\n",
      "                    ],\n",
      "                    \"expected_outputs\": \"YES\"\n",
      "                },\n",
      "                \"test_3\": {\n",
      "                    \"tests_inputs\": [\n",
      "                        \"1\",\n",
      "                        \"7\",\n",
      "                        \"0111011\"\n",
      "                    ],\n",
      "                    \"expected_outputs\": \"NO\"\n",
      "                },\n",
      "                \"test_4\": {\n",
      "                    \"tests_inputs\": [\n",
      "                        \"3\",\n",
      "                        \"4\",\n",
      "                        \"1001\",\n",
      "                        \"5\",\n",
      "                        \"10010\",\n",
      "                        \"7\",\n",
      "                        \"0111011\",\n",
      "                        \"\"\n",
      "                    ],\n",
      "                    \"expected_outputs\": \"YES\\nYES\\nNO\\n\"\n",
      "                }\n",
      "            },\n",
      "            \"timeout_seconds\": 10,\n",
      "            \"use_test_input_as_key\": false\n",
      "        },\n",
      "        \"ProgramDBFlow\": {\n",
      "            \"name\": \"ProgramDB\",\n",
      "            \"description\": \"A flow that registers samples and evaluations in a database\",\n",
      "            \"evaluate_function\": \"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\",\n",
      "            \"evaluate_file_full_content\": \"\\\"\\\"\\\"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) \\u2014 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    \",\n",
      "            \"artifact_to_evolve_name\": \"solve_function\",\n",
      "            \"num_islands\": 3,\n",
      "            \"template\": {\n",
      "                \"preface\": \"\"\n",
      "            }\n",
      "        }\n",
      "    }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(json.dumps(funsearch_cfg, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's start serving FunSearch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "serving.recursive_serve_flow(\n",
    "        cl=cl,\n",
    "        flow_class_name=\"flow_modules.aiflows.FunSearchFlowModule.FunSearch\",\n",
    "        flow_endpoint=\"FunSearch\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also configure how many workers we want to use for FunSearch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_workers = 10\n",
    "for i in range(n_workers):\n",
    "    run_dispatch_worker_thread(cl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can get an instance of FunSearch set up to solve problem 1789B:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "api_information = [ApiInfo(backend_used=\"openai\",api_key = os.getenv(\"OPENAI_API_KEY\"))]\n",
    "quick_load_api_keys(funsearch_cfg, api_information, key=\"api_infos\")\n",
    "funsearch_proxy = serving.get_flow_instance(\n",
    "        cl=cl,\n",
    "        flow_endpoint=\"FunSearch\",\n",
    "        config_overrides=funsearch_cfg,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Interacting with FunSearch\n",
    "\n",
    "In this subsection we will be showing how to interact with FunSearch. More specifically, we will be showing how to:\n",
    "- Register a solution to FunSearch (i.e., provide a solution to FunSearch, typically a dummy solution)\n",
    "- Get the best solutions produced by FunSearch\n",
    "- Start the search process\n",
    "- Stop the search process\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Registering a Solution\n",
    "\n",
    "To get started with FunSearch, the first thing you need to do is register a solution. This can be a simple placeholder solution that FunSearch will try to make better. Just remember, FunSearch won't be able to begin searching without a solution provided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "##How to Register a Program to FunSearch\n",
    "\n",
    "#Register the first program\n",
    "#You can register a program by sending a message to the FunSearch flow (pretending be a SamplerFlow)\n",
    "data = {\n",
    "        \"from\": \"SamplerFlow\", #Keep this as is\n",
    "        \"operation\": \"register_program\", #Keep this as is\n",
    "        \"api_output\": dummy_solution #This is the solution you want to register\n",
    "    }\n",
    "#Package the message\n",
    "input_message = funsearch_proxy.package_input_message(data = data)\n",
    "#send the message (without expecting a response)\n",
    "funsearch_proxy.send_message(input_message)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Get the best solutions produced by FunSearch\n",
    "\n",
    "Here below we show how to get the best solutions produced by FunSearch. Since we haven't actually started FunSearch and have only registered a dummy solution, we will only see the dummy solution.\n",
    "\n",
    "The `get_best_solutions` method returns a list of dictionaries containing the best solutions found by FunSearch so far on each island of the ProgramDB. Each item in the list is a dictionary with the following keys:\n",
    "- `rank`: The rank of the solution (among the best solutions found by FunSearch). Best rank is 1\n",
    "- `score`: The score of the solution (for codeforces problems, the maximal score is 1.0)\n",
    "- `island_id`: The id of the island where the solution was found in the ProgramDB\n",
    "- `program`: The solution itself"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "waiting for response....\n",
      "\n",
      "### Rank: 3\n",
      "### Score: 0.25\n",
      "### Island ID: 0\n",
      "### Program:\n",
      "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",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "### Rank: 2\n",
      "### Score: 0.25\n",
      "### Island ID: 1\n",
      "### Program:\n",
      "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",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "### Rank: 1\n",
      "### Score: 0.25\n",
      "### Island ID: 2\n",
      "### Program:\n",
      "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",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#Fetch Best Programs\n",
    "data = {\n",
    "        \"from\": \"FunSearch\", #Keep this as is\n",
    "        \"operation\": \"get_best_programs_per_island\", #Keep this as is\n",
    "        \"content\": {} #Keep this as is\n",
    "}\n",
    "#Package the message\n",
    "input_message = funsearch_proxy.package_input_message(data = data)\n",
    "#send the message (expecting a response)\n",
    "future = funsearch_proxy.get_reply_future(input_message)\n",
    "print(\"waiting for response....\")\n",
    "#wait for the response (this will block until the response is received)\n",
    "response = future.get_data()\n",
    "#Print the response\n",
    "nice_response = make_best_program_per_island_nice(response) #This is a helper function to make the response more readable\n",
    "print(\"\\n\\n\".join(nice_response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Start the search process of FunSearch\n",
    "\n",
    "Here below is how to start the search process of FunSearch. FunSearch will try to improve the solution you provided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Start FunSearch\n",
    "data = {\n",
    "        \"from\": \"FunSearch\", #Keep this as is\n",
    "        \"operation\": \"start\", #Keep this as is\n",
    "        \"content\": {\"num_samplers\": 5}, #This is the number of samplers you want to start\n",
    "    }\n",
    "#NOTE: You are not assured that all the samplers will start. \n",
    "#       The number of samplers that will start is dependent on the number of available workers\n",
    "\n",
    "#Package the message\n",
    "input_message = funsearch_proxy.package_input_message(data = data)\n",
    "#send the message (without expecting a response)\n",
    "funsearch_proxy.send_message(input_message)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can keep on checking the best solutions produced by FunSearch by calling `get_best_solutions` method (just run the cell below).\n",
    "In my experience, **FunSearch takes not more than a couple minutes to find an optimal solution** for this codeforces problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "waiting for response....\n",
      "\n",
      "### Rank: 3\n",
      "### Score: 0.5\n",
      "### Island ID: 0\n",
      "### Program:\n",
      "def solve_function_v1(input) -> str:\n",
      "    \"\"\"Improved version of solve_function_v0\"\"\"\n",
      "    t = int(input[0])\n",
      "    results = []\n",
      "    for i in range(t):\n",
      "        n = int(input[i*2+1])\n",
      "        s = input[i*2+2]\n",
      "        if s == s[::-1]:\n",
      "            results.append('YES')\n",
      "        else:\n",
      "            results.append('NO')\n",
      "    return ' '.join(results)\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "### Rank: 2\n",
      "### Score: 0.5\n",
      "### Island ID: 1\n",
      "### Program:\n",
      "def solve_function_v2(input) -> str:\n",
      "    \"\"\"Improved version of solve_function_v1\"\"\"\n",
      "    def is_palindrome(s):\n",
      "        return s == s[::-1]\n",
      "\n",
      "    def solve_case(s):\n",
      "        zero_count = s.count('0')\n",
      "        one_count = s.count('1')\n",
      "        if zero_count % 2 == 1 and one_count % 2 == 1:\n",
      "            return 'NO'\n",
      "        else:\n",
      "            return 'YES' if is_palindrome(s) else 'NO'\n",
      "\n",
      "    t = int(input[0])\n",
      "    results = []\n",
      "    for i in range(t):\n",
      "        n = int(input[i*2+1])\n",
      "        s = input[i*2+2]\n",
      "        results.append(solve_case(s))\n",
      "    return ' '.join(results)\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "### Rank: 1\n",
      "### Score: 1.0\n",
      "### Island ID: 2\n",
      "### Program:\n",
      "def solve_function_v2(input) -> str:\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': 1.0, 'feedback': 'No feedback available.'} test_4:{'score': 1.0, 'feedback': 'No feedback available.'}\n",
      "    Improved version of solve_function_v1\"\"\"\n",
      "    t = int(input[0])\n",
      "    results = []\n",
      "    for i in range(t):\n",
      "        n = int(input[i*2+1])\n",
      "        s = input[i*2+2]\n",
      "        ones = [i for i, c in enumerate(s) if c == '1']\n",
      "        if not ones:\n",
      "            results.append('YES')\n",
      "        else:\n",
      "            first, last = ones[0], ones[-1]\n",
      "            if s[first:last+1] == s[first:last+1][::-1]:\n",
      "                results.append('YES')\n",
      "            else:\n",
      "                results.append('NO')\n",
      "    return ' '.join(results)\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#Fetch Best Programs\n",
    "data = {\n",
    "        \"from\": \"FunSearch\",\n",
    "        \"operation\": \"get_best_programs_per_island\",\n",
    "        \"content\": {}\n",
    "}\n",
    "    \n",
    "input_message = funsearch_proxy.package_input_message(data = data)\n",
    "\n",
    "future = funsearch_proxy.get_reply_future(input_message)\n",
    "print(\"waiting for response....\")\n",
    "response = future.get_data()\n",
    "nice_response = make_best_program_per_island_nice(response)\n",
    "print(\"\\n\".join(nice_response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Stop the search process of FunSearch\n",
    "\n",
    "Here below is how to stop the search process of FunSearch. More specifally, it will let samplers who are ongoing to finish their work and then stop the search process. Note that you can always start the search process again by calling `start` method (see above)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Stop FunSearch\n",
    "input_message = funsearch_proxy.package_input_message(data = data)\n",
    "    \n",
    "funsearch_proxy.send_message(input_message)\n",
    "\n",
    "data = {\n",
    "    \"from\": \"FunSearch\",\n",
    "    \"operation\": \"stop\",\n",
    "    \"content\": {},\n",
    "}\n",
    "\n",
    "input_message = funsearch_proxy.package_input_message(data = data)\n",
    "\n",
    "funsearch_proxy.send_message(input_message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 2: Adapting FunSearch to your own Problem\n",
    "\n",
    "In this section, we will show how to adapt FunSearch to your own problem. We will showcase this by adapting FunSearch to solve a problem presented in the paper by DeepMind [[2]](https://www.nature.com/articles/s41586-023-06924-6): The capset problem.\n",
    "\n",
    "For details on the capset problem, please refer to the original paper [[2]](https://www.nature.com/articles/s41586-023-06924-6) or [Wikipedia](https://en.wikipedia.org/wiki/Cap_set)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining your own problem\n",
    "\n",
    "The first thing you need to do is to **define your evaluate in a py file**. You can chose the name of your evaluate function. In the example below we chose the name `evalute`. This function should take as input:\n",
    "- `function_str`: The function to evaluate passed as a string\n",
    "- any other arguments you need. In this case, we add the `n` too\n",
    "and return a score. Note that you can define multiple functions that the evaluate function can call.\n",
    "\n",
    "Then you must **define a dummy solution of your evolve function in a seperate py file**. You can chose the name of your evolve function. In the example below we chose the name `priority`. The inputs of this function depend on your problem. In this case, the input is `el` corresponding to a candidate vector to be added to the capset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%compile_and_writefile ./Capset_problem/evaluate_functions.py\n",
    "\n",
    "\"\"\"Finds large cap sets.\"\"\"\n",
    "\n",
    "import ast\n",
    "import itertools\n",
    "import numpy as np\n",
    "from typing import List, Callable\n",
    "\n",
    "### !!! Don't Touch this Function !!! ###\n",
    "def 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",
    "### !!! Don't Touch this Function !!! ###\n",
    "def get_function(function_str: str):\n",
    "    local_namespace = {}\n",
    "    \n",
    "    exec(function_str,local_namespace)\n",
    "    found_name, program_name = get_function_name_from_code(function_str)\n",
    "    \n",
    "    if not found_name:\n",
    "        raise ValueError(f\"Function name not found in program: {function_str}\")\n",
    "    \n",
    "    priority_fn = local_namespace.get(program_name)\n",
    "    return priority_fn\n",
    "\n",
    "# Define Your evaluate function here\n",
    "def evaluate(function_str: str, n: int) -> int:\n",
    "    \"\"\"Returns the size of an `n`-dimensional cap set.\"\"\"\n",
    "    \n",
    "    # This line must always be in evaluate function (converts generated string sample to callable)\n",
    "    priority_fn = get_function(function_str) \n",
    "    capset = solve(priority_fn,n)\n",
    "    return len(capset)\n",
    "\n",
    "# Solve function, used in evaluate\n",
    "def solve(priority_fn: Callable, n: int) -> np.ndarray:\n",
    "    \"\"\"Returns a large cap set in `n` dimensions.\"\"\"\n",
    "    \n",
    "    all_vectors = np.array(list(itertools.product((0, 1, 2), repeat=n)), dtype=np.int32)\n",
    "\n",
    "    # Powers in decreasing order for compatibility with `itertools.product`, so\n",
    "    # that the relationship `i = all_vectors[i] @ powers` holds for all `i`.\n",
    "    powers = 3 ** np.arange(n - 1, -1, -1)\n",
    "\n",
    "    # Precompute all priorities.\n",
    "    priorities = np.array([priority_fn(tuple(vector), n) for vector in all_vectors])\n",
    "\n",
    "    # Build `capset` greedily, using priorities for prioritization.\n",
    "    capset = np.empty(shape=(0, n), dtype=np.int32)\n",
    "    while np.any(priorities != -np.inf):\n",
    "        # Add a vector with maximum priority to `capset`, and set priorities of\n",
    "        # invalidated vectors to `-inf`, so that they never get selected.\n",
    "        max_index = np.argmax(priorities)\n",
    "        vector = all_vectors[None, max_index]  # [1, n]\n",
    "        blocking = np.einsum('cn,n->c', (- capset - vector) % 3, powers)  # [C]\n",
    "        priorities[blocking] = -np.inf\n",
    "        priorities[max_index] = -np.inf\n",
    "        capset = np.concatenate([capset, vector], axis=0)\n",
    "\n",
    "    return capset\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%compile_and_writefile ./Capset_problem/evolve_functions.py\n",
    "\n",
    "# The function we wish to evolve with Funsearch\n",
    "def priority(el: tuple[int, ...], n: int) -> float:\n",
    "  \"\"\"Returns the priority with which we want to add `element` to the cap set.\"\"\"\n",
    "  return 0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading your own problem to FunSearch\n",
    "\n",
    "Next we will show how to change the configuration files in order to run FunSearch on your own problem.\n",
    "\n",
    "You should start by loading the functions you just defined in your py files. You can use the `Loader` class from the `FunSearchFlowModule` to do so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flow_modules.aiflows.FunSearchFlowModule.Loader import Loader\n",
    "\n",
    "# path where we saved the evaluate function\n",
    "path_to_evaluate_functions_file = os.path.join(\".\", \"Capset_problem\", \"evaluate_functions.py\")\n",
    "# name we gave to the evaluate function\n",
    "name_of_evaluate_function = \"evaluate\"\n",
    "# path where we saved the evolve function\n",
    "path_to_evolve_functions_file = os.path.join(\".\", \"Capset_problem\", \"evolve_functions.py\")\n",
    "# name we gave to the evolve function\n",
    "name_of_evolve_function = \"priority\"\n",
    "\n",
    "loader = Loader(file_path = path_to_evaluate_functions_file, target_name = name_of_evaluate_function)\n",
    "evaluate_function = loader.load_target()\n",
    "evaluate_file_full_content = loader.load_full_file()\n",
    "\n",
    "loader = Loader(file_path = path_to_evolve_functions_file, target_name = name_of_evolve_function)\n",
    "evolve_function = loader.load_target()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we've loaded the functions, we can define our configuration file for the problem and adjust any parameters we need.\n",
    "\n",
    "You can override more stuff: Check out https://huggingface.co/aiflows/FunSearchFlowModule for full configuration parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"subflows_config\": {\n",
      "        \"SamplerFlow\": {\n",
      "            \"system_message_prompt_template\": {\n",
      "                \"partial_variables\": {\n",
      "                    \"evaluate_name\": \"evaluate\",\n",
      "                    \"evolve_name\": \"priority\",\n",
      "                    \"artifacts_per_prompt\": 2\n",
      "                }\n",
      "            },\n",
      "            \"backend\": {\n",
      "                \"api_infos\": \"???\"\n",
      "            }\n",
      "        },\n",
      "        \"EvaluatorFlow\": {\n",
      "            \"run_error_score\": -1,\n",
      "            \"py_file\": \"\\n\\\"\\\"\\\"Finds large cap sets.\\\"\\\"\\\"\\n\\nimport ast\\nimport itertools\\nimport numpy as np\\nfrom typing import List, Callable\\n\\n###\\u00a0!!! Don't Touch this Function !!! ###\\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### !!! Don't Touch this Function !!! ###\\ndef get_function(function_str: str):\\n    local_namespace = {}\\n    \\n    exec(function_str,local_namespace)\\n    found_name, program_name = get_function_name_from_code(function_str)\\n    \\n    if not found_name:\\n        raise ValueError(f\\\"Function name not found in program: {function_str}\\\")\\n    \\n    priority_fn = local_namespace.get(program_name)\\n    return priority_fn\\n\\n# Define Your evaluate function here\\ndef evaluate(function_str: str, n: int) -> int:\\n    \\\"\\\"\\\"Returns the size of an `n`-dimensional cap set.\\\"\\\"\\\"\\n    \\n    # This line must always be in evaluate function (converts generated string sample to callable)\\n    priority_fn = get_function(function_str) \\n    capset = solve(priority_fn,n)\\n    return len(capset)\\n\\n# Solve function, used in evaluate\\ndef solve(priority_fn: Callable, n: int) -> np.ndarray:\\n    \\\"\\\"\\\"Returns a large cap set in `n` dimensions.\\\"\\\"\\\"\\n    \\n    all_vectors = np.array(list(itertools.product((0, 1, 2), repeat=n)), dtype=np.int32)\\n\\n    # Powers in decreasing order for compatibility with `itertools.product`, so\\n    # that the relationship `i = all_vectors[i] @ powers` holds for all `i`.\\n    powers = 3 ** np.arange(n - 1, -1, -1)\\n\\n    # Precompute all priorities.\\n    priorities = np.array([priority_fn(tuple(vector), n) for vector in all_vectors])\\n\\n    # Build `capset` greedily, using priorities for prioritization.\\n    capset = np.empty(shape=(0, n), dtype=np.int32)\\n    while np.any(priorities != -np.inf):\\n        # Add a vector with maximum priority to `capset`, and set priorities of\\n        # invalidated vectors to `-inf`, so that they never get selected.\\n        max_index = np.argmax(priorities)\\n        vector = all_vectors[None, max_index]  # [1, n]\\n        blocking = np.einsum('cn,n->c', (- capset - vector) % 3, powers)  # [C]\\n        priorities[blocking] = -np.inf\\n        priorities[max_index] = -np.inf\\n        capset = np.concatenate([capset, vector], axis=0)\\n\\n    return capset\\n\\n\\n\\n\\n\",\n",
      "            \"function_to_run_name\": \"evaluate\",\n",
      "            \"test_inputs\": {\n",
      "                \"n=3\": {\n",
      "                    \"n\": 3\n",
      "                }\n",
      "            },\n",
      "            \"timeout_seconds\": 10\n",
      "        },\n",
      "        \"ProgramDBFlow\": {\n",
      "            \"evaluate_function\": \"def evaluate(function_str: str, n: int) -> int:\\n    \\\"\\\"\\\"Returns the size of an `n`-dimensional cap set.\\\"\\\"\\\"\\n    priority_fn = get_function(function_str)\\n    capset = solve(priority_fn, n)\\n    return len(capset)\",\n",
      "            \"evaluate_file_full_content\": \"\\n\\\"\\\"\\\"Finds large cap sets.\\\"\\\"\\\"\\n\\nimport ast\\nimport itertools\\nimport numpy as np\\nfrom typing import List, Callable\\n\\n###\\u00a0!!! Don't Touch this Function !!! ###\\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### !!! Don't Touch this Function !!! ###\\ndef get_function(function_str: str):\\n    local_namespace = {}\\n    \\n    exec(function_str,local_namespace)\\n    found_name, program_name = get_function_name_from_code(function_str)\\n    \\n    if not found_name:\\n        raise ValueError(f\\\"Function name not found in program: {function_str}\\\")\\n    \\n    priority_fn = local_namespace.get(program_name)\\n    return priority_fn\\n\\n# Define Your evaluate function here\\ndef evaluate(function_str: str, n: int) -> int:\\n    \\\"\\\"\\\"Returns the size of an `n`-dimensional cap set.\\\"\\\"\\\"\\n    \\n    # This line must always be in evaluate function (converts generated string sample to callable)\\n    priority_fn = get_function(function_str) \\n    capset = solve(priority_fn,n)\\n    return len(capset)\\n\\n# Solve function, used in evaluate\\ndef solve(priority_fn: Callable, n: int) -> np.ndarray:\\n    \\\"\\\"\\\"Returns a large cap set in `n` dimensions.\\\"\\\"\\\"\\n    \\n    all_vectors = np.array(list(itertools.product((0, 1, 2), repeat=n)), dtype=np.int32)\\n\\n    # Powers in decreasing order for compatibility with `itertools.product`, so\\n    # that the relationship `i = all_vectors[i] @ powers` holds for all `i`.\\n    powers = 3 ** np.arange(n - 1, -1, -1)\\n\\n    # Precompute all priorities.\\n    priorities = np.array([priority_fn(tuple(vector), n) for vector in all_vectors])\\n\\n    # Build `capset` greedily, using priorities for prioritization.\\n    capset = np.empty(shape=(0, n), dtype=np.int32)\\n    while np.any(priorities != -np.inf):\\n        # Add a vector with maximum priority to `capset`, and set priorities of\\n        # invalidated vectors to `-inf`, so that they never get selected.\\n        max_index = np.argmax(priorities)\\n        vector = all_vectors[None, max_index]  # [1, n]\\n        blocking = np.einsum('cn,n->c', (- capset - vector) % 3, powers)  # [C]\\n        priorities[blocking] = -np.inf\\n        priorities[max_index] = -np.inf\\n        capset = np.concatenate([capset, vector], axis=0)\\n\\n    return capset\\n\\n\\n\\n\\n\",\n",
      "            \"artifact_to_evolve_name\": \"priority\",\n",
      "            \"num_islands\": 3,\n",
      "            \"artifacts_per_prompt\": 2\n",
      "        }\n",
      "    }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "#Note:\n",
    "artifacts_per_prompt = 2\n",
    "overrides = {\n",
    "    \"subflows_config\": {\n",
    "        #Override configuration of SamplerFlow\n",
    "        \"SamplerFlow\": {\n",
    "            #Arguments that must be passed to the system prompt template\n",
    "            \"system_message_prompt_template\": {\n",
    "                \"partial_variables\": {\n",
    "                    \"evaluate_name\": name_of_evaluate_function,\n",
    "                    \"evolve_name\": name_of_evolve_function,\n",
    "                    \"artifacts_per_prompt\": artifacts_per_prompt\n",
    "                }\n",
    "            },\n",
    "            \"backend\": {\n",
    "                \"api_infos\": \"???\", # We will add API info after\n",
    "                \"model_name\": {\n",
    "                    \"openai\": \"gpt-4\", #Follows litellm model naming (see ChatFlowModule)\n",
    "                }\n",
    "            }\n",
    "        },\n",
    "        \"EvaluatorFlow\": {\n",
    "            #Score to give to a program that fails to run\n",
    "            \"run_error_score\": -1,\n",
    "            #Necessary code to run evaluation function \n",
    "            \"py_file\": evaluate_file_full_content,\n",
    "            #Name of the evaluate function\n",
    "            \"function_to_run_name\": name_of_evaluate_function,\n",
    "            #Inputs to evaluate function on which we will score (can be more than one test)\n",
    "            \"test_inputs\": {\n",
    "                \"n=3\": {\n",
    "                    \"n\": 3,\n",
    "                },\n",
    "            },\n",
    "            #Timeout in seconds for each evaluation\n",
    "            \"timeout_seconds\": 10,\n",
    "        },\n",
    "        \"ProgramDBFlow\": {\n",
    "            #evaluate function to use to evaluate programs\n",
    "            \"evaluate_function\": evaluate_function,\n",
    "            #Full content of evaluate function\n",
    "            \"evaluate_file_full_content\": evaluate_file_full_content,\n",
    "            #Name of function to evolve\n",
    "            \"artifact_to_evolve_name\": name_of_evolve_function,\n",
    "            #Number of islands. !! Default for FunSearch experiments was 10 but we are using 3 for this example\n",
    "            \"num_islands\": 3,\n",
    "            #Number of evolve examples from previous generations to include in prompt\n",
    "            \"artifacts_per_prompt\": artifacts_per_prompt\n",
    "        }\n",
    "    }\n",
    "}\n",
    "print(json.dumps(overrides, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### !!!!!!! Uncomment code below if you didn't run section 1 !!!!!! ####\n",
    "# cl = start_colink_server()\n",
    "# n_workers = 10\n",
    "# for i in range(n_workers):\n",
    "#     run_dispatch_worker_thread(cl)\n",
    "    \n",
    "# serving.recursive_serve_flow(\n",
    "#         cl=cl,\n",
    "#         flow_class_name=\"flow_modules.aiflows.FunSearchFlowModule.FunSearch\",\n",
    "#         flow_endpoint=\"FunSearch\",\n",
    "# )\n",
    "\n",
    "#Load api information\n",
    "api_information = [ApiInfo(backend_used=\"openai\",api_key = os.getenv(\"OPENAI_API_KEY\"))]\n",
    "quick_load_api_keys(overrides, api_information, key=\"api_infos\")\n",
    "custom_funsearch_proxy = serving.get_flow_instance(\n",
    "        cl=cl,\n",
    "        flow_endpoint=\"FunSearch\",\n",
    "        config_overrides=overrides,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Register dummy solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "##How to Register a Program to FunSearch\n",
    "\n",
    "#Register the first program\n",
    "#You can register a program by sending a message to the FunSearch flow (pretending be a SamplerFlow)\n",
    "data = {\n",
    "        \"from\": \"SamplerFlow\", #Keep this as is\n",
    "        \"operation\": \"register_program\", #Keep this as is\n",
    "        \"api_output\": evolve_function #This is the solution you want to register\n",
    "    }\n",
    "#Package the message\n",
    "input_message = custom_funsearch_proxy.package_input_message(data = data)\n",
    "#send the message (without expecting a response)\n",
    "custom_funsearch_proxy.send_message(input_message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Start the search process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Start FunSearch\n",
    "data = {\n",
    "        \"from\": \"FunSearch\", #Keep this as is\n",
    "        \"operation\": \"start\", #Keep this as is\n",
    "        \"content\": {\"num_samplers\": 5}, #This is the number of samplers you want to start\n",
    "    }\n",
    "#NOTE: You are not assured that all the samplers will start. \n",
    "#       The number of samplers that will start is dependent on the number of available workers\n",
    "\n",
    "#Package the message\n",
    "input_message = custom_funsearch_proxy.package_input_message(data = data)\n",
    "#send the message (without expecting a response)\n",
    "custom_funsearch_proxy.send_message(input_message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Query Best Solutions (Run as many times as you like)\n",
    "\n",
    "For cost reasons, we will not run the search process here. But you can run the search process by calling the `start` method. We might run the search process in the future and update this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "waiting for response....\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "### Rank: 3\n",
      "### Score: 8.0\n",
      "### Island ID: 0\n",
      "### Program:\n",
      "def priority(el: tuple[int, ...], n: int) -> float:\n",
      "    \"\"\"Returns the priority with which we want to add `element` to the cap set.\"\"\"\n",
      "    return 0.0\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "### Rank: 2\n",
      "### Score: 8.0\n",
      "### Island ID: 1\n",
      "### Program:\n",
      "def priority(el: tuple[int, ...], n: int) -> float:\n",
      "    \"\"\"Returns the priority with which we want to add `element` to the cap set.\"\"\"\n",
      "    return 0.0\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "### Rank: 1\n",
      "### Score: 8.0\n",
      "### Island ID: 2\n",
      "### Program:\n",
      "def priority(el: tuple[int, ...], n: int) -> float:\n",
      "    \"\"\"Returns the priority with which we want to add `element` to the cap set.\"\"\"\n",
      "    return 0.0\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#Fetch Best Programs\n",
    "data = {\n",
    "        \"from\": \"FunSearch\", #Keep this as is\n",
    "        \"operation\": \"get_best_programs_per_island\", #Keep this as is\n",
    "        \"content\": {} #Keep this as is\n",
    "}\n",
    "#Package the message\n",
    "input_message = custom_funsearch_proxy.package_input_message(data = data)\n",
    "#send the message (expecting a response)\n",
    "future = custom_funsearch_proxy.get_reply_future(input_message)\n",
    "print(\"waiting for response....\")\n",
    "#wait for the response (this will block until the response is received)\n",
    "response = future.get_data()\n",
    "#Print the response\n",
    "nice_response = make_best_program_per_island_nice(response) #This is a helper function to make the response more readable\n",
    "print(\"\\n\\n\".join(nice_response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Stop the search process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Stop FunSearch\n",
    "input_message = custom_funsearch_proxy.package_input_message(data = data)\n",
    "    \n",
    "custom_funsearch_proxy.send_message(input_message)\n",
    "\n",
    "data = {\n",
    "    \"from\": \"FunSearch\",\n",
    "    \"operation\": \"stop\",\n",
    "    \"content\": {},\n",
    "}\n",
    "\n",
    "input_message = custom_funsearch_proxy.package_input_message(data = data)\n",
    "\n",
    "custom_funsearch_proxy.send_message(input_message)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}